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

refactor: Update ChatLiteLLMModelComponent to use Message from langflow.schema.message

Browse source
This commit is contained in:
ogabrielluiz 2024-06-18 15:00:19 -03:00
commit c04cdb3412
1 changed files with 4 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

7
src/backend/base/langflow/components/models/ChatLiteLLMModel.py
View file

@ -4,8 +4,9 @@ from langchain_community.chat_models.litellm import ChatLiteLLM, ChatLiteLLMExce
from langflow.base.constants import STREAM_INFO_TEXT from langflow.base.constants import STREAM_INFO_TEXT
from langflow.base.models.model import LCModelComponent from langflow.base.models.model import LCModelComponent
from langflow.field_typing import BaseLanguageModel
from langflow.inputs import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, MessageInput, StrInput from langflow.inputs import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, MessageInput, StrInput
from langflow.template import Output from langflow.schema.message import Message
from langflow.template import Output from langflow.template import Output
@ -13,7 +14,7 @@ class ChatLiteLLMModelComponent(LCModelComponent):
display_name = "LiteLLM" display_name = "LiteLLM"
description = "`LiteLLM` collection of large language models." description = "`LiteLLM` collection of large language models."
documentation = "https://python.langchain.com/docs/integrations/chat/litellm" documentation = "https://python.langchain.com/docs/integrations/chat/litellm"
icon = "🚅" icon = "LiteLLM"
inputs = [ inputs = [
MessageInput(name="input_value", display_name="Input"), MessageInput(name="input_value", display_name="Input"),
@ -119,7 +120,7 @@ class ChatLiteLLMModelComponent(LCModelComponent):
Output(display_name="Language Model", name="model_output", method="build_model"), Output(display_name="Language Model", name="model_output", method="build_model"),
] ]
def text_response(self) -> Text: def text_response(self) -> Message:
input_value = self.input_value input_value = self.input_value
stream = self.stream stream = self.stream
system_message = self.system_message system_message = self.system_message