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.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-07-29 17:38:26 -03:00 committed by GitHub
commit 72a880df1e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -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()