refactor: Improve handling of artifacts in vertex base class

This commit is contained in:
ogabrielluiz 2024-06-07 15:21:03 -03:00
commit b8a361c546
2 changed files with 4 additions and 4 deletions

View file

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

View file

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