From f823d5ccf7b9fce8b578a609e07a53ff0fce7b98 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Thu, 13 Jun 2024 00:16:44 -0300 Subject: [PATCH] refactor: Update ChatInput and ChatOutput to handle Message input Update the ChatInput and ChatOutput components in ChatInput.py and ChatOutput.py to handle Message input. If the input value is an instance of the Message class, extract the text attribute and use it as the input value. This change ensures compatibility with the langflow schema and improves the flexibility of the components. --- .../base/langflow/components/inputs/ChatInput.py | 2 +- .../langflow/components/outputs/ChatOutput.py | 15 +++++++++------ src/backend/base/langflow/schema/schema.py | 3 --- 3 files changed, 10 insertions(+), 10 deletions(-) diff --git a/src/backend/base/langflow/components/inputs/ChatInput.py b/src/backend/base/langflow/components/inputs/ChatInput.py index f6db2e6d3..3ea68e6c7 100644 --- a/src/backend/base/langflow/components/inputs/ChatInput.py +++ b/src/backend/base/langflow/components/inputs/ChatInput.py @@ -49,7 +49,7 @@ class ChatInput(ChatComponent): sender_name=self.sender_name, session_id=self.session_id, ) - if self.session_id and isinstance(message, (Message, str)): + if self.session_id and isinstance(message, (Message, str)) and isinstance(message.text, str): self.store_message(message) self.status = message return message diff --git a/src/backend/base/langflow/components/outputs/ChatOutput.py b/src/backend/base/langflow/components/outputs/ChatOutput.py index 67fd9189e..c028cfa0e 100644 --- a/src/backend/base/langflow/components/outputs/ChatOutput.py +++ b/src/backend/base/langflow/components/outputs/ChatOutput.py @@ -39,12 +39,15 @@ class ChatOutput(ChatComponent): ] def message_response(self) -> Message: - message = Message( - text=self.input_value, - sender=self.sender, - sender_name=self.sender_name, - session_id=self.session_id, - ) + if isinstance(self.input_value, Message): + message = self.input_value + else: + message = Message( + text=self.input_value, + sender=self.sender, + sender_name=self.sender_name, + session_id=self.session_id, + ) if self.session_id and isinstance(message, Message) and isinstance(message.text, str): self.store_message(message) self.status = message diff --git a/src/backend/base/langflow/schema/schema.py b/src/backend/base/langflow/schema/schema.py index 8ac89b805..56a2bf21e 100644 --- a/src/backend/base/langflow/schema/schema.py +++ b/src/backend/base/langflow/schema/schema.py @@ -23,9 +23,6 @@ def build_logs_from_artifacts(artifacts: dict) -> dict: message = artifacts[key]["raw"] _type = artifacts[key]["type"] - if not isinstance(message, dict): - message = {"message": message} - if "stream_url" in message and "type" in message: stream_url = StreamURL(location=message["stream_url"]) log = Log(message=stream_url, type=_type)