From 72a880df1ed00138dc2ddb604780bcdb167dd3e7 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 29 Jul 2024 17:38:26 -0300 Subject: [PATCH] fix: add support for custom component message storage (#3042) * fix: remove create task call * feat: Add support for custom component message storage The code changes in `types.py` add support for storing messages in custom components. If a custom component has the attributes `should_store_message` and `store_message`, the `store_message` method will be called to store the message. This enhancement improves the functionality of the interface vertex in the graph. --- .../base/langflow/graph/vertex/types.py | 24 ++++++++++--------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/src/backend/base/langflow/graph/vertex/types.py b/src/backend/base/langflow/graph/vertex/types.py index 9acac07b4..c36487dd9 100644 --- a/src/backend/base/langflow/graph/vertex/types.py +++ b/src/backend/base/langflow/graph/vertex/types.py @@ -7,7 +7,7 @@ from langchain_core.messages import AIMessage, AIMessageChunk from loguru import logger from langflow.graph.schema import CHAT_COMPONENTS, RECORDS_COMPONENTS, InterfaceComponentTypes, ResultData -from langflow.graph.utils import UnbuiltObject, serialize_field, log_transaction, log_vertex_build +from langflow.graph.utils import UnbuiltObject, log_transaction, log_vertex_build, serialize_field from langflow.graph.vertex.base import Vertex from langflow.schema import Data from langflow.schema.artifact import ArtifactType @@ -387,16 +387,18 @@ class InterfaceVertex(ComponentVertex): for key, value in origin_vertex.results.items(): if isinstance(value, (AsyncIterator, Iterator)): origin_vertex.results[key] = complete_message - - asyncio.create_task( - log_vertex_build( - flow_id=self.graph.flow_id, - vertex_id=self.id, - valid=True, - params=self._built_object_repr(), - data=self.result, - artifacts=self.artifacts, - ) + if self._custom_component: + if hasattr(self._custom_component, "should_store_message") and hasattr( + self._custom_component, "store_message" + ): + self._custom_component.store_message(message) + log_vertex_build( + flow_id=self.graph.flow_id, + vertex_id=self.id, + valid=True, + params=self._built_object_repr(), + data=self.result, + artifacts=self.artifacts, ) self._validate_built_object()