vellum-ai 0.14.0__py3-none-any.whl → 0.14.1__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 = []
 
@@ -244,6 +267,7 @@ class VellumWorkflowDisplay(
                 },
             },
             "input_variables": input_variables,
+            "state_variables": state_variables,
             "output_variables": output_variables,
         }
 
@@ -283,6 +307,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],