From 2cd0316506861055fddb95eefadbd564a961f583 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Mon, 17 Jun 2024 17:12:24 -0300 Subject: [PATCH] refactor: Update get_artifact_type function in artifact.py --- src/backend/base/langflow/schema/artifact.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/backend/base/langflow/schema/artifact.py b/src/backend/base/langflow/schema/artifact.py index 3855195bb..c1a0d287c 100644 --- a/src/backend/base/langflow/schema/artifact.py +++ b/src/backend/base/langflow/schema/artifact.py @@ -18,8 +18,12 @@ class ArtifactType(str, Enum): def get_artifact_type(value, build_result=None) -> str: result = ArtifactType.UNKNOWN match value: + case Message(): + enum_value = get_artifact_type(value.text) + result = ArtifactType(enum_value) case Data(): - result = ArtifactType.DATA + enum_value = get_artifact_type(value.data) + result = ArtifactType(enum_value) case str(): result = ArtifactType.TEXT @@ -30,9 +34,6 @@ def get_artifact_type(value, build_result=None) -> str: case list(): result = ArtifactType.ARRAY - case Message(): - result = ArtifactType.MESSAGE - if result == ArtifactType.UNKNOWN: if build_result and isinstance(build_result, Generator): result = ArtifactType.STREAM