From b8a361c5462eb738a0da07ea47b853588b314036 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Fri, 7 Jun 2024 15:21:03 -0300 Subject: [PATCH] refactor: Improve handling of artifacts in vertex base class --- src/backend/base/langflow/graph/schema.py | 4 ++-- src/backend/base/langflow/graph/vertex/base.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/backend/base/langflow/graph/schema.py b/src/backend/base/langflow/graph/schema.py index 82cfa2930..72b810555 100644 --- a/src/backend/base/langflow/graph/schema.py +++ b/src/backend/base/langflow/graph/schema.py @@ -31,10 +31,10 @@ class ResultData(BaseModel): if not values.get("logs") and values.get("artifacts"): # Build the log from the artifacts message = values["artifacts"] - if "stream_url" in message: + if "stream_url" in message and "type" in message: stream_url = StreamURL(location=message["stream_url"]) values["logs"] = [Log(message=stream_url, type=message["type"])] - else: + elif "type" in message: values["logs"] = [Log(message=message, type=message["type"])] return values diff --git a/src/backend/base/langflow/graph/vertex/base.py b/src/backend/base/langflow/graph/vertex/base.py index d3c9b1a9a..1086ee337 100644 --- a/src/backend/base/langflow/graph/vertex/base.py +++ b/src/backend/base/langflow/graph/vertex/base.py @@ -628,8 +628,8 @@ class Vertex: self._built_object, self.artifacts = result elif len(result) == 3: self._custom_component, self._built_object, self.artifacts = result - self.artifacts_raw = self.artifacts.pop("raw", None) - self.artifacts_type = self.artifacts.pop("type", None) or ArtifactType.UNKNOWN.value + self.artifacts_raw = self.artifacts.get("raw", None) + self.artifacts_type = self.artifacts.get("type", None) or ArtifactType.UNKNOWN.value else: self._built_object = result