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

fix: add support for custom component message storage (#3042)

Browse source 
* 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
1 changed files with 13 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

10
src/backend/base/langflow/graph/vertex/types.py
View file

@ -7,7 +7,7 @@ from langchain_core.messages import AIMessage, AIMessageChunk
from loguru import logger from loguru import logger
from langflow.graph.schema import CHAT_COMPONENTS, RECORDS_COMPONENTS, InterfaceComponentTypes, ResultData 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.graph.vertex.base import Vertex
from langflow.schema import Data from langflow.schema import Data
from langflow.schema.artifact import ArtifactType from langflow.schema.artifact import ArtifactType
@ -387,8 +387,11 @@ class InterfaceVertex(ComponentVertex):
for key, value in origin_vertex.results.items(): for key, value in origin_vertex.results.items():
if isinstance(value, (AsyncIterator, Iterator)): if isinstance(value, (AsyncIterator, Iterator)):
origin_vertex.results[key] = complete_message origin_vertex.results[key] = complete_message
if self._custom_component:
asyncio.create_task( if hasattr(self._custom_component, "should_store_message") and hasattr(
self._custom_component, "store_message"
):
self._custom_component.store_message(message)
log_vertex_build( log_vertex_build(
flow_id=self.graph.flow_id, flow_id=self.graph.flow_id,
vertex_id=self.id, vertex_id=self.id,
@ -397,7 +400,6 @@ class InterfaceVertex(ComponentVertex):
data=self.result, data=self.result,
artifacts=self.artifacts, artifacts=self.artifacts,
) )
)
self._validate_built_object() self._validate_built_object()
self._built = True self._built = True