From 6afbd32621f3c81cfc3bce6c0cd6f59fff47fc33 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 4 Dec 2023 11:44:56 -0300 Subject: [PATCH] Update AnthropicLLM to use BaseLanguageModel and SecretStr --- src/backend/langflow/components/llms/AnthropicLLM.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/components/llms/AnthropicLLM.py b/src/backend/langflow/components/llms/AnthropicLLM.py index ff6a24a0e..34017bd41 100644 --- a/src/backend/langflow/components/llms/AnthropicLLM.py +++ b/src/backend/langflow/components/llms/AnthropicLLM.py @@ -1,7 +1,8 @@ from typing import Optional 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 @@ -55,7 +56,7 @@ class AnthropicLLM(CustomComponent): max_tokens: Optional[int] = None, temperature: Optional[float] = None, api_endpoint: Optional[str] = None, - ) -> BaseLLM: + ) -> BaseLanguageModel: # Set default API endpoint if not provided if not api_endpoint: api_endpoint = "https://api.anthropic.com" @@ -63,8 +64,8 @@ class AnthropicLLM(CustomComponent): try: output = ChatAnthropic( model_name=model, - anthropic_api_key=anthropic_api_key, - max_tokens_to_sample=max_tokens, + anthropic_api_key=SecretStr(anthropic_api_key), + max_tokens_to_sample=max_tokens, # type: ignore temperature=temperature, anthropic_api_url=api_endpoint, )