jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Update AnthropicLLM to use BaseLanguageModel and

Browse source 
SecretStr
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-12-04 11:44:56 -03:00
commit 6afbd32621
1 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

9
src/backend/langflow/components/llms/AnthropicLLM.py
View file

@ -1,7 +1,8 @@
from typing import Optional from typing import Optional
from langchain.chat_models.anthropic import ChatAnthropic from langchain.chat_models.anthropic import ChatAnthropic
from langchain.llms.base import BaseLLM from langchain.llms.base import BaseLanguageModel
from pydantic import SecretStr
from langflow import CustomComponent from langflow import CustomComponent
@ -55,7 +56,7 @@ class AnthropicLLM(CustomComponent):
max_tokens: Optional[int] = None, max_tokens: Optional[int] = None,
temperature: Optional[float] = None, temperature: Optional[float] = None,
api_endpoint: Optional[str] = None, api_endpoint: Optional[str] = None,
) -> BaseLLM: ) -> BaseLanguageModel:
# Set default API endpoint if not provided # Set default API endpoint if not provided
if not api_endpoint: if not api_endpoint:
api_endpoint = "https://api.anthropic.com" api_endpoint = "https://api.anthropic.com"
@ -63,8 +64,8 @@ class AnthropicLLM(CustomComponent):
try: try:
output = ChatAnthropic( output = ChatAnthropic(
model_name=model, model_name=model,
anthropic_api_key=anthropic_api_key, anthropic_api_key=SecretStr(anthropic_api_key),
max_tokens_to_sample=max_tokens, max_tokens_to_sample=max_tokens, # type: ignore
temperature=temperature, temperature=temperature,
anthropic_api_url=api_endpoint, anthropic_api_url=api_endpoint,
) )