Add ChatLiteLLMModelComponent to models package
This commit is contained in:
Gabriel Luiz Freitas Almeida
parent
4c04b1ef1c
commit
6e131120af
3 changed files with 194 additions and 1 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
191
src/backend/base/langflow/components/models/ChatLiteLLMModel.py
Normal file
191
src/backend/base/langflow/components/models/ChatLiteLLMModel.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue