Fix vertex iteration bug and add component metadata

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-11 18:09:33 -03:00
commit 4d45f47698
6 changed files with 15 additions and 5 deletions

View file

@ -171,8 +171,7 @@ class Graph:
raise ValueError(f"Error running graph: {exc}") from exc
# Get the outputs
vertex_outputs = []
for vertex_id in self.vertices:
vertex = self.get_vertex(vertex_id)
for vertex in self.vertices:
if vertex is None:
raise ValueError(f"Vertex {vertex_id} not found")

View file

@ -13,6 +13,8 @@ class ResultData(BaseModel):
messages: Optional[list[ChatOutputResponse]] = Field(default_factory=list)
timedelta: Optional[float] = None
duration: Optional[str] = None
component_display_name: Optional[str] = None
component_id: Optional[str] = None
@field_serializer("results")
def serialize_results(self, value):

View file

@ -428,6 +428,8 @@ class Vertex:
results=result_dict,
artifacts=artifacts,
messages=messages,
component_display_name=self.display_name,
component_id=self.id,
)
self.set_result(result_dict)

View file

@ -272,6 +272,9 @@ def process_tweaks(graph_data: Dict[str, Any], tweaks: Union["Tweaks", Dict[str,
:raises ValueError: If the input is not in the expected format.
"""
if not isinstance(tweaks, dict):
tweaks = tweaks.model_dump()
nodes = validate_input(graph_data, tweaks)
nodes_map = {node.get("id"): node for node in nodes}