From 6e131120af399b605c3117ff0e420c5a15938126 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 3 Apr 2024 19:33:03 -0300 Subject: [PATCH] Add ChatLiteLLMModelComponent to models package --- .../model_specs/ChatLiteLLMSpecs.py | 2 +- .../components/models/ChatLiteLLMModel.py | 191 ++++++++++++++++++ .../langflow/components/models/__init__.py | 2 + 3 files changed, 194 insertions(+), 1 deletion(-) create mode 100644 src/backend/base/langflow/components/models/ChatLiteLLMModel.py diff --git a/src/backend/base/langflow/components/model_specs/ChatLiteLLMSpecs.py b/src/backend/base/langflow/components/model_specs/ChatLiteLLMSpecs.py index 5c1463b71..439266fbb 100644 --- a/src/backend/base/langflow/components/model_specs/ChatLiteLLMSpecs.py +++ b/src/backend/base/langflow/components/model_specs/ChatLiteLLMSpecs.py @@ -118,7 +118,7 @@ class ChatLiteLLMComponent(CustomComponent): max_tokens: int = 256, max_retries: int = 6, verbose: bool = False, - ) -> Union[BaseLanguageModel, Callable]: + ) -> BaseLanguageModel: try: import litellm # type: ignore diff --git a/src/backend/base/langflow/components/models/ChatLiteLLMModel.py b/src/backend/base/langflow/components/models/ChatLiteLLMModel.py new file mode 100644 index 000000000..5266c6935 --- /dev/null +++ b/src/backend/base/langflow/components/models/ChatLiteLLMModel.py @@ -0,0 +1,191 @@ +from typing import Any, Dict, Optional + +from langchain_community.chat_models.litellm import ChatLiteLLM, ChatLiteLLMException + +from langflow.base.constants import STREAM_INFO_TEXT +from langflow.base.models.model import LCModelComponent +from langflow.field_typing import BaseLanguageModel, Text + + +class ChatLiteLLMModelComponent(LCModelComponent): + display_name = "LiteLLM" + description = "`LiteLLM` collection of large language models." + documentation = "https://python.langchain.com/docs/integrations/chat/litellm" + field_order = [ + "model", + "api_key", + "provider", + "temperature", + "model_kwargs", + "top_p", + "top_k", + "n", + "max_tokens", + "max_retries", + "verbose", + "stream", + "input_value", + "system_message", + ] + + def build_config(self): + return { + "model": { + "display_name": "Model name", + "field_type": "str", + "advanced": False, + "required": True, + "info": "The name of the model to use. For example, `gpt-3.5-turbo`.", + }, + "api_key": { + "display_name": "API key", + "field_type": "str", + "advanced": False, + "required": False, + "password": True, + }, + "provider": { + "display_name": "Provider", + "info": "The provider of the API key.", + "options": [ + "OpenAI", + "Azure", + "Anthropic", + "Replicate", + "Cohere", + "OpenRouter", + ], + }, + "temperature": { + "display_name": "Temperature", + "field_type": "float", + "advanced": False, + "required": False, + "default": 0.7, + }, + "model_kwargs": { + "display_name": "Model kwargs", + "field_type": "dict", + "advanced": True, + "required": False, + "default": {}, + }, + "top_p": { + "display_name": "Top p", + "field_type": "float", + "advanced": True, + "required": False, + }, + "top_k": { + "display_name": "Top k", + "field_type": "int", + "advanced": True, + "required": False, + }, + "n": { + "display_name": "N", + "field_type": "int", + "advanced": True, + "required": False, + "info": "Number of chat completions to generate for each prompt. " + "Note that the API may not return the full n completions if duplicates are generated.", + "default": 1, + }, + "max_tokens": { + "display_name": "Max tokens", + "field_type": "int", + "advanced": False, + "required": False, + "default": 256, + "info": "The maximum number of tokens to generate for each chat completion.", + }, + "max_retries": { + "display_name": "Max retries", + "field_type": "int", + "advanced": True, + "required": False, + "default": 6, + }, + "verbose": { + "display_name": "Verbose", + "field_type": "bool", + "advanced": True, + "required": False, + "default": False, + }, + "input_value": {"display_name": "Input"}, + "stream": { + "display_name": "Stream", + "info": STREAM_INFO_TEXT, + "advanced": True, + }, + "system_message": { + "display_name": "System Message", + "info": "System message to pass to the model.", + "advanced": True, + }, + } + + def build( + self, + input_value: Text, + model: str, + provider: str, + api_key: Optional[str] = None, + stream: bool = False, + temperature: Optional[float] = 0.7, + model_kwargs: Optional[Dict[str, Any]] = {}, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + n: int = 1, + max_tokens: int = 256, + max_retries: int = 6, + verbose: bool = False, + system_message: Optional[str] = None, + ) -> BaseLanguageModel: + try: + import litellm # type: ignore + + litellm.drop_params = True + litellm.set_verbose = verbose + except ImportError: + raise ChatLiteLLMException( + "Could not import litellm python package. " "Please install it with `pip install litellm`" + ) + provider_map = { + "OpenAI": "openai_api_key", + "Azure": "azure_api_key", + "Anthropic": "anthropic_api_key", + "Replicate": "replicate_api_key", + "Cohere": "cohere_api_key", + "OpenRouter": "openrouter_api_key", + } + # Set the API key based on the provider + api_keys: dict[str, Optional[str]] = {v: None for v in provider_map.values()} + + if variable_name := provider_map.get(provider): + api_keys[variable_name] = api_key + else: + raise ChatLiteLLMException( + f"Provider {provider} is not supported. Supported providers are: {', '.join(provider_map.keys())}" + ) + + output = ChatLiteLLM( + model=model, + client=None, + streaming=stream, + temperature=temperature, + model_kwargs=model_kwargs if model_kwargs is not None else {}, + top_p=top_p, + top_k=top_k, + n=n, + max_tokens=max_tokens, + max_retries=max_retries, + openai_api_key=api_keys["openai_api_key"], + azure_api_key=api_keys["azure_api_key"], + anthropic_api_key=api_keys["anthropic_api_key"], + replicate_api_key=api_keys["replicate_api_key"], + cohere_api_key=api_keys["cohere_api_key"], + openrouter_api_key=api_keys["openrouter_api_key"], + ) + return self.get_chat_result(output, stream, input_value, system_message) diff --git a/src/backend/base/langflow/components/models/__init__.py b/src/backend/base/langflow/components/models/__init__.py index 3efbe4b6f..9db6caa26 100644 --- a/src/backend/base/langflow/components/models/__init__.py +++ b/src/backend/base/langflow/components/models/__init__.py @@ -2,6 +2,7 @@ from .AmazonBedrockModel import AmazonBedrockComponent from .AnthropicModel import AnthropicLLM from .AzureOpenAIModel import AzureChatOpenAIComponent from .BaiduQianfanChatModel import QianfanChatEndpointComponent +from .ChatLiteLLMModel import ChatLiteLLMModelComponent from .CohereModel import CohereComponent from .GoogleGenerativeAIModel import GoogleGenerativeAIComponent from .HuggingFaceModel import HuggingFaceEndpointsComponent @@ -10,6 +11,7 @@ from .OpenAIModel import OpenAIModelComponent from .VertexAiModel import ChatVertexAIComponent __all__ = [ + "ChatLiteLLMModelComponent", "AmazonBedrockComponent", "AnthropicLLM", "AzureChatOpenAIComponent",