vellum-ai 0.14.48__tar.gz → 0.14.50__tar.gz

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.48
+Version: 0.14.50
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/nodes/base_node_display.py

@@ -44,18 +44,40 @@ if TYPE_CHECKING:
 _NodeDisplayAttrType = TypeVar("_NodeDisplayAttrType")
 
 
+def _get_node_input_ids_by_ref(node_class: Type[BaseNode], path: str, inst: Any):
+    if isinstance(inst, dict):
+        node_input_ids_by_name: Dict[str, UUID] = {}
+        for key, value in inst.items():
+            node_input_ids_by_name.update(_get_node_input_ids_by_ref(node_class, f"{path}.{key}", value))
+        return node_input_ids_by_name
+
+    return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
+
+
 class BaseNodeDisplayMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         cls = cast(Type["BaseNodeDisplay"], super().__new__(mcs, name, bases, dct))
+        # This cast shouldn't be necessary, but it's a workaround for a mypy bug
+        node_class = cast(Type[BaseNode], cls.infer_node_class() if name != "BaseNodeDisplay" else BaseNode)
 
         if not dct.get("output_display"):
-            node_class = cls.infer_node_class()
             cls.output_display = {
                 ref: NodeOutputDisplay(id=node_class.__output_ids__[ref.name], name=ref.name)
                 for ref in node_class.Outputs
                 if ref.name in node_class.__output_ids__
             }
 
+        if not dct.get("node_input_ids_by_name"):
+            node_input_ids_by_name: Dict[str, UUID] = {}
+            for ref in node_class:
+                if ref not in cls.__serializable_inputs__:
+                    continue
+
+                node_input_ids_by_name.update(_get_node_input_ids_by_ref(node_class, ref.name, ref.instance))
+
+            if node_input_ids_by_name:
+                cls.node_input_ids_by_name = node_input_ids_by_name
+
         return cls.__annotate_node__()
 
     def __annotate_node__(cls):

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/nodes/get_node_display_class.py

@@ -1,9 +1,7 @@
 import types
-from uuid import UUID
-from typing import TYPE_CHECKING, Any, Dict, Generic, Type, TypeVar
+from typing import TYPE_CHECKING, Generic, Type, TypeVar
 
 from vellum.workflows.types.generics import NodeType
-from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.utils.registry import get_from_node_display_registry
 
 if TYPE_CHECKING:
@@ -22,30 +20,9 @@ def get_node_display_class(node_class: Type[NodeType]) -> Type["BaseNodeDisplay"
     # `base_node_display_class` is always a Generic class, so it's safe to index into it
     NodeDisplayBaseClass = base_node_display_class[_NodeClassType]  # type: ignore[index]
 
-    def _get_node_input_ids_by_ref(path: str, inst: Any):
-        if isinstance(inst, dict):
-            node_input_ids_by_name: Dict[str, UUID] = {}
-            for key, value in inst.items():
-                node_input_ids_by_name.update(_get_node_input_ids_by_ref(f"{path}.{key}", value))
-            return node_input_ids_by_name
-
-        return {path: uuid4_from_hash(f"{node_class.__id__}|{path}")}
-
-    def exec_body(ns: Dict):
-        node_input_ids_by_name: Dict[str, UUID] = {}
-        for ref in node_class:
-            if ref not in base_node_display_class.__serializable_inputs__:
-                continue
-
-            node_input_ids_by_name.update(_get_node_input_ids_by_ref(ref.name, ref.instance))
-
-        if node_input_ids_by_name:
-            ns["node_input_ids_by_name"] = node_input_ids_by_name
-
     NodeDisplayClass = types.new_class(
         f"{node_class.__name__}Display",
         bases=(NodeDisplayBaseClass, Generic[_NodeClassType]),
-        exec_body=exec_body,
     )
 
     return NodeDisplayClass

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/nodes/vellum/code_execution_node.py

@@ -79,8 +79,8 @@ class BaseCodeExecutionNodeDisplay(BaseNodeDisplay[_CodeExecutionNodeType], Gene
 
         packages = raise_if_descriptor(node.packages)
 
-        _, output_display = display_context.global_node_output_displays[node.Outputs.result]
-        _, log_output_display = display_context.global_node_output_displays[node.Outputs.log]
+        output_display = self.output_display[node.Outputs.result]
+        log_output_display = self.output_display[node.Outputs.log]
 
         output_type = primitive_type_to_vellum_variable_type(node.get_output_type())
 

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_code_execution_node.py

@@ -2,7 +2,9 @@ import pytest
 from uuid import UUID
 from typing import Type
 
+from vellum.client.core.api_error import ApiError
 from vellum.workflows.nodes.displayable.code_execution_node.node import CodeExecutionNode
+from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.vellum.code_execution_node import BaseCodeExecutionNodeDisplay
 from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
@@ -110,3 +112,44 @@ def test_serialize_node__code_node_inputs(GetDisplayClass, expected_input_id):
             },
         },
     ]
+
+
+def test_serialize_node__with_unresolved_secret_references(vellum_client):
+    # GIVEN a node has access to a secret reference
+    class MyNode(CodeExecutionNode):
+        code_inputs = {"api_key": VellumSecretReference("MY_API_KEY")}
+
+    # AND the secret is not found
+    vellum_client.workspace_secrets.retrieve.side_effect = ApiError(
+        status_code=404,
+        body="Secret not found",
+    )
+
+    # AND a workflow with the code node
+    class Workflow(BaseWorkflow):
+        graph = MyNode
+
+    # WHEN we serialize the workflow
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    data: dict = workflow_display.serialize()
+
+    # THEN the condition should be serialized correctly
+    node = next(node for node in data["workflow_raw_data"]["nodes"] if node["type"] == "CODE_EXECUTION")
+    assert node["inputs"][0]["value"] == {
+        "combinator": "OR",
+        "rules": [
+            {
+                "type": "WORKSPACE_SECRET",
+                "data": {
+                    "type": "STRING",
+                    "workspace_secret_id": None,
+                },
+            }
+        ],
+    }
+
+    # AND we should have a warning of the invalid reference
+    # TODO: Come up with a proposal for how nodes should propagate warnings
+    # warnings = list(workflow_display.errors)
+    # assert len(warnings) == 1
+    # assert "Failed to resolve secret reference 'MY_API_KEY'" in str(warnings[0])

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_attributes_serialization.py

@@ -1,4 +1,3 @@
-import pytest
 from uuid import uuid4
 from typing import List
 
@@ -257,11 +256,58 @@ def test_serialize_node__workflow_input_as_nested_chat_history():
 
     # WHEN the workflow is serialized
     workflow_display = get_workflow_display(workflow_class=Workflow)
-    with pytest.raises(Exception) as exc_info:
-        workflow_display.serialize()
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the attribute reference
+    generic_node = next(
+        node for node in serialized_workflow["workflow_raw_data"]["nodes"] if node["id"] == str(GenericNode.__id__)
+    )
 
-    # THEN we should raise a user facing error
-    assert str(exc_info.value) == "Failed to serialize attribute 'attr': Nested references are not supported."
+    assert not DeepDiff(
+        {
+            "id": "11be9d37-0069-4695-a317-14a3b6519d4e",
+            "label": "test_serialize_node__workflow_input_as_nested_chat_history.<locals>.GenericNode",
+            "type": "GENERIC",
+            "display_data": {"position": {"x": 0.0, "y": 0.0}},
+            "base": {"name": "BaseNode", "module": ["vellum", "workflows", "nodes", "bases", "base"]},
+            "definition": {
+                "name": "GenericNode",
+                "module": [
+                    "vellum_ee",
+                    "workflows",
+                    "display",
+                    "tests",
+                    "workflow_serialization",
+                    "generic_nodes",
+                    "test_attributes_serialization",
+                ],
+            },
+            "trigger": {"id": "d4548468-85a4-449e-92e0-e2d8b8fd097c", "merge_behavior": "AWAIT_ATTRIBUTES"},
+            "ports": [{"id": "c4a9a57d-1380-4689-8500-e8a0b6477291", "name": "default", "type": "DEFAULT"}],
+            "adornments": None,
+            "attributes": [
+                {
+                    "id": "e878bbc9-1231-461e-9e9d-947604da116e",
+                    "name": "attr",
+                    "value": {
+                        "type": "DICTIONARY_REFERENCE",
+                        "entries": [
+                            {
+                                "key": "hello",
+                                "value": {
+                                    "type": "WORKFLOW_INPUT",
+                                    "input_variable_id": "f727c3f9-f27f-4ac9-abd7-12bf612a094e",
+                                },
+                            }
+                        ],
+                    },
+                }
+            ],
+            "outputs": [],
+        },
+        generic_node,
+        ignore_order=True,
+    )
 
 
 def test_serialize_node__node_output(serialize_node):

vellum_ai-0.14.50/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py

@@ -0,0 +1,177 @@
+from deepdiff import DeepDiff
+
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_tool_calling_node.workflow import BasicToolCallingNodeWorkflow
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a generic node
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=BasicToolCallingNodeWorkflow)
+
+    serialized_workflow: dict = workflow_display.serialize()
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 0
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+    assert not DeepDiff(
+        [
+            {"id": "8e7c0147-930d-4b7f-b6b1-6d79641cd3eb", "key": "text", "type": "STRING"},
+            {"id": "01a07e6d-7269-4f45-8b44-ef0227a2e88d", "key": "chat_history", "type": "CHAT_HISTORY"},
+        ],
+        output_variables,
+        ignore_order=True,
+    )
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    tool_calling_node = workflow_raw_data["nodes"][1]
+    assert tool_calling_node == {
+        "id": "21f29cac-da87-495f-bba1-093d423f4e46",
+        "label": "GetCurrentWeatherNode",
+        "type": "GENERIC",
+        "display_data": {
+            "position": {"x": 0.0, "y": 0.0},
+            "comment": {"value": "\n    A tool calling node that calls the get_current_weather function.\n    "},
+        },
+        "base": {
+            "name": "ToolCallingNode",
+            "module": ["vellum", "workflows", "nodes", "experimental", "tool_calling_node", "node"],
+        },
+        "definition": {
+            "name": "GetCurrentWeatherNode",
+            "module": ["tests", "workflows", "basic_tool_calling_node", "workflow"],
+        },
+        "trigger": {"id": "2414743b-b1dd-4552-8abf-9b7481df9762", "merge_behavior": "AWAIT_ATTRIBUTES"},
+        "ports": [{"id": "3cd6d78c-9dad-42aa-ad38-31f67057c379", "name": "default", "type": "DEFAULT"}],
+        "adornments": None,
+        "attributes": [
+            {
+                "id": "44420e39-966f-4c59-bdf8-6365a61c5d2a",
+                "name": "ml_model",
+                "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": "gpt-4o-mini"}},
+            },
+            {
+                "id": "669cfb4b-8c25-460e-8952-b63d91302cbc",
+                "name": "blocks",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": [
+                            {
+                                "block_type": "CHAT_MESSAGE",
+                                "state": None,
+                                "cache_config": None,
+                                "chat_role": "SYSTEM",
+                                "chat_source": None,
+                                "chat_message_unterminated": None,
+                                "blocks": [
+                                    {
+                                        "block_type": "RICH_TEXT",
+                                        "state": None,
+                                        "cache_config": None,
+                                        "blocks": [
+                                            {
+                                                "block_type": "PLAIN_TEXT",
+                                                "state": None,
+                                                "cache_config": None,
+                                                "text": "You are a weather expert",
+                                            }
+                                        ],
+                                    }
+                                ],
+                            },
+                            {
+                                "block_type": "CHAT_MESSAGE",
+                                "state": None,
+                                "cache_config": None,
+                                "chat_role": "USER",
+                                "chat_source": None,
+                                "chat_message_unterminated": None,
+                                "blocks": [
+                                    {
+                                        "block_type": "RICH_TEXT",
+                                        "state": None,
+                                        "cache_config": None,
+                                        "blocks": [
+                                            {
+                                                "block_type": "VARIABLE",
+                                                "state": None,
+                                                "cache_config": None,
+                                                "input_variable": "question",
+                                            }
+                                        ],
+                                    }
+                                ],
+                            },
+                        ],
+                    },
+                },
+            },
+            {
+                "id": "78324739-ff89-47a5-902b-10da0cb95c6d",
+                "name": "functions",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": [
+                            {
+                                "state": None,
+                                "cache_config": None,
+                                "name": "get_current_weather",
+                                "description": None,
+                                "parameters": {
+                                    "type": "object",
+                                    "properties": {"location": {"type": "string"}, "unit": {"type": "string"}},
+                                    "required": ["location", "unit"],
+                                },
+                                "forced": None,
+                                "strict": None,
+                            }
+                        ],
+                    },
+                },
+            },
+            {
+                "id": "0f6dc102-3460-4963-91fa-7ba85d65ef7a",
+                "name": "prompt_inputs",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {"type": "JSON", "value": {"question": "What's the weather like in San Francisco?"}},
+                },
+            },
+            {
+                "id": "5c041b7d-732c-4773-a93a-32211f2af0b3",
+                "name": "max_tool_calls",
+                "value": {"type": "CONSTANT_VALUE", "value": {"type": "NUMBER", "value": 1.0}},
+            },
+        ],
+        "outputs": [
+            {
+                "id": "e62bc785-a914-4066-b79e-8c89a5d0ec6c",
+                "name": "text",
+                "type": "STRING",
+                "value": {"type": "CONSTANT_VALUE", "value": {"type": "STRING", "value": ""}},
+            },
+            {
+                "id": "4674f1d9-e3af-411f-8a55-40a3a3ab5394",
+                "name": "chat_history",
+                "type": "CHAT_HISTORY",
+                "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": []}},
+            },
+        ],
+    }

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -253,7 +253,7 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
                 "type": "CONSTANT_VALUE",
                 "value": {
                     "type": "JSON",
-                    "items": constant_values,
+                    "value": constant_values,
                 },
             }
         else:
@@ -262,8 +262,27 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
                 "items": cast(JsonArray, serialized_items),  # list[JsonObject] -> JsonArray
             }
 
-    if isinstance(value, dict) and any(isinstance(v, BaseDescriptor) for v in value.values()):
-        raise ValueError("Nested references are not supported.")
+    if isinstance(value, dict):
+        serialized_entries = [
+            {"key": key, "value": serialize_value(display_context, val)} for key, val in value.items()
+        ]
+
+        # Check if all entries have constant values
+        if all(entry["value"]["type"] == "CONSTANT_VALUE" for entry in serialized_entries):
+            constant_entries = {}
+            for entry in serialized_entries:
+                entry_value = entry["value"]["value"]
+                constant_entries[entry["key"]] = entry_value["value"]
+
+            return {
+                "type": "CONSTANT_VALUE",
+                "value": {
+                    "type": "JSON",
+                    "value": constant_entries,
+                },
+            }
+        else:
+            return {"type": "DICTIONARY_REFERENCE", "entries": cast(JsonArray, serialized_entries)}
 
     if not isinstance(value, BaseDescriptor):
         vellum_value = primitive_to_vellum_value(value)

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/utils/vellum.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
+from vellum.client.core.api_error import ApiError
 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
@@ -117,13 +118,18 @@ def create_node_input_value_pointer_rule(
         workflow_input_display = display_context.global_workflow_input_displays[value]
         return InputVariablePointer(data=InputVariableData(input_variable_id=str(workflow_input_display.id)))
     if isinstance(value, VellumSecretReference):
-        workspace_secret = display_context.client.workspace_secrets.retrieve(
-            id=value.name,
-        )
+        try:
+            workspace_secret = display_context.client.workspace_secrets.retrieve(
+                id=value.name,
+            )
+            workspace_secret_id: Optional[str] = str(workspace_secret.id)
+        except ApiError:
+            workspace_secret_id = None
+
         return WorkspaceSecretPointer(
             data=WorkspaceSecretData(
                 type="STRING",
-                workspace_secret_id=str(workspace_secret.id),
+                workspace_secret_id=workspace_secret_id,
             ),
         )
     if isinstance(value, ExecutionCountReference):

{vellum_ai-0.14.48 → vellum_ai-0.14.50}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -153,20 +153,8 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             "definition": None,
         }
 
-        # Add all the nodes in the workflow
-        for node in self._workflow.get_nodes():
-            node_display = self.display_context.node_displays[node]
-
-            try:
-                serialized_node = node_display.serialize(self.display_context)
-            except NotImplementedError as e:
-                self.add_error(e)
-                continue
-
-            serialized_nodes[node_display.node_id] = serialized_node
-
-        # Add all unused nodes in the workflow
-        for node in self._workflow.get_unused_nodes():
+        # Add all the nodes in the workflows
+        for node in self._workflow.get_all_nodes():
             node_display = self.display_context.node_displays[node]
 
             try:
@@ -417,16 +405,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
 
         port_displays: PortDisplays = {}
 
-        for node in self._workflow.get_nodes():
-            self._enrich_node_displays(
-                node=node,
-                node_displays=node_displays,
-                global_node_displays=global_node_displays,
-                global_node_output_displays=global_node_output_displays,
-                port_displays=port_displays,
-            )
-
-        for node in self._workflow.get_unused_nodes():
+        for node in self._workflow.get_all_nodes():
             self._enrich_node_displays(
                 node=node,
                 node_displays=node_displays,