From d7e1270436313063750f4ba35a611f9020754222 Mon Sep 17 00:00:00 2001 From: italojohnny Date: Fri, 31 May 2024 19:15:23 -0300 Subject: [PATCH] add function to determine artifact type --- src/backend/base/langflow/graph/utils.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/backend/base/langflow/graph/utils.py b/src/backend/base/langflow/graph/utils.py index f6a574114..86fbebc4a 100644 --- a/src/backend/base/langflow/graph/utils.py +++ b/src/backend/base/langflow/graph/utils.py @@ -1,4 +1,4 @@ -from typing import Any, Union +from typing import Any, Union, Generator from enum import Enum from langchain_core.documents import Document @@ -20,6 +20,7 @@ class ArtifactType(str, Enum): TEXT = "text" RECORD = "record" OBJECT = "object" + STREAM = "stream" UNKNOWN = "unknown" @@ -61,8 +62,9 @@ def serialize_field(value): return value -def get_artifact_type(value: Any) -> str: +def get_artifact_type(custom_component, build_result) -> str: result = ArtifactType.UNKNOWN + value = custom_component.repr_value match value: case Record(): result = ArtifactType.RECORD @@ -73,4 +75,8 @@ def get_artifact_type(value: Any) -> str: case dict(): result = ArtifactType.OBJECT + if result == ArtifactType.UNKNOWN: + if isinstance(build_result, Generator): + result = ArtifactType.STREAM + return result.value