vellum-ai 0.13.0__py3-none-any.whl → 0.13.2__py3-none-any.whl

vellum_ee/workflows/display/tests/workflow_serialization/generic_nodes/test_trigger_serialization.py

@@ -10,8 +10,7 @@ class Inputs(BaseInputs):
 
 
 class BasicGenericNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
+    pass
 
 
 def test_serialize_node__basic(serialize_node):
@@ -39,11 +38,13 @@ def test_serialize_node__basic(serialize_node):
             "ports": [
                 {
                     "id": "4fbf0fff-a42e-4410-852a-238b5059198e",
+                    "name": "default",
                     "type": "DEFAULT",
                 }
             ],
             "adornments": None,
             "attributes": [],
+            "outputs": [],
         },
         serialized_node,
         ignore_order=True,
@@ -51,9 +52,6 @@ def test_serialize_node__basic(serialize_node):
 
 
 class AwaitAnyGenericNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
-
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ANY
 
@@ -83,11 +81,13 @@ def test_serialize_node__await_any(serialize_node):
             "ports": [
                 {
                     "id": "3e219c0a-e5f8-443a-ac78-1a458b189009",
+                    "name": "default",
                     "type": "DEFAULT",
                 }
             ],
             "adornments": None,
             "attributes": [],
+            "outputs": [],
         },
         serialized_node,
         ignore_order=True,
@@ -95,9 +95,6 @@ def test_serialize_node__await_any(serialize_node):
 
 
 class AwaitAllGenericNode(BaseNode):
-    class Outputs(BaseNode.Outputs):
-        output = Inputs.input
-
     class Trigger(BaseNode.Trigger):
         merge_behavior = MergeBehavior.AWAIT_ALL
 
@@ -127,11 +124,13 @@ def test_serialize_node__await_all(serialize_node):
             "ports": [
                 {
                     "id": "9797e93f-9fe6-48b2-a48e-476abbd20e32",
+                    "name": "default",
                     "type": "DEFAULT",
                 }
             ],
             "adornments": None,
             "attributes": [],
+            "outputs": [],
         },
         serialized_node,
         ignore_order=True,

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_inline_subworkflow_serialization.py

@@ -141,9 +141,42 @@ def test_serialize_workflow():
                                 "module": ["tests", "workflows", "basic_inline_subworkflow", "workflow"],
                             },
                             "trigger": {"id": "a95a34f2-e894-4fb6-a2c9-15d12c1e3135", "merge_behavior": "AWAIT_ANY"},
-                            "ports": [{"id": "4a62dea3-6a4a-4390-88d8-4af5d2c81474", "type": "DEFAULT"}],
+                            "ports": [
+                                {"id": "4a62dea3-6a4a-4390-88d8-4af5d2c81474", "type": "DEFAULT", "name": "default"}
+                            ],
                             "adornments": None,
-                            "attributes": [],
+                            "attributes": [
+                                {
+                                    "id": "b0ac6b50-22a8-42ba-a707-1aa09a653205",
+                                    "name": "metro",
+                                    "value": {
+                                        "type": "WORKFLOW_INPUT",
+                                        "input_variable_id": "f2f5da15-026d-4905-bfe7-7d16bda20eed",
+                                    },
+                                },
+                                {
+                                    "id": "c5f2d66c-5bb6-4d2a-8e4d-5356318cd3ba",
+                                    "name": "date",
+                                    "value": {
+                                        "type": "WORKFLOW_INPUT",
+                                        "input_variable_id": "aba1e6e0-dfa7-4c15-a4e6-aec6feebfaca",
+                                    },
+                                },
+                            ],
+                            "outputs": [
+                                {
+                                    "id": "3f4c753e-f057-47bb-9748-7968283cc8aa",
+                                    "name": "temperature",
+                                    "type": "NUMBER",
+                                    "value": None,
+                                },
+                                {
+                                    "id": "2a4a62b3-cd26-4d2c-b3f1-eaa5f9dd22dd",
+                                    "name": "reasoning",
+                                    "type": "STRING",
+                                    "value": None,
+                                },
+                            ],
                         },
                         {
                             "id": "a773c3a5-78cb-4250-8d29-7282e8a579d3",

vellum_ee/workflows/display/tests/workflow_serialization/test_basic_map_node_serialization.py

@@ -130,9 +130,36 @@ def test_serialize_workflow():
                                 "module": ["tests", "workflows", "basic_map_node", "workflow"],
                             },
                             "trigger": {"id": "01324747-9bc0-4ecd-a8ab-40dca5a94e2e", "merge_behavior": "AWAIT_ANY"},
-                            "ports": [{"id": "36791877-95b0-4390-a794-10accd8f548e", "type": "DEFAULT"}],
+                            "ports": [
+                                {"id": "36791877-95b0-4390-a794-10accd8f548e", "type": "DEFAULT", "name": "default"}
+                            ],
                             "adornments": None,
-                            "attributes": [],
+                            "attributes": [
+                                {
+                                    "id": "4306b0ce-f31c-4c8f-8122-659d40e60ffe",
+                                    "name": "item",
+                                    "value": {
+                                        "type": "WORKFLOW_INPUT",
+                                        "input_variable_id": "b29bb546-9bc8-4136-857d-8c7a464ba9d4",
+                                    },
+                                },
+                                {
+                                    "id": "bca02056-0fd9-44c6-8aa9-743dd1180d48",
+                                    "name": "index",
+                                    "value": {
+                                        "type": "WORKFLOW_INPUT",
+                                        "input_variable_id": "17e7ca49-668f-450d-a792-e1f97d13db67",
+                                    },
+                                },
+                            ],
+                            "outputs": [
+                                {
+                                    "id": "a7bcb362-a2b8-4476-b0de-a361efeec204",
+                                    "name": "count",
+                                    "type": "NUMBER",
+                                    "value": None,
+                                }
+                            ],
                         },
                         {
                             "id": "6f4883b2-70b1-4e1c-ae15-7d0f5aec810b",

vellum_ee/workflows/display/utils/vellum.py

@@ -1,10 +1,8 @@
-import enum
-import json
-import typing
-from typing import Any, List, Union, cast
+from typing import Any, TypeVar
 
-from vellum import ChatMessage, SearchResult, SearchResultRequest, VellumVariableType
+from vellum.client.types.vellum_variable_type import VellumVariableType
 from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.references import OutputReference, WorkflowInputReference
 from vellum.workflows.references.execution_count import ExecutionCountReference
 from vellum.workflows.references.node import NodeReference
@@ -13,25 +11,19 @@ from vellum.workflows.utils.vellum_variables import primitive_type_to_vellum_var
 from vellum.workflows.vellum_client import create_vellum_client
 from vellum_ee.workflows.display.types import WorkflowDisplayContext
 from vellum_ee.workflows.display.vellum import (
-    ChatHistoryVellumValue,
     ConstantValuePointer,
     ExecutionCounterData,
     ExecutionCounterPointer,
     InputVariableData,
     InputVariablePointer,
-    JsonVellumValue,
     NodeInputValuePointerRule,
     NodeOutputData,
     NodeOutputPointer,
-    NumberVellumValue,
-    SearchResultsVellumValue,
-    StringVellumValue,
-    VellumValue,
     WorkspaceSecretData,
     WorkspaceSecretPointer,
 )
 
-_T = typing.TypeVar("_T")
+_T = TypeVar("_T")
 
 
 def infer_vellum_variable_type(value: Any) -> VellumVariableType:
@@ -90,33 +82,3 @@ def create_node_input_value_pointer_rule(
         return ConstantValuePointer(type="CONSTANT_VALUE", data=vellum_value)
 
     raise ValueError(f"Unsupported descriptor type: {value.__class__.__name__}")
-
-
-def primitive_to_vellum_value(value: Any) -> VellumValue:
-    """Converts a python primitive to a VellumVariableValue"""
-
-    if isinstance(value, str):
-        return StringVellumValue(value=value)
-    elif isinstance(value, enum.Enum):
-        return StringVellumValue(value=value.value)
-    elif isinstance(value, (int, float)):
-        return NumberVellumValue(value=value)
-    elif isinstance(value, list) and (
-        all(isinstance(message, ChatMessage) for message in value)
-        or all(isinstance(message, ChatMessage) for message in value)
-    ):
-        chat_messages = cast(Union[List[ChatMessage], List[ChatMessage]], value)
-        return ChatHistoryVellumValue(value=chat_messages)
-    elif isinstance(value, list) and (
-        all(isinstance(search_result, SearchResultRequest) for search_result in value)
-        or all(isinstance(search_result, SearchResult) for search_result in value)
-    ):
-        search_results = cast(Union[List[SearchResultRequest], List[SearchResult]], value)
-        return SearchResultsVellumValue(value=search_results)
-
-    try:
-        json_value = json.dumps(value)
-    except json.JSONDecodeError:
-        raise ValueError(f"Unsupported variable type: {value.__class__.__name__}")
-
-    return JsonVellumValue(value=json.loads(json_value))

vellum_ee/workflows/display/vellum.py

@@ -1,11 +1,13 @@
 from dataclasses import dataclass, field
 from enum import Enum
 from uuid import UUID
-from typing import Any, List, Literal, Optional, Union
+from typing import List, Literal, Optional, Union
 
 from pydantic import Field
 
-from vellum import ChatMessage, PromptParameters, SearchResult, SearchResultRequest, VellumVariable, VellumVariableType
+from vellum import PromptParameters, VellumVariable, VellumVariableType
+from vellum.client.types.array_vellum_value import ArrayVellumValue
+from vellum.client.types.vellum_value import VellumValue
 from vellum.core import UniversalBaseModel
 from vellum_ee.workflows.display.base import (
     EdgeDisplay,
@@ -140,45 +142,14 @@ class WorkflowNodeType(str, Enum):
     ERROR = "ERROR"
 
 
-class StringVellumValue(UniversalBaseModel):
-    type: Literal["STRING"] = "STRING"
-    value: str
-
-
-class NumberVellumValue(UniversalBaseModel):
-    type: Literal["NUMBER"] = "NUMBER"
-    value: Union[int, float]
-
-
-class ChatHistoryVellumValue(UniversalBaseModel):
-    type: Literal["CHAT_HISTORY"] = "CHAT_HISTORY"
-    value: Union[List[ChatMessage], List[ChatMessage]]
-
-
-class SearchResultsVellumValue(UniversalBaseModel):
-    type: Literal["SEARCH_RESULTS"] = "SEARCH_RESULTS"
-    value: Union[List[SearchResultRequest], List[SearchResult]]
-
-
-class JsonVellumValue(UniversalBaseModel):
-    type: Literal["JSON"] = "JSON"
-    value: Optional[Any] = None
-
-
-VellumValue = Union[
-    StringVellumValue,
-    NumberVellumValue,
-    ChatHistoryVellumValue,
-    SearchResultsVellumValue,
-    JsonVellumValue,
-]
-
-
 class ConstantValuePointer(UniversalBaseModel):
     type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
     data: VellumValue
 
 
+ArrayVellumValue.model_rebuild()
+
+
 class NodeOutputData(UniversalBaseModel):
     node_id: str
     output_id: str

vellum_ee/workflows/display/workflows/vellum_workflow_display.py

@@ -5,6 +5,7 @@ from typing import Dict, List, Optional, Type, cast
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges import Edge
 from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
 from vellum.workflows.nodes.utils import get_unadorned_node, get_unadorned_port, get_wrapped_node, has_wrapped_node
 from vellum.workflows.ports import Port
@@ -16,7 +17,7 @@ from vellum.workflows.utils.uuids import uuid4_from_hash
 from vellum_ee.workflows.display.nodes.base_node_vellum_display import BaseNodeVellumDisplay
 from vellum_ee.workflows.display.nodes.types import PortDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
-from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type, primitive_to_vellum_value
+from vellum_ee.workflows.display.utils.vellum import infer_vellum_variable_type
 from vellum_ee.workflows.display.vellum import (
     EdgeVellumDisplay,
     EdgeVellumDisplayOverrides,
@@ -264,16 +265,18 @@ class VellumWorkflowDisplay(
         self, workflow_input: WorkflowInputReference, overrides: Optional[WorkflowInputsVellumDisplayOverrides] = None
     ) -> WorkflowInputsVellumDisplay:
         workflow_input_id: UUID
+        name = None
         required = None
         color = None
         if overrides:
             workflow_input_id = overrides.id
+            name = overrides.name
             required = overrides.required
             color = overrides.color
         else:
             workflow_input_id = uuid4_from_hash(f"{self.workflow_id}|inputs|id|{workflow_input.name}")
 
-        return WorkflowInputsVellumDisplay(id=workflow_input_id, required=required, color=color)
+        return WorkflowInputsVellumDisplay(id=workflow_input_id, name=name, required=required, color=color)
 
     def _generate_entrypoint_display(
         self,