From a142b45160004b5c173bf023cea70d673550f86b Mon Sep 17 00:00:00 2001 From: anovazzi1 Date: Mon, 6 Jan 2025 18:01:09 -0300 Subject: [PATCH] feat: make Agents not send message if not connected to an output component (#5521) * Refactor: Commented out TextInput and TextOutput components * Refactor: Add function to check if graph has output vertex connected to it * Refactor: Add function to skip messages based on vertex configuration and message type * Refactor: Update function to check for chat output in graph vertices * Refactor: Enhance message skipping logic to include vertex existence check --------- Co-authored-by: Gabriel Luiz Freitas Almeida --- .../langflow/custom/custom_component/component.py | 12 ++++++++++++ src/backend/base/langflow/graph/schema.py | 4 ++-- src/backend/base/langflow/graph/utils.py | 10 ++++++++++ 3 files changed, 24 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/custom/custom_component/component.py b/src/backend/base/langflow/custom/custom_component/component.py index 4cb0bbb8d..aa51a223c 100644 --- a/src/backend/base/langflow/custom/custom_component/component.py +++ b/src/backend/base/langflow/custom/custom_component/component.py @@ -25,6 +25,7 @@ from langflow.custom.tree_visitor import RequiredInputsVisitor from langflow.exceptions.component import StreamingError from langflow.field_typing import Tool # noqa: TC001 Needed by _add_toolkit_output from langflow.graph.state.model import create_state_model +from langflow.graph.utils import has_chat_output from langflow.helpers.custom import format_type from langflow.memory import astore_message, aupdate_messages, delete_message from langflow.schema.artifact import get_artifact_type, post_process_raw @@ -1019,7 +1020,18 @@ class Component(CustomComponent): ) ) + def _should_skip_message(self, message: Message) -> bool: + """Check if the message should be skipped based on vertex configuration and message type.""" + return ( + self._vertex is not None + and not (self._vertex.is_output or self._vertex.is_input) + and not has_chat_output(self.graph.get_vertex_neighbors(self._vertex)) + and not isinstance(message, ErrorMessage) + ) + async def send_message(self, message: Message, id_: str | None = None): + if self._should_skip_message(message): + return message if (hasattr(self, "graph") and self.graph.session_id) and (message is not None and not message.session_id): session_id = ( UUID(self.graph.session_id) if isinstance(self.graph.session_id, str) else self.graph.session_id diff --git a/src/backend/base/langflow/graph/schema.py b/src/backend/base/langflow/graph/schema.py index dc643d647..d792868c1 100644 --- a/src/backend/base/langflow/graph/schema.py +++ b/src/backend/base/langflow/graph/schema.py @@ -60,13 +60,13 @@ CHAT_COMPONENTS = [InterfaceComponentTypes.ChatInput, InterfaceComponentTypes.Ch RECORDS_COMPONENTS = [InterfaceComponentTypes.DataOutput] INPUT_COMPONENTS = [ InterfaceComponentTypes.ChatInput, - InterfaceComponentTypes.TextInput, InterfaceComponentTypes.WebhookInput, + InterfaceComponentTypes.TextInput, ] OUTPUT_COMPONENTS = [ InterfaceComponentTypes.ChatOutput, - InterfaceComponentTypes.TextOutput, InterfaceComponentTypes.DataOutput, + InterfaceComponentTypes.TextOutput, ] diff --git a/src/backend/base/langflow/graph/utils.py b/src/backend/base/langflow/graph/utils.py index 0a83c8aa1..fe4840715 100644 --- a/src/backend/base/langflow/graph/utils.py +++ b/src/backend/base/langflow/graph/utils.py @@ -208,3 +208,13 @@ def rewrite_file_path(file_path: str): consistent_file_path = "/".join(file_path_split) return [consistent_file_path] + + +def has_output_vertex(vertices: dict[Vertex, int]): + return any(vertex.is_output for vertex in vertices) + + +def has_chat_output(vertices: dict[Vertex, int]): + from langflow.graph.schema import InterfaceComponentTypes + + return any(InterfaceComponentTypes.ChatOutput in vertex.id for vertex in vertices)