vellum-ai 0.14.0__py3-none-any.whl → 0.14.2__py3-none-any.whl

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -25,6 +25,8 @@ from vellum_ee.workflows.display.vellum import (
     EntrypointVellumDisplay,
     EntrypointVellumDisplayOverrides,
     NodeDisplayData,
+    StateValueVellumDisplay,
+    StateValueVellumDisplayOverrides,
     WorkflowInputsVellumDisplay,
     WorkflowInputsVellumDisplayOverrides,
     WorkflowMetaVellumDisplay,
@@ -44,6 +46,8 @@ class VellumWorkflowDisplay(
         WorkflowMetaVellumDisplayOverrides,
         WorkflowInputsVellumDisplay,
         WorkflowInputsVellumDisplayOverrides,
+        StateValueVellumDisplay,
+        StateValueVellumDisplayOverrides,
         BaseNodeDisplay,
         EntrypointVellumDisplay,
         EntrypointVellumDisplayOverrides,
@@ -76,6 +80,25 @@ class VellumWorkflowDisplay(
                 }
             )
 
+        state_variables: JsonArray = []
+        for state_value, state_value_display in self.display_context.state_value_displays.items():
+            default = primitive_to_vellum_value(state_value.instance) if state_value.instance else None
+            required = (
+                state_value_display.required
+                if state_value_display.required is not None
+                else type(None) not in state_value.types
+            )
+            state_variables.append(
+                {
+                    "id": str(state_value_display.id),
+                    "key": state_value_display.name or state_value.name,
+                    "type": infer_vellum_variable_type(state_value),
+                    "default": default.dict() if default else None,
+                    "required": required,
+                    "extensions": {"color": state_value_display.color},
+                }
+            )
+
         nodes: JsonArray = []
         edges: JsonArray = []
 
@@ -153,6 +176,16 @@ class VellumWorkflowDisplay(
                     except IndexError:
                         source_node_display = None
 
+                synthetic_target_handle_id = (
+                    str(workflow_output_display.target_handle_id)
+                    if workflow_output_display.target_handle_id
+                    else str(uuid4_from_hash(f"{self.workflow_id}|target_handle_id|{workflow_output_display.name}"))
+                )
+                synthetic_display_data = (
+                    workflow_output_display.display_data.dict()
+                    if workflow_output_display.display_data
+                    else NodeDisplayData().dict()
+                )
                 nodes.append(
                     {
                         "id": str(final_output_node_id),
@@ -160,13 +193,13 @@ class VellumWorkflowDisplay(
                         "data": {
                             "label": workflow_output_display.label,
                             "name": workflow_output_display.name,
-                            "target_handle_id": str(workflow_output_display.target_handle_id),
+                            "target_handle_id": synthetic_target_handle_id,
                             "output_id": str(workflow_output_display.id),
                             "output_type": inferred_type,
                             "node_input_id": str(node_input.id),
                         },
                         "inputs": [node_input.dict()],
-                        "display_data": workflow_output_display.display_data.dict(),
+                        "display_data": synthetic_display_data,
                         "base": final_output_node_base,
                         "definition": None,
                     }
@@ -188,7 +221,7 @@ class VellumWorkflowDisplay(
                             "source_node_id": str(source_node_display.node_id),
                             "source_handle_id": str(source_handle_id),
                             "target_node_id": str(workflow_output_display.node_id),
-                            "target_handle_id": str(workflow_output_display.target_handle_id),
+                            "target_handle_id": synthetic_target_handle_id,
                             "type": "DEFAULT",
                         }
                     )
@@ -244,6 +277,7 @@ class VellumWorkflowDisplay(
                 },
             },
             "input_variables": input_variables,
+            "state_variables": state_variables,
             "output_variables": output_variables,
         }
 
@@ -283,6 +317,23 @@ class VellumWorkflowDisplay(
 
         return WorkflowInputsVellumDisplay(id=workflow_input_id, name=name, required=required, color=color)
 
+    def _generate_state_value_display(
+        self, state_value: BaseDescriptor, overrides: Optional[StateValueVellumDisplayOverrides] = None
+    ) -> StateValueVellumDisplay:
+        state_value_id: UUID
+        name = None
+        required = None
+        color = None
+        if overrides:
+            state_value_id = overrides.id
+            name = overrides.name
+            required = overrides.required
+            color = overrides.color
+        else:
+            state_value_id = uuid4_from_hash(f"{self.workflow_id}|state_values|id|{state_value.name}")
+
+        return StateValueVellumDisplay(id=state_value_id, name=name, required=required, color=color)
+
     def _generate_entrypoint_display(
         self,
         entrypoint: Type[BaseNode],
@@ -331,15 +382,12 @@ class VellumWorkflowDisplay(
 
         output_id = uuid4_from_hash(f"{self.workflow_id}|id|{output.name}")
         node_id = uuid4_from_hash(f"{self.workflow_id}|node_id|{output.name}")
-        target_handle_id = uuid4_from_hash(f"{self.workflow_id}|target_handle_id|{output.name}")
 
         return WorkflowOutputVellumDisplay(
             id=output_id,
             node_id=node_id,
             name=output.name,
             label="Final Output",
-            target_handle_id=target_handle_id,
-            display_data=NodeDisplayData(),
         )
 
     def _generate_edge_display(

vellum_ee/workflows/tests/local_workflow/display/workflow.py

@@ -50,7 +50,5 @@ class WorkflowDisplay(VellumWorkflowDisplay[Workflow]):
             node_id=UUID("f3ef4b2b-fec9-4026-9cc6-e5eac295307f"),
             name="final-output",
             label="Final Output",
-            target_handle_id=UUID("3ec34f6e-da48-40d5-a65b-a48fefa75763"),
-            display_data=NodeDisplayData(position=NodeDisplayPosition(x=2750, y=210), width=459, height=234),
         )
     }