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

vellum_ee/workflows/display/nodes/vellum/tests/test_utils.py

@@ -1,4 +1,5 @@
 import pytest
+from unittest.mock import patch
 from uuid import UUID, uuid4
 from typing import List
 
@@ -8,11 +9,15 @@ from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.references import LazyReference
-from vellum_ee.workflows.display.base import WorkflowInputsDisplay, WorkflowMetaDisplay
+from vellum.workflows.references.state_value import StateValueReference
+from vellum.workflows.references.trigger import TriggerAttributeReference
+from vellum.workflows.state import BaseState
+from vellum.workflows.triggers.base import BaseTrigger
+from vellum_ee.workflows.display.base import StateValueDisplay, WorkflowInputsDisplay, WorkflowMetaDisplay
 from vellum_ee.workflows.display.editor.types import NodeDisplayData
 from vellum_ee.workflows.display.nodes.base_node_display import BaseNodeDisplay
 from vellum_ee.workflows.display.nodes.types import NodeOutputDisplay
-from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input_value_pointer_rules
+from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input, create_node_input_value_pointer_rules
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.utils.vellum import (
     ConstantValuePointer,
@@ -21,6 +26,11 @@ from vellum_ee.workflows.display.utils.vellum import (
     NodeInputValuePointerRule,
     NodeOutputData,
     NodeOutputPointer,
+    TriggerAttributeData,
+    TriggerAttributePointer,
+    WorkflowStateData,
+    WorkflowStatePointer,
+    create_node_input_value_pointer_rule,
 )
 from vellum_ee.workflows.display.workflows.base_workflow_display import BaseWorkflowDisplay
 
@@ -129,5 +139,100 @@ def test_create_node_input_value_pointer_rules(
                 MyNodeA: MyNodeADisplay(),
             },
         ),
+        uuid4(),
     )
     assert rules == expected_rules
+
+
+class MyState(BaseState):
+    my_attribute: str
+
+
+def test_create_node_input_value_pointer_rule__state_value_reference() -> None:
+    """
+    Tests that StateValueReference is serialized to WorkflowStatePointer using the display override ID.
+    """
+
+    # GIVEN a StateValueReference
+    state_value_reference: StateValueReference[str] = MyState.my_attribute  # type: ignore[assignment]
+
+    # AND a display context with a state value display override
+    override_id = uuid4()
+    display_context = WorkflowDisplayContext(
+        global_state_value_displays={
+            state_value_reference: StateValueDisplay(id=override_id),
+        },
+    )
+
+    # WHEN we create a node input value pointer rule
+    result = create_node_input_value_pointer_rule(state_value_reference, display_context)
+
+    # THEN we should get a WorkflowStatePointer with the overridden display ID
+    assert isinstance(result, WorkflowStatePointer)
+    assert result.type == "WORKFLOW_STATE"
+    assert isinstance(result.data, WorkflowStateData)
+    assert result.data.state_variable_id == str(override_id)
+
+
+class MyTrigger(BaseTrigger):
+    my_attribute: str
+
+
+def test_create_node_input_value_pointer_rule__trigger_attribute_reference() -> None:
+    """
+    Tests that TriggerAttributeReference is serialized to TriggerAttributePointer.
+    """
+
+    # GIVEN a TriggerAttributeReference
+    trigger_attribute_reference: TriggerAttributeReference[str] = MyTrigger.my_attribute  # type: ignore[assignment]
+
+    # AND a display context
+    display_context = WorkflowDisplayContext()
+
+    # WHEN we create a node input value pointer rule
+    result = create_node_input_value_pointer_rule(trigger_attribute_reference, display_context)
+
+    # THEN we should get a TriggerAttributePointer with the correct data
+    assert isinstance(result, TriggerAttributePointer)
+    assert result.type == "TRIGGER_ATTRIBUTE"
+    assert isinstance(result.data, TriggerAttributeData)
+    assert result.data.trigger_id == str(MyTrigger.__id__)
+    assert result.data.attribute_id == str(trigger_attribute_reference.id)
+
+
+def test_create_node_input__unexpected_error_returns_node_input_with_empty_rules() -> None:
+    """
+    Tests that when create_node_input_value_pointer_rules raises an unexpected error,
+    create_node_input still returns a NodeInput with empty rules and adds the error to display context.
+    """
+
+    # GIVEN a display context
+    display_context = WorkflowDisplayContext()
+
+    # AND a node_id and input_name
+    node_id = uuid4()
+    input_name = "test_input"
+
+    # AND create_node_input_value_pointer_rules will raise an unexpected error
+    unexpected_error = RuntimeError("Unexpected error during serialization")
+    with patch(
+        "vellum_ee.workflows.display.nodes.vellum.utils.create_node_input_value_pointer_rules",
+        side_effect=unexpected_error,
+    ):
+        # WHEN we call create_node_input
+        result = create_node_input(
+            node_id=node_id,
+            input_name=input_name,
+            value="test_value",
+            display_context=display_context,
+        )
+
+    # THEN we should still get a NodeInput with empty rules
+    assert result.key == input_name
+    assert result.value.rules == []
+    assert result.value.combinator == "OR"
+
+    # AND the error should be added to the display context
+    errors = list(display_context.errors)
+    assert len(errors) == 1
+    assert errors[0] is unexpected_error

vellum_ee/workflows/display/nodes/vellum/utils.py

@@ -1,6 +1,8 @@
+import logging
 from uuid import UUID
 from typing import Any, List, Optional, Type, Union, cast
 
+from vellum.client.types.json_vellum_value import JsonVellumValue
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.expressions.coalesce_expression import CoalesceExpression
 from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
@@ -8,8 +10,12 @@ from vellum.workflows.references import NodeReference
 from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
-from vellum_ee.workflows.display.utils.exceptions import UnsupportedSerializationException
-from vellum_ee.workflows.display.utils.expressions import get_child_descriptor
+from vellum_ee.workflows.display.utils.exceptions import (
+    InvalidOutputReferenceError,
+    UnsupportedSerializationException,
+    UserFacingException,
+)
+from vellum_ee.workflows.display.utils.expressions import get_child_descriptor, serialize_value
 from vellum_ee.workflows.display.utils.vellum import (
     ConstantValuePointer,
     ExecutionCounterData,
@@ -23,6 +29,8 @@ from vellum_ee.workflows.display.utils.vellum import (
 )
 from vellum_ee.workflows.display.vellum import NodeInput, NodeInputValuePointer
 
+logger = logging.getLogger(__name__)
+
 
 def create_node_input(
     node_id: UUID,
@@ -33,7 +41,13 @@ def create_node_input(
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = ConstantValuePointer,
 ) -> NodeInput:
     input_id = str(input_id) if input_id else str(uuid4_from_hash(f"{node_id}|{input_name}"))
-    rules = create_node_input_value_pointer_rules(value, display_context, pointer_type=pointer_type)
+    try:
+        rules = create_node_input_value_pointer_rules(value, display_context, node_id, pointer_type=pointer_type)
+    except UserFacingException:
+        raise
+    except Exception as e:
+        display_context.add_validation_error(e)
+        rules = []
     return NodeInput(
         id=input_id,
         key=input_name,
@@ -44,9 +58,21 @@ def create_node_input(
     )
 
 
+def _contains_descriptor(value: Any) -> bool:
+    """Check if a value contains any BaseDescriptor objects."""
+    if isinstance(value, BaseDescriptor):
+        return True
+    if isinstance(value, dict):
+        return any(_contains_descriptor(v) for v in value.values())
+    if isinstance(value, list):
+        return any(_contains_descriptor(item) for item in value)
+    return False
+
+
 def create_node_input_value_pointer_rules(
     value: Any,
     display_context: WorkflowDisplayContext,
+    node_id: UUID,
     existing_rules: Optional[List[NodeInputValuePointerRule]] = None,
     pointer_type: Optional[Type[NodeInputValuePointerRule]] = None,
 ) -> List[NodeInputValuePointerRule]:
@@ -59,33 +85,49 @@ def create_node_input_value_pointer_rules(
             value = cast(BaseDescriptor, value.instance)
 
         if isinstance(value, LazyReference):
-            child_descriptor = get_child_descriptor(value, display_context)
+            try:
+                child_descriptor = get_child_descriptor(value, display_context)
+            except InvalidOutputReferenceError as e:
+                logger.warning("Failed to parse lazy reference '%s', skipping serialization", value.name)
+                display_context.add_validation_error(e)
+                return node_input_value_pointer_rules
             return create_node_input_value_pointer_rules(
-                child_descriptor, display_context, [], pointer_type=pointer_type
+                child_descriptor, display_context, node_id, existing_rules=[], pointer_type=pointer_type
             )
 
         if isinstance(value, CoalesceExpression):
-            # Recursively handle the left-hand side
-            lhs_rules = create_node_input_value_pointer_rules(value.lhs, display_context, [], pointer_type=pointer_type)
+            lhs_rules = create_node_input_value_pointer_rules(
+                value.lhs, display_context, node_id, existing_rules=[], pointer_type=pointer_type
+            )
             node_input_value_pointer_rules.extend(lhs_rules)
 
-            # Handle the right-hand side
             if not isinstance(value.rhs, CoalesceExpression):
                 rhs_rules = create_node_input_value_pointer_rules(
-                    value.rhs, display_context, [], pointer_type=pointer_type
+                    value.rhs, display_context, node_id, existing_rules=[], pointer_type=pointer_type
                 )
                 node_input_value_pointer_rules.extend(rhs_rules)
         else:
-            # Non-CoalesceExpression case
             try:
                 rule = create_node_input_value_pointer_rule(value, display_context)
             except UnsupportedSerializationException:
                 return node_input_value_pointer_rules
 
             node_input_value_pointer_rules.append(rule)
+    elif isinstance(value, (dict, list)) and _contains_descriptor(value):
+        display_context.add_validation_error(
+            UnsupportedSerializationException(
+                "The Vellum UI does not support nested references for this Node attribute. "
+                "Consider flattening the references."
+            )
+        )
+        serialized = serialize_value(node_id, display_context, value)
+        serialized_pointer = ConstantValuePointer(
+            data=JsonVellumValue(value=serialized),
+        )
+        node_input_value_pointer_rules.append(serialized_pointer)
     else:
-        pointer = create_pointer(value, pointer_type)
-        node_input_value_pointer_rules.append(pointer)
+        constant_pointer = create_pointer(value, pointer_type)
+        node_input_value_pointer_rules.append(constant_pointer)
 
     return node_input_value_pointer_rules
 

vellum_ee/workflows/display/tests/test_base_workflow_display.py

@@ -92,6 +92,7 @@ def test_vellum_workflow_display__serialize_input_variables_with_capitalized_var
             "default": None,
             "required": True,
             "extensions": {"color": None},
+            "schema": {"type": "string"},
         }
     ]
 
@@ -129,24 +130,19 @@ def test_vellum_workflow_display_serialize_valid_handle_ids_for_base_nodes():
     edge_source_handle_ids = {edge.get("source_handle_id") for edge in edges if isinstance(edge, dict)}
     edge_target_handle_ids = {edge.get("target_handle_id") for edge in edges if isinstance(edge, dict)}
 
-    for node in nodes:
-        assert isinstance(node, dict)
-
-        if node["type"] in {"ENTRYPOINT", "TERMINAL"}:
-            continue
+    start_node = next(
+        node for node in nodes if isinstance(node, dict) and node["type"] == "GENERIC" and node["label"] == "Start Node"
+    )
+    end_node = next(
+        node for node in nodes if isinstance(node, dict) and node["type"] == "GENERIC" and node["label"] == "End Node"
+    )
 
-        ports = node.get("ports")
-        assert isinstance(ports, list)
-        for port in ports:
-            assert isinstance(port, dict)
-            assert (
-                port["id"] in edge_source_handle_ids
-            ), f"Port {port['id']} from node {node['label']} not found in edge source handle ids"
+    assert isinstance(start_node["ports"], list)
+    assert isinstance(start_node["ports"][0], dict)
+    assert start_node["ports"][0]["id"] in edge_source_handle_ids
 
-        assert isinstance(node["trigger"], dict)
-        assert (
-            node["trigger"]["id"] in edge_target_handle_ids
-        ), f"Trigger {node['trigger']['id']} from node {node['label']} not found in edge target handle ids"
+    assert isinstance(end_node["trigger"], dict)
+    assert end_node["trigger"]["id"] in edge_target_handle_ids
 
 
 def test_vellum_workflow_display__serialize_with_unused_nodes_and_edges():
@@ -280,6 +276,7 @@ def test_vellum_workflow_display__serialize_with_parse_json_expression():
         "id": "e73fd6b1-1109-4a97-8510-c9ba8e6f5dbe",
         "name": "json_result",
         "type": "JSON",
+        "schema": {},
         "value": {
             "type": "UNARY_EXPRESSION",
             "lhs": {

vellum_ee/workflows/display/tests/test_json_schema_validation.py

@@ -0,0 +1,190 @@
+"""
+Tests for JSON schema validation in InlinePromptNode during serialization.
+
+This uses the jsonschema library to validate schemas against the JSON Schema meta-schema.
+The validation is spec-compliant, meaning some schemas that might seem incomplete
+(like arrays without items) are valid according to the JSON Schema specification.
+"""
+
+import pytest
+
+import jsonschema
+
+from vellum import PromptParameters
+from vellum.workflows.nodes import InlinePromptNode
+
+
+@pytest.mark.parametrize(
+    "json_schema",
+    [
+        pytest.param(
+            {"type": "object", "properties": "invalid"},
+            id="object with non-dict properties",
+        ),
+        pytest.param(
+            {"anyOf": {"type": "string"}},
+            id="anyOf not list",
+        ),
+        pytest.param(
+            {"type": "object", "properties": {"name": "string"}},
+            id="object with non-schema property value",
+        ),
+    ],
+)
+def test_inline_prompt_node_validation__invalid_schemas_raise_error(
+    json_schema: dict,
+) -> None:
+    """
+    Tests that InlinePromptNode validation rejects structurally invalid JSON Schemas.
+    """
+
+    # GIVEN an InlinePromptNode configured with an invalid JSON Schema
+    class MyPromptNode(InlinePromptNode):
+        ml_model = "gpt-4"
+        blocks = []
+        parameters = PromptParameters(custom_parameters={"json_schema": json_schema})
+
+    # WHEN we call __validate__ on the node
+    # THEN it should raise a SchemaError
+    with pytest.raises(jsonschema.exceptions.SchemaError):
+        MyPromptNode.__validate__()
+
+
+@pytest.mark.parametrize(
+    "json_schema",
+    [
+        pytest.param(
+            {"type": "array"},
+            id="array without items (valid per JSON Schema spec)",
+        ),
+        pytest.param(
+            {
+                "type": "array",
+                "prefixItems": [{"type": "string"}, {"type": "number"}],
+            },
+            id="array with prefixItems only",
+        ),
+        pytest.param(
+            {"type": "array", "items": {"type": "string"}},
+            id="valid array with items",
+        ),
+        pytest.param(
+            {"type": "array", "items": {}},
+            id="array with empty items object",
+        ),
+        pytest.param(
+            {"type": "array", "prefixItems": []},
+            id="array with empty prefixItems",
+        ),
+        pytest.param(
+            {
+                "type": "array",
+                "prefixItems": [{"type": "string"}],
+                "items": {"type": "number"},
+            },
+            id="array with both items and prefixItems",
+        ),
+        pytest.param(
+            {
+                "type": "array",
+                "prefixItems": [{"type": "array", "items": {"type": "string"}}],
+            },
+            id="valid nested array in prefixItems",
+        ),
+        pytest.param(
+            {
+                "type": "array",
+                "items": [
+                    {"type": "string"},
+                    {"type": "array", "items": {"type": "number"}},
+                ],
+            },
+            id="valid nested array in list-valued items",
+        ),
+        pytest.param(
+            {
+                "type": "object",
+                "properties": {
+                    "name": {"type": "string"},
+                    "age": {"type": "number"},
+                    "tags": {"type": "array", "items": {"type": "string"}},
+                    "address": {
+                        "type": "object",
+                        "properties": {
+                            "street": {"type": "string"},
+                            "city": {"type": "string"},
+                        },
+                    },
+                },
+                "anyOf": [{"required": ["name"]}, {"required": ["age"]}],
+            },
+            id="valid complex object schema with anyOf",
+        ),
+        pytest.param(
+            {
+                "name": "match_scorer_schema",
+                "schema": {
+                    "type": "object",
+                    "title": "MatchScorerSchema",
+                    "required": ["recommendation", "score", "remarks"],
+                    "properties": {
+                        "score": {
+                            "type": "integer",
+                            "title": "Match Score",
+                            "description": "Match score out of 10",
+                        },
+                        "remarks": {
+                            "type": "string",
+                            "title": "Remarks",
+                        },
+                        "recommendation": {
+                            "enum": ["Advance", "Defer", "Reject"],
+                            "type": "string",
+                            "title": "Status",
+                        },
+                    },
+                },
+            },
+            id="wrapper with valid nested schema",
+        ),
+        pytest.param(
+            {"type": "array", "items": True},
+            id="array with boolean items (valid per JSON Schema spec)",
+        ),
+        pytest.param(
+            {"anyOf": [True, {"type": "string"}]},
+            id="anyOf with boolean element (valid per JSON Schema spec)",
+        ),
+    ],
+)
+def test_inline_prompt_node_validation__valid_schemas_succeed(
+    json_schema: dict,
+) -> None:
+    """
+    Tests that InlinePromptNode validation accepts structurally valid JSON Schemas.
+    """
+
+    # GIVEN an InlinePromptNode configured with a valid JSON Schema
+    class MyPromptNode(InlinePromptNode):
+        ml_model = "gpt-4"
+        blocks = []
+        parameters = PromptParameters(custom_parameters={"json_schema": json_schema})
+
+    # WHEN we call __validate__ on the node
+    # THEN it should not raise an error
+    MyPromptNode.__validate__()
+
+
+def test_inline_prompt_node_validation__no_json_schema__succeeds():
+    """
+    Tests that InlinePromptNode without json_schema passes validation.
+    """
+
+    # GIVEN an InlinePromptNode that has no json_schema
+    class MyPromptNode(InlinePromptNode):
+        ml_model = "gpt-4"
+        blocks = []
+
+    # WHEN we call __validate__() on the node
+    # THEN it should not raise any errors
+    MyPromptNode.__validate__()