vellum-ai 0.14.44__py3-none-any.whl → 0.14.45__py3-none-any.whl

vellum_ee/workflows/display/utils/expressions.py

@@ -1,12 +1,98 @@
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from vellum.client.types.logical_operator import LogicalOperator
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.expressions.accessor import AccessorExpression
+from vellum.workflows.expressions.and_ import AndExpression
+from vellum.workflows.expressions.begins_with import BeginsWithExpression
+from vellum.workflows.expressions.between import BetweenExpression
+from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
+from vellum.workflows.expressions.contains import ContainsExpression
+from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
+from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
+from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
+from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
+from vellum.workflows.expressions.ends_with import EndsWithExpression
+from vellum.workflows.expressions.equals import EqualsExpression
+from vellum.workflows.expressions.greater_than import GreaterThanExpression
+from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
+from vellum.workflows.expressions.in_ import InExpression
+from vellum.workflows.expressions.is_nil import IsNilExpression
+from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
+from vellum.workflows.expressions.is_not_null import IsNotNullExpression
+from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
+from vellum.workflows.expressions.is_null import IsNullExpression
+from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
+from vellum.workflows.expressions.less_than import LessThanExpression
+from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
+from vellum.workflows.expressions.not_between import NotBetweenExpression
+from vellum.workflows.expressions.not_in import NotInExpression
+from vellum.workflows.expressions.or_ import OrExpression
+from vellum.workflows.expressions.parse_json import ParseJsonExpression
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
+from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.references.output import OutputReference
+from vellum.workflows.references.state_value import StateValueReference
+from vellum.workflows.references.vellum_secret import VellumSecretReference
+from vellum.workflows.references.workflow_input import WorkflowInputReference
+from vellum.workflows.types.core import JsonObject
+from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
 
 if TYPE_CHECKING:
     from vellum_ee.workflows.display.types import WorkflowDisplayContext
 
 
+def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
+    if isinstance(descriptor, EqualsExpression):
+        return "="
+    elif isinstance(descriptor, DoesNotEqualExpression):
+        return "!="
+    elif isinstance(descriptor, LessThanExpression):
+        return "<"
+    elif isinstance(descriptor, GreaterThanExpression):
+        return ">"
+    elif isinstance(descriptor, LessThanOrEqualToExpression):
+        return "<="
+    elif isinstance(descriptor, GreaterThanOrEqualToExpression):
+        return ">="
+    elif isinstance(descriptor, ContainsExpression):
+        return "contains"
+    elif isinstance(descriptor, BeginsWithExpression):
+        return "beginsWith"
+    elif isinstance(descriptor, EndsWithExpression):
+        return "endsWith"
+    elif isinstance(descriptor, DoesNotContainExpression):
+        return "doesNotContain"
+    elif isinstance(descriptor, DoesNotBeginWithExpression):
+        return "doesNotBeginWith"
+    elif isinstance(descriptor, DoesNotEndWithExpression):
+        return "doesNotEndWith"
+    elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
+        return "null"
+    elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
+        return "notNull"
+    elif isinstance(descriptor, InExpression):
+        return "in"
+    elif isinstance(descriptor, NotInExpression):
+        return "notIn"
+    elif isinstance(descriptor, BetweenExpression):
+        return "between"
+    elif isinstance(descriptor, NotBetweenExpression):
+        return "notBetween"
+    elif isinstance(descriptor, AndExpression):
+        return "and"
+    elif isinstance(descriptor, OrExpression):
+        return "or"
+    elif isinstance(descriptor, CoalesceExpression):
+        return "coalesce"
+    elif isinstance(descriptor, ParseJsonExpression):
+        return "parseJson"
+    else:
+        raise ValueError(f"Unsupported descriptor type: {descriptor}")
+
+
 def get_child_descriptor(value: LazyReference, display_context: "WorkflowDisplayContext") -> BaseDescriptor:
     if isinstance(value._get, str):
         reference_parts = value._get.split(".")
@@ -28,3 +114,137 @@ def get_child_descriptor(value: LazyReference, display_context: "WorkflowDisplay
         raise Exception(f"Failed to parse lazy reference: {value._get}")
 
     return value._get()
+
+
+def serialize_condition(display_context: "WorkflowDisplayContext", condition: BaseDescriptor) -> JsonObject:
+    if isinstance(
+        condition,
+        (
+            IsNullExpression,
+            IsNotNullExpression,
+            IsNilExpression,
+            IsNotNilExpression,
+            IsUndefinedExpression,
+            IsNotUndefinedExpression,
+            ParseJsonExpression,
+        ),
+    ):
+        lhs = serialize_value(display_context, condition._expression)
+        return {
+            "type": "UNARY_EXPRESSION",
+            "lhs": lhs,
+            "operator": convert_descriptor_to_operator(condition),
+        }
+    elif isinstance(condition, (BetweenExpression, NotBetweenExpression)):
+        base = serialize_value(display_context, condition._value)
+        lhs = serialize_value(display_context, condition._start)
+        rhs = serialize_value(display_context, condition._end)
+
+        return {
+            "type": "TERNARY_EXPRESSION",
+            "base": base,
+            "operator": convert_descriptor_to_operator(condition),
+            "lhs": lhs,
+            "rhs": rhs,
+        }
+    elif isinstance(
+        condition,
+        (
+            AndExpression,
+            BeginsWithExpression,
+            CoalesceExpression,
+            ContainsExpression,
+            DoesNotBeginWithExpression,
+            DoesNotContainExpression,
+            DoesNotEndWithExpression,
+            DoesNotEqualExpression,
+            EndsWithExpression,
+            EqualsExpression,
+            GreaterThanExpression,
+            GreaterThanOrEqualToExpression,
+            InExpression,
+            LessThanExpression,
+            LessThanOrEqualToExpression,
+            NotInExpression,
+            OrExpression,
+        ),
+    ):
+        lhs = serialize_value(display_context, condition._lhs)
+        rhs = serialize_value(display_context, condition._rhs)
+
+        return {
+            "type": "BINARY_EXPRESSION",
+            "lhs": lhs,
+            "operator": convert_descriptor_to_operator(condition),
+            "rhs": rhs,
+        }
+    elif isinstance(condition, AccessorExpression):
+        return {
+            "type": "BINARY_EXPRESSION",
+            "lhs": serialize_value(display_context, condition._base),
+            "operator": "accessField",
+            "rhs": serialize_value(display_context, condition._field),
+        }
+
+    raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
+
+
+def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
+    if isinstance(value, ConstantValueReference):
+        return serialize_value(display_context, value._value)
+
+    if isinstance(value, LazyReference):
+        child_descriptor = get_child_descriptor(value, display_context)
+        return serialize_value(display_context, child_descriptor)
+
+    if isinstance(value, WorkflowInputReference):
+        workflow_input_display = display_context.global_workflow_input_displays[value]
+        return {
+            "type": "WORKFLOW_INPUT",
+            "input_variable_id": str(workflow_input_display.id),
+        }
+
+    if isinstance(value, StateValueReference):
+        state_value_display = display_context.global_state_value_displays[value]
+        return {
+            "type": "STATE_VALUE",
+            "state_variable_id": str(state_value_display.id),
+        }
+
+    if isinstance(value, OutputReference):
+        upstream_node, output_display = display_context.global_node_output_displays[value]
+        upstream_node_display = display_context.global_node_displays[upstream_node]
+
+        return {
+            "type": "NODE_OUTPUT",
+            "node_id": str(upstream_node_display.node_id),
+            "node_output_id": str(output_display.id),
+        }
+
+    if isinstance(value, VellumSecretReference):
+        return {
+            "type": "VELLUM_SECRET",
+            "vellum_secret_name": value.name,
+        }
+
+    if isinstance(value, ExecutionCountReference):
+        node_class_display = display_context.global_node_displays[value.node_class]
+
+        return {
+            "type": "EXECUTION_COUNTER",
+            "node_id": str(node_class_display.node_id),
+        }
+
+    if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
+        raise ValueError("Nested references are not supported.")
+
+    if not isinstance(value, BaseDescriptor):
+        vellum_value = primitive_to_vellum_value(value)
+        return {
+            "type": "CONSTANT_VALUE",
+            "value": vellum_value.dict(),
+        }
+
+    # If it's not any of the references we know about,
+    # then try to serialize it as a nested value
+    return serialize_condition(display_context, value)

vellum_ee/workflows/display/utils/vellum.py

@@ -2,36 +2,9 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
-from vellum.client.types.logical_operator import LogicalOperator
 from vellum.client.types.vellum_value import VellumValue
 from vellum.client.types.vellum_variable_type import VellumVariableType
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.expressions.and_ import AndExpression
-from vellum.workflows.expressions.begins_with import BeginsWithExpression
-from vellum.workflows.expressions.between import BetweenExpression
-from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
-from vellum.workflows.expressions.contains import ContainsExpression
-from vellum.workflows.expressions.does_not_begin_with import DoesNotBeginWithExpression
-from vellum.workflows.expressions.does_not_contain import DoesNotContainExpression
-from vellum.workflows.expressions.does_not_end_with import DoesNotEndWithExpression
-from vellum.workflows.expressions.does_not_equal import DoesNotEqualExpression
-from vellum.workflows.expressions.ends_with import EndsWithExpression
-from vellum.workflows.expressions.equals import EqualsExpression
-from vellum.workflows.expressions.greater_than import GreaterThanExpression
-from vellum.workflows.expressions.greater_than_or_equal_to import GreaterThanOrEqualToExpression
-from vellum.workflows.expressions.in_ import InExpression
-from vellum.workflows.expressions.is_nil import IsNilExpression
-from vellum.workflows.expressions.is_not_nil import IsNotNilExpression
-from vellum.workflows.expressions.is_not_null import IsNotNullExpression
-from vellum.workflows.expressions.is_not_undefined import IsNotUndefinedExpression
-from vellum.workflows.expressions.is_null import IsNullExpression
-from vellum.workflows.expressions.is_undefined import IsUndefinedExpression
-from vellum.workflows.expressions.less_than import LessThanExpression
-from vellum.workflows.expressions.less_than_or_equal_to import LessThanOrEqualToExpression
-from vellum.workflows.expressions.not_between import NotBetweenExpression
-from vellum.workflows.expressions.not_in import NotInExpression
-from vellum.workflows.expressions.or_ import OrExpression
-from vellum.workflows.expressions.parse_json import ParseJsonExpression
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import OutputReference, WorkflowInputReference
@@ -166,52 +139,3 @@ def create_node_input_value_pointer_rule(
         return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
 
     raise UnsupportedSerializationException(f"Unsupported descriptor type: {value.__class__.__name__}")
-
-
-def convert_descriptor_to_operator(descriptor: BaseDescriptor) -> LogicalOperator:
-    if isinstance(descriptor, EqualsExpression):
-        return "="
-    elif isinstance(descriptor, DoesNotEqualExpression):
-        return "!="
-    elif isinstance(descriptor, LessThanExpression):
-        return "<"
-    elif isinstance(descriptor, GreaterThanExpression):
-        return ">"
-    elif isinstance(descriptor, LessThanOrEqualToExpression):
-        return "<="
-    elif isinstance(descriptor, GreaterThanOrEqualToExpression):
-        return ">="
-    elif isinstance(descriptor, ContainsExpression):
-        return "contains"
-    elif isinstance(descriptor, BeginsWithExpression):
-        return "beginsWith"
-    elif isinstance(descriptor, EndsWithExpression):
-        return "endsWith"
-    elif isinstance(descriptor, DoesNotContainExpression):
-        return "doesNotContain"
-    elif isinstance(descriptor, DoesNotBeginWithExpression):
-        return "doesNotBeginWith"
-    elif isinstance(descriptor, DoesNotEndWithExpression):
-        return "doesNotEndWith"
-    elif isinstance(descriptor, (IsNullExpression, IsNilExpression, IsUndefinedExpression)):
-        return "null"
-    elif isinstance(descriptor, (IsNotNullExpression, IsNotNilExpression, IsNotUndefinedExpression)):
-        return "notNull"
-    elif isinstance(descriptor, InExpression):
-        return "in"
-    elif isinstance(descriptor, NotInExpression):
-        return "notIn"
-    elif isinstance(descriptor, BetweenExpression):
-        return "between"
-    elif isinstance(descriptor, NotBetweenExpression):
-        return "notBetween"
-    elif isinstance(descriptor, AndExpression):
-        return "and"
-    elif isinstance(descriptor, OrExpression):
-        return "or"
-    elif isinstance(descriptor, CoalesceExpression):
-        return "coalesce"
-    elif isinstance(descriptor, ParseJsonExpression):
-        return "parseJson"
-    else:
-        raise ValueError(f"Unsupported descriptor type: {descriptor}")

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -46,6 +46,7 @@ from vellum_ee.workflows.display.types import (
     WorkflowInputsDisplays,
     WorkflowOutputDisplays,
 )
+from vellum_ee.workflows.display.utils.expressions import serialize_value
 from vellum_ee.workflows.display.utils.registry import register_workflow_display_class
 from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -124,26 +125,24 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 }
             )
 
-        nodes: JsonArray = []
+        serialized_nodes: Dict[UUID, JsonObject] = {}
         edges: JsonArray = []
 
         # Add a single synthetic node for the workflow entrypoint
         entrypoint_node_id = self.display_context.workflow_display.entrypoint_node_id
         entrypoint_node_source_handle_id = self.display_context.workflow_display.entrypoint_node_source_handle_id
-        nodes.append(
-            {
-                "id": str(entrypoint_node_id),
-                "type": "ENTRYPOINT",
-                "inputs": [],
-                "data": {
-                    "label": "Entrypoint Node",
-                    "source_handle_id": str(entrypoint_node_source_handle_id),
-                },
-                "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
-                "base": None,
-                "definition": None,
+        serialized_nodes[entrypoint_node_id] = {
+            "id": str(entrypoint_node_id),
+            "type": "ENTRYPOINT",
+            "inputs": [],
+            "data": {
+                "label": "Entrypoint Node",
+                "source_handle_id": str(entrypoint_node_source_handle_id),
             },
-        )
+            "display_data": self.display_context.workflow_display.entrypoint_node_display.dict(),
+            "base": None,
+            "definition": None,
+        }
 
         # Add all the nodes in the workflow
         for node in self._workflow.get_nodes():
@@ -155,7 +154,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.add_error(e)
                 continue
 
-            nodes.append(serialized_node)
+            serialized_nodes[node_display.node_id] = serialized_node
 
         # Add all unused nodes in the workflow
         for node in self._workflow.get_unused_nodes():
@@ -167,10 +166,11 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.add_error(e)
                 continue
 
-            nodes.append(serialized_node)
+            serialized_nodes[node_display.node_id] = serialized_node
 
         synthetic_output_edges: JsonArray = []
         output_variables: JsonArray = []
+        output_values: JsonArray = []
         final_output_nodes = [
             node for node in self.display_context.node_displays.keys() if issubclass(node, FinalOutputNode)
         ]
@@ -185,9 +185,9 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
         for workflow_output, workflow_output_display in self.display_context.workflow_output_displays.items():
             final_output_node_id = uuid4_from_hash(f"{self.workflow_id}|node_id|{workflow_output.name}")
             inferred_type = infer_vellum_variable_type(workflow_output)
-
             # Remove the terminal node output from the unreferenced set
-            unreferenced_final_output_node_outputs.discard(cast(OutputReference, workflow_output.instance))
+            if isinstance(workflow_output.instance, OutputReference):
+                unreferenced_final_output_node_outputs.discard(workflow_output.instance)
 
             if workflow_output.instance not in final_output_node_outputs:
                 # Create a synthetic terminal node only if there is no terminal node for this output
@@ -220,24 +220,22 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 )
                 synthetic_display_data = NodeDisplayData().dict()
                 synthetic_node_label = "Final Output"
-                nodes.append(
-                    {
-                        "id": str(final_output_node_id),
-                        "type": "TERMINAL",
-                        "data": {
-                            "label": synthetic_node_label,
-                            "name": workflow_output_display.name,
-                            "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": synthetic_display_data,
-                        "base": final_output_node_base,
-                        "definition": None,
-                    }
-                )
+                serialized_nodes[final_output_node_id] = {
+                    "id": str(final_output_node_id),
+                    "type": "TERMINAL",
+                    "data": {
+                        "label": synthetic_node_label,
+                        "name": workflow_output_display.name,
+                        "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": synthetic_display_data,
+                    "base": final_output_node_base,
+                    "definition": None,
+                }
 
                 if source_node_display:
                     source_handle_id = source_node_display.get_source_handle_id(
@@ -255,6 +253,19 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                         }
                     )
 
+            elif isinstance(workflow_output.instance, OutputReference):
+                terminal_node_id = workflow_output.instance.outputs_class._node_class.__id__
+                serialized_terminal_node = serialized_nodes.get(terminal_node_id)
+                if serialized_terminal_node and isinstance(serialized_terminal_node["data"], dict):
+                    serialized_terminal_node["data"]["name"] = workflow_output_display.name
+
+            output_values.append(
+                {
+                    "output_variable_id": str(workflow_output_display.id),
+                    "value": serialize_value(self.display_context, workflow_output.instance),
+                }
+            )
+
             output_variables.append(
                 {
                     "id": str(workflow_output_display.id),
@@ -309,13 +320,14 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
         return {
             "workflow_raw_data": {
-                "nodes": nodes,
+                "nodes": list(serialized_nodes.values()),
                 "edges": edges,
                 "display_data": self.display_context.workflow_display.display_data.dict(),
                 "definition": {
                     "name": self._workflow.__name__,
                     "module": cast(JsonArray, self._workflow.__module__.split(".")),
                 },
+                "output_values": output_values,
             },
             "input_variables": input_variables,
             "state_variables": state_variables,

vellum_ee/workflows/display/workflows/tests/test_workflow_display.py

@@ -6,6 +6,7 @@ from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubwo
 from vellum.workflows.nodes.core.retry_node.node import RetryNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 from vellum.workflows.nodes.core.try_node.node import TryNode
+from vellum.workflows.nodes.displayable.final_output_node.node import FinalOutputNode
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.editor.types import NodeDisplayData, NodeDisplayPosition
 from vellum_ee.workflows.display.nodes import BaseNodeDisplay
@@ -327,3 +328,47 @@ def test_serialize_workflow__inherited_workflow_display_class_not_registered():
 
     # THEN it should should succeed
     assert data is not None
+
+
+def test_serialize_workflow__terminal_node_mismatches_workflow_output_name():
+    # GIVEN a node
+    class ExitNode(FinalOutputNode):
+        class Outputs(FinalOutputNode.Outputs):
+            value = "hello"
+
+    # AND a workflow that uses the node
+    class MyWorkflow(BaseWorkflow):
+        graph = ExitNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            answer = ExitNode.Outputs.value
+
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=MyWorkflow)
+    data = workflow_display.serialize()
+
+    # THEN it should have an output name that matches the workflow output
+    assert isinstance(data["workflow_raw_data"], dict)
+    assert isinstance(data["workflow_raw_data"]["nodes"], list)
+    terminal_node = [
+        node for node in data["workflow_raw_data"]["nodes"] if isinstance(node, dict) and node["type"] == "TERMINAL"
+    ][0]
+    assert isinstance(terminal_node["data"], dict)
+    assert terminal_node["data"]["name"] == "answer"
+
+    # AND the output variable should have the correct name
+    assert isinstance(data["output_variables"], list)
+    assert isinstance(data["output_variables"][0], dict)
+    assert data["output_variables"][0]["key"] == "answer"
+    assert data["output_variables"][0]["type"] == "STRING"
+
+    # AND the output value should have the correct name
+    output_variable_id = data["output_variables"][0]["id"]
+    assert isinstance(data["workflow_raw_data"]["output_values"], list)
+    assert isinstance(data["workflow_raw_data"]["output_values"][0], dict)
+    assert data["workflow_raw_data"]["output_values"][0]["output_variable_id"] == output_variable_id
+    assert data["workflow_raw_data"]["output_values"][0]["value"] == {
+        "type": "NODE_OUTPUT",
+        "node_id": str(ExitNode.__id__),
+        "node_output_id": str(ExitNode.__output_ids__["value"]),
+    }