vellum-ai 0.14.63__py3-none-any.whl → 1.13.5__py3-none-any.whl

vellum/workflows/nodes/mocks.py

@@ -1,242 +1,107 @@
-from functools import reduce
-from uuid import UUID
-from typing import TYPE_CHECKING, Any, List, Literal, Optional, Sequence, Type, Union
+import logging
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Sequence, Type, Union
 
-from pydantic import ConfigDict, ValidationError
+from pydantic import ConfigDict, SerializationInfo, ValidationError, field_serializer, model_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.client.types.array_vellum_value import ArrayVellumValue
-from vellum.client.types.vellum_value import VellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import WorkflowErrorCode
-from vellum.workflows.exceptions import WorkflowInitializationException
+from vellum.workflows.events.types import default_serializer
+from vellum.workflows.exceptions import NodeException, WorkflowInitializationException
+from vellum.workflows.nodes.utils import get_unadorned_node
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
 
 if TYPE_CHECKING:
     from vellum.workflows import BaseWorkflow
 
-import logging
-
 logger = logging.getLogger(__name__)
 
 
-class _RawLogicalCondition(UniversalBaseModel):
-    type: Literal["LOGICAL_CONDITION"] = "LOGICAL_CONDITION"
-    lhs_variable_id: UUID
-    operator: Literal["==", ">", ">=", "<", "<=", "!="]
-    rhs_variable_id: UUID
-
-
-class _RawLogicalConditionGroup(UniversalBaseModel):
-    type: Literal["LOGICAL_CONDITION_GROUP"] = "LOGICAL_CONDITION_GROUP"
-    conditions: List["_RawLogicalExpression"]
-    combinator: Literal["AND", "OR"]
-    negated: bool
-
-
-_RawLogicalExpression = Union[_RawLogicalCondition, _RawLogicalConditionGroup]
-
-
-class _RawLogicalExpressionVariable(UniversalBaseModel):
-    id: UUID
-
-
-class _RawMockWorkflowNodeExecutionConstantValuePointer(_RawLogicalExpressionVariable):
-    type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
-    variable_value: VellumValue
-
-
-class _RawMockWorkflowNodeExecutionNodeExecutionCounterPointer(_RawLogicalExpressionVariable):
-    type: Literal["EXECUTION_COUNTER"] = "EXECUTION_COUNTER"
-    node_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionInputVariablePointer(_RawLogicalExpressionVariable):
-    type: Literal["INPUT_VARIABLE"] = "INPUT_VARIABLE"
-    input_variable_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionNodeOutputPointer(_RawLogicalExpressionVariable):
-    type: Literal["NODE_OUTPUT"] = "NODE_OUTPUT"
-    node_id: UUID
-    input_id: UUID
-
-
-class _RawMockWorkflowNodeExecutionNodeInputPointer(_RawLogicalExpressionVariable):
-    type: Literal["NODE_INPUT"] = "NODE_INPUT"
-    node_id: UUID
-    input_id: UUID
-
-
-_RawMockWorkflowNodeExecutionValuePointer = Union[
-    _RawMockWorkflowNodeExecutionConstantValuePointer,
-    _RawMockWorkflowNodeExecutionNodeExecutionCounterPointer,
-    _RawMockWorkflowNodeExecutionInputVariablePointer,
-    _RawMockWorkflowNodeExecutionNodeOutputPointer,
-    _RawMockWorkflowNodeExecutionNodeInputPointer,
-]
-
-
-class _RawMockWorkflowNodeWhenCondition(UniversalBaseModel):
-    expression: _RawLogicalExpression
-    variables: List[_RawMockWorkflowNodeExecutionValuePointer]
-
-
-class _RawMockWorkflowNodeThenOutput(UniversalBaseModel):
-    output_id: UUID
-    value: _RawMockWorkflowNodeExecutionValuePointer
-
-
-class _RawMockWorkflowNodeExecution(UniversalBaseModel):
-    when_condition: _RawMockWorkflowNodeWhenCondition
-    then_outputs: List[_RawMockWorkflowNodeThenOutput]
-
-
-class _RawMockWorkflowNodeConfig(UniversalBaseModel):
-    type: Literal["WORKFLOW_NODE_OUTPUT"] = "WORKFLOW_NODE_OUTPUT"
-    node_id: UUID
-    mock_executions: List[_RawMockWorkflowNodeExecution]
-
-
 class MockNodeExecution(UniversalBaseModel):
     when_condition: BaseDescriptor
     then_outputs: BaseOutputs
+    disabled: Optional[bool] = None
 
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
+    @model_serializer(mode="wrap")
+    def serialize_full_model(self, handler: Callable[[Any], Any], info: SerializationInfo) -> Dict[str, Any]:
+        """Serialize the model and add node_id field computed from then_outputs."""
+        serialized = handler(self)
+        node_class = self.then_outputs.__class__.__parent_class__
+        unadorned_node = get_unadorned_node(node_class)  # type: ignore[arg-type]
+        serialized["node_id"] = str(unadorned_node.__id__)
+        serialized["type"] = "NODE_EXECUTION"
+        if self.disabled is None:
+            del serialized["disabled"]
+        return serialized
+
+    @field_serializer("then_outputs")
+    def serialize_then_outputs(self, then_outputs: BaseOutputs, _info: SerializationInfo) -> Dict[str, Any]:
+        return default_serializer(then_outputs)
+
+    @field_serializer("when_condition")
+    def serialize_when_condition(self, when_condition: BaseDescriptor, info: SerializationInfo) -> Dict[str, Any]:
+        # This allows `ee` to pass in workflow display context through the pydantic context
+        if isinstance(info.context, dict) and "serializer" in info.context and callable(info.context["serializer"]):
+            return info.context["serializer"](when_condition)
+
+        return default_serializer(when_condition)
+
     @staticmethod
     def validate_all(
         raw_mock_workflow_node_configs: Optional[List[Any]],
         workflow: Type["BaseWorkflow"],
+        descriptor_validator: Optional[Callable[[dict, Type["BaseWorkflow"]], BaseDescriptor]] = None,
     ) -> Optional[List["MockNodeExecution"]]:
         if not raw_mock_workflow_node_configs:
             return None
 
         ArrayVellumValue.model_rebuild()
-        try:
-            mock_workflow_node_configs = [
-                _RawMockWorkflowNodeConfig.model_validate(raw_mock_workflow_node_config)
-                for raw_mock_workflow_node_config in raw_mock_workflow_node_configs
-            ]
-        except ValidationError as e:
-            raise WorkflowInitializationException(
-                message="Failed to validate mock node executions",
-                code=WorkflowErrorCode.INVALID_INPUTS,
-            ) from e
-
-        nodes = {node.__id__: node for node in workflow.get_nodes()}
-        node_output_name_by_id = {
-            node.__output_ids__[output.name]: output.name for node in workflow.get_nodes() for output in node.Outputs
-        }
-
-        # We need to support the old way that the Vellum App's WorkflowRunner used to define Node Mocks in order to
-        # avoid needing to update the mock resolution strategy that it and the frontend uses. The path towards
-        # cleaning this up will go as follows:
-        # 1. Release Mock support in SDK-Enabled Workflows
-        # 2. Deprecate Mock support in non-SDK enabled Workflows, encouraging users to migrate to SDK-enabled Workflows
-        # 3. Remove the old mock resolution strategy
-        # 4. Update this SDK to handle the new mock resolution strategy with WorkflowValueDescriptors
-        # 5. Cutover the Vellum App to the new mock resolution strategy
-        # 6. Remove the old mock resolution strategy from this SDK
-        def _translate_raw_logical_expression(
-            raw_logical_expression: _RawLogicalExpression,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            if raw_logical_expression.type == "LOGICAL_CONDITION":
-                return _translate_raw_logical_condition(raw_logical_expression, raw_variables)
-            else:
-                return _translate_raw_logical_condition_group(raw_logical_expression, raw_variables)
-
-        def _translate_raw_logical_condition_group(
-            raw_logical_condition_group: _RawLogicalConditionGroup,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            if not raw_logical_condition_group.conditions:
-                return ConstantValueReference(True)
-
-            conditions = [
-                _translate_raw_logical_expression(condition, raw_variables)
-                for condition in raw_logical_condition_group.conditions
-            ]
-            return reduce(
-                lambda acc, condition: (
-                    acc and condition if raw_logical_condition_group.combinator == "AND" else acc or condition
-                ),
-                conditions,
-            )
-
-        def _translate_raw_logical_condition(
-            raw_logical_condition: _RawLogicalCondition,
-            raw_variables: List[_RawMockWorkflowNodeExecutionValuePointer],
-        ) -> BaseDescriptor:
-            variable_by_id = {v.id: v for v in raw_variables}
-            lhs = _translate_raw_logical_expression_variable(variable_by_id[raw_logical_condition.lhs_variable_id])
-            rhs = _translate_raw_logical_expression_variable(variable_by_id[raw_logical_condition.rhs_variable_id])
-            if raw_logical_condition.operator == ">":
-                return lhs.greater_than(rhs)
-            elif raw_logical_condition.operator == ">=":
-                return lhs.greater_than_or_equal_to(rhs)
-            elif raw_logical_condition.operator == "<":
-                return lhs.less_than(rhs)
-            elif raw_logical_condition.operator == "<=":
-                return lhs.less_than_or_equal_to(rhs)
-            elif raw_logical_condition.operator == "==":
-                return lhs.equals(rhs)
-            elif raw_logical_condition.operator == "!=":
-                return lhs.does_not_equal(rhs)
-            else:
-                raise WorkflowInitializationException(f"Unsupported logical operator: {raw_logical_condition.operator}")
-
-        def _translate_raw_logical_expression_variable(
-            raw_variable: _RawMockWorkflowNodeExecutionValuePointer,
-        ) -> BaseDescriptor:
-            if raw_variable.type == "CONSTANT_VALUE":
-                return ConstantValueReference(raw_variable.variable_value.value)
-            elif raw_variable.type == "EXECUTION_COUNTER":
-                node = nodes[raw_variable.node_id]
-                return node.Execution.count
-            else:
-                raise WorkflowInitializationException(f"Unsupported logical expression type: {raw_variable.type}")
-
-        mock_node_executions = []
-        for mock_workflow_node_config in mock_workflow_node_configs:
-            for mock_execution in mock_workflow_node_config.mock_executions:
-                try:
-                    when_condition = _translate_raw_logical_expression(
-                        mock_execution.when_condition.expression,
-                        mock_execution.when_condition.variables,
+        mock_node_executions: list[MockNodeExecution] = []
+        for raw_mock_workflow_node_config in raw_mock_workflow_node_configs:
+            if not isinstance(raw_mock_workflow_node_config, dict):
+                raise WorkflowInitializationException(
+                    message=f"Invalid mock node execution type: {type(raw_mock_workflow_node_config)}",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    workflow_definition=workflow,
+                )
+
+            try:
+                mock_node_executions.append(
+                    MockNodeExecution.model_validate(
+                        raw_mock_workflow_node_config,
+                        context={
+                            "workflow": workflow,
+                            "node_id": raw_mock_workflow_node_config.get("node_id"),
+                            "descriptor_validator": lambda value: (
+                                descriptor_validator(value, workflow)
+                                if descriptor_validator
+                                else ConstantValueReference(False)
+                            ),
+                        },
                     )
-
-                    then_outputs = nodes[mock_workflow_node_config.node_id].Outputs()
-                    for then_output in mock_execution.then_outputs:
-                        node_output_name = node_output_name_by_id.get(then_output.output_id)
-                        if node_output_name is None:
-                            raise WorkflowInitializationException(
-                                f"Output {then_output.output_id} not found in node {mock_workflow_node_config.node_id}"
-                            )
-
-                        resolved_output_reference = _translate_raw_logical_expression_variable(then_output.value)
-                        if isinstance(resolved_output_reference, ConstantValueReference):
-                            setattr(
-                                then_outputs,
-                                node_output_name,
-                                resolved_output_reference._value,
-                            )
-                        else:
-                            raise WorkflowInitializationException(
-                                f"Unsupported resolved output reference type: {type(resolved_output_reference)}"
-                            )
-
-                    mock_node_executions.append(
-                        MockNodeExecution(
-                            when_condition=when_condition,
-                            then_outputs=then_outputs,
-                        )
+                )
+            except WorkflowInitializationException as e:
+                # If the node is not found in the workflow, skip it with a warning
+                node_id = raw_mock_workflow_node_config.get("node_id")
+                raw_data = e.raw_data or {}
+                if raw_data.get("node_ref") == node_id:
+                    logger.warning(
+                        "Skipping mock for node %s: node not found in workflow %s",
+                        node_id,
+                        workflow.__name__,
                     )
-                except Exception as e:
-                    logger.exception("Failed to validate mock node execution", exc_info=e)
                     continue
+                raise
+            except (ValidationError, NodeException) as e:
+                raise WorkflowInitializationException(
+                    message="Failed to validate mock node executions",
+                    code=WorkflowErrorCode.INVALID_INPUTS,
+                    workflow_definition=workflow,
+                ) from e
 
         return mock_node_executions
 

vellum/workflows/nodes/tests/test_mocks.py

@@ -1,98 +1,6 @@
-import uuid
-
-from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows import BaseWorkflow
-from vellum.workflows.nodes import InlinePromptNode
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.mocks import MockNodeExecution
-from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
-from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
-
-
-def test_mocks__parse_from_app():
-    # GIVEN a PromptNode
-    class PromptNode(InlinePromptNode):
-        pass
-
-    # AND a workflow class with that PromptNode
-    class MyWorkflow(BaseWorkflow):
-        graph = PromptNode
-
-    # AND a mock workflow node execution from the app
-    raw_mock_workflow_node_execution = [
-        {
-            "type": "WORKFLOW_NODE_OUTPUT",
-            "node_id": str(PromptNode.__id__),
-            "mock_executions": [
-                {
-                    "when_condition": {
-                        "expression": {
-                            "type": "LOGICAL_CONDITION_GROUP",
-                            "combinator": "AND",
-                            "negated": False,
-                            "conditions": [
-                                {
-                                    "type": "LOGICAL_CONDITION",
-                                    "lhs_variable_id": "e60902d5-6892-4916-80c1-f0130af52322",
-                                    "operator": ">=",
-                                    "rhs_variable_id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                                }
-                            ],
-                        },
-                        "variables": [
-                            {
-                                "type": "EXECUTION_COUNTER",
-                                "node_id": str(PromptNode.__id__),
-                                "id": "e60902d5-6892-4916-80c1-f0130af52322",
-                            },
-                            {
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "NUMBER", "value": 0},
-                                "id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                            },
-                        ],
-                    },
-                    "then_outputs": [
-                        {
-                            "output_id": "9e6dc5d3-8ea0-4346-8a2a-7cce5495755b",
-                            "value": {
-                                "id": "27006b2a-fa81-430c-a0b2-c66a9351fc68",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "STRING", "value": "Hello"},
-                            },
-                        },
-                        {
-                            "output_id": "60305ffd-60b0-42aa-b54e-4fdae0f8c28a",
-                            "value": {
-                                "id": "4559c778-6e27-4cfe-a460-734ba62a5082",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "ARRAY", "value": [{"type": "STRING", "value": "Hello"}]},
-                            },
-                        },
-                    ],
-                }
-            ],
-        }
-    ]
-
-    # WHEN we parse the mock workflow node execution
-    node_output_mocks = MockNodeExecution.validate_all(
-        raw_mock_workflow_node_execution,
-        MyWorkflow,
-    )
-
-    # THEN we get a list of MockNodeExecution objects
-    assert node_output_mocks
-    assert len(node_output_mocks) == 1
-    assert node_output_mocks[0] == MockNodeExecution(
-        when_condition=PromptNode.Execution.count.greater_than_or_equal_to(0.0),
-        then_outputs=PromptNode.Outputs(
-            text="Hello",
-            results=[
-                StringVellumValue(value="Hello"),
-            ],
-        ),
-    )
 
 
 def test_mocks__parse_none_still_runs():
@@ -120,88 +28,3 @@ def test_mocks__parse_none_still_runs():
 
     # THEN it was successful
     assert final_event.name == "workflow.execution.fulfilled"
-
-
-def test_mocks__use_id_from_display():
-    # GIVEN a Base Node
-    class StartNode(BaseNode):
-        class Outputs(BaseNode.Outputs):
-            foo: str
-
-    # AND a workflow class with that Node
-    class MyWorkflow(BaseWorkflow):
-        graph = StartNode
-
-        class Outputs(BaseWorkflow.Outputs):
-            final_value = StartNode.Outputs.foo
-
-    # AND a display class on that Base Node
-    node_output_id = uuid.uuid4()
-
-    class StartNodeDisplay(BaseNodeDisplay[StartNode]):
-        output_display = {StartNode.Outputs.foo: NodeOutputDisplay(id=node_output_id, name="foo")}
-
-    # AND a mock workflow node execution from the app
-    raw_mock_workflow_node_execution = [
-        {
-            "type": "WORKFLOW_NODE_OUTPUT",
-            "node_id": str(StartNode.__id__),
-            "mock_executions": [
-                {
-                    "when_condition": {
-                        "expression": {
-                            "type": "LOGICAL_CONDITION_GROUP",
-                            "combinator": "AND",
-                            "negated": False,
-                            "conditions": [
-                                {
-                                    "type": "LOGICAL_CONDITION",
-                                    "lhs_variable_id": "e60902d5-6892-4916-80c1-f0130af52322",
-                                    "operator": ">=",
-                                    "rhs_variable_id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                                }
-                            ],
-                        },
-                        "variables": [
-                            {
-                                "type": "EXECUTION_COUNTER",
-                                "node_id": str(StartNode.__id__),
-                                "id": "e60902d5-6892-4916-80c1-f0130af52322",
-                            },
-                            {
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "NUMBER", "value": 0},
-                                "id": "5c1bbb24-c288-49cb-a9b7-0c6f38a86037",
-                            },
-                        ],
-                    },
-                    "then_outputs": [
-                        {
-                            "output_id": str(node_output_id),
-                            "value": {
-                                "id": "27006b2a-fa81-430c-a0b2-c66a9351fc68",
-                                "type": "CONSTANT_VALUE",
-                                "variable_value": {"type": "STRING", "value": "Hello"},
-                            },
-                        },
-                    ],
-                }
-            ],
-        }
-    ]
-
-    # WHEN we parsed the raw data on `MockNodeExecution`
-    node_output_mocks = MockNodeExecution.validate_all(
-        raw_mock_workflow_node_execution,
-        MyWorkflow,
-    )
-
-    # THEN we get the expected list of MockNodeExecution objects
-    assert node_output_mocks
-    assert len(node_output_mocks) == 1
-    assert node_output_mocks[0] == MockNodeExecution(
-        when_condition=StartNode.Execution.count.greater_than_or_equal_to(0.0),
-        then_outputs=StartNode.Outputs(
-            foo="Hello",
-        ),
-    )

vellum/workflows/nodes/tests/test_utils.py

@@ -1,5 +1,5 @@
 import pytest
-from typing import Any, List, Union
+from typing import Any, Dict, List, Union
 
 from pydantic import BaseModel
 
@@ -34,6 +34,9 @@ class Person(BaseModel):
             Json,
             {"name": "test", "args": [1, 2]},
         ),
+        ('{"name": "Alice", "age": 30}', dict, {"name": "Alice", "age": 30}),
+        ('{"name": "Alice", "age": 30}', Dict[str, Any], {"name": "Alice", "age": 30}),
+        ('{"name": "Alice", "age": 30}', Any, {"name": "Alice", "age": 30}),
         ("42", Union[int, str], 42),
         ("hello", Union[int, str], "hello"),
     ],
@@ -47,6 +50,9 @@ class Person(BaseModel):
         "list_of_str",
         "simple_json",
         "function_call_json",
+        "dict_type",
+        "dict_str_any_type",
+        "any_type",
         "union_int",
         "union_str",
     ],
@@ -86,23 +92,23 @@ def test_parse_type_from_str_list_of_models():
             List[str],
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON Array format for result_as_str",
+            "Invalid JSON format: expected a valid JSON array",
         ),
         (
             "{invalid json}",
             Person,
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON format for result_as_str",
+            "Invalid JSON format: unable to parse the provided data",
         ),
         (
             "{invalid json}",
             Json,
             NodeException,
             WorkflowErrorCode.INVALID_OUTPUTS,
-            "Invalid JSON format for result_as_str",
+            "Invalid JSON format: unable to parse the provided data",
         ),
-        ('{"name": "Alice"}', List[str], ValueError, None, "Expected a list of items for result_as_str, received dict"),
+        ('{"name": "Alice"}', List[str], ValueError, None, "Expected a list of items, but received dict"),
         ("data", object, ValueError, None, "Unsupported output type: <class 'object'>"),
     ],
     ids=[
@@ -150,3 +156,26 @@ def test_cast_to_output_type_none_value(output_type, expected_result):
     """Test that cast_to_output_type returns appropriate default values when None is provided."""
     result = cast_to_output_type(None, output_type)
     assert result == expected_result
+
+
+@pytest.mark.parametrize(
+    "input_value,expected_result",
+    [
+        ('{"name": "Alice", "age": 30}', {"name": "Alice", "age": 30}),
+        ("[1, 2, 3]", [1, 2, 3]),
+        ("invalid json", "invalid json"),
+        ([1, 2, 3], [1, 2, 3]),
+        ({"already": "dict"}, {"already": "dict"}),
+    ],
+    ids=[
+        "valid_json_object",
+        "valid_json_array",
+        "invalid_json_string",
+        "non_string_list",
+        "non_string_dict",
+    ],
+)
+def test_cast_to_output_type_any_json_parsing(input_value, expected_result):
+    """Test that cast_to_output_type attempts JSON parsing for Any type and falls back gracefully."""
+    result = cast_to_output_type(input_value, Any)
+    assert result == expected_result