vellum-ai 0.14.17__py3-none-any.whl → 0.14.19__py3-none-any.whl

vellum_ee/workflows/display/vellum.py

@@ -1,11 +1,10 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from uuid import UUID
 from typing import List, Literal, Optional, Union
 
 from pydantic import Field
 
-from vellum import PromptParameters, VellumVariable, VellumVariableType
+from vellum import VellumVariableType
 from vellum.client.types.array_vellum_value import ArrayVellumValue
 from vellum.client.types.vellum_value import VellumValue
 from vellum.core import UniversalBaseModel
@@ -138,24 +137,6 @@ class WorkflowOutputVellumDisplay(WorkflowOutputVellumDisplayOverrides):
     pass
 
 
-class WorkflowNodeType(str, Enum):
-    PROMPT = "PROMPT"
-    TEMPLATING = "TEMPLATING"
-    NOTE = "NOTE"
-    CODE_EXECUTION = "CODE_EXECUTION"
-    METRIC = "METRIC"
-    SEARCH = "SEARCH"
-    WEBHOOK = "WEBHOOK"
-    MERGE = "MERGE"
-    CONDITIONAL = "CONDITIONAL"
-    API = "API"
-    ENTRYPOINT = "ENTRYPOINT"
-    TERMINAL = "TERMINAL"
-    SUBWORKFLOW = "SUBWORKFLOW"
-    MAP = "MAP"
-    ERROR = "ERROR"
-
-
 class ConstantValuePointer(UniversalBaseModel):
     type: Literal["CONSTANT_VALUE"] = "CONSTANT_VALUE"
     data: VellumValue
@@ -220,131 +201,3 @@ class NodeInput(UniversalBaseModel):
     id: str
     key: str
     value: NodeInputValuePointer
-
-
-class BaseWorkflowNode(UniversalBaseModel):
-    id: str
-    inputs: List[NodeInput]
-    type: str
-    display_data: Optional[NodeDisplayData] = None
-    base: CodeResourceDefinition
-    definition: CodeResourceDefinition
-
-
-class EntrypointNodeData(UniversalBaseModel):
-    source_handle_id: str
-
-
-class EntrypointNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.ENTRYPOINT] = WorkflowNodeType.ENTRYPOINT
-    data: EntrypointNodeData
-
-
-class PromptTemplateBlockData(UniversalBaseModel):
-    version: Literal[1] = 1
-    # blocks: List[PromptBlockRequest]
-
-
-class PromptVersionExecConfig(UniversalBaseModel):
-    parameters: PromptParameters
-    input_variables: List[VellumVariable]
-    prompt_template_block_data: PromptTemplateBlockData
-
-
-class BasePromptNodeData(UniversalBaseModel):
-    label: str
-    output_id: str
-    error_output_id: Optional[str] = None
-    array_output_id: str
-    source_handle_id: str
-    target_handle_id: str
-
-
-class InlinePromptNodeData(BasePromptNodeData):
-    variant: Literal["INLINE"] = "INLINE"
-    exec_config: PromptVersionExecConfig
-    ml_model_name: str
-
-
-class DeploymentPromptNodeData(BasePromptNodeData):
-    variant: Literal["DEPLOYMENT"] = "DEPLOYMENT"
-    deployment_id: str
-    release_tag: str
-
-
-PromptNodeData = Union[
-    InlinePromptNodeData,
-    DeploymentPromptNodeData,
-]
-
-
-class PromptNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.PROMPT] = WorkflowNodeType.PROMPT
-    data: PromptNodeData
-
-
-class SearchNodeData(UniversalBaseModel):
-    label: str
-
-    results_output_id: str
-    text_output_id: str
-    error_output_id: Optional[str] = None
-
-    source_handle_id: str
-    target_handle_id: str
-
-    query_node_input_id: str
-    document_index_node_input_id: str
-    weights_node_input_id: str
-    limit_node_input_id: str
-    separator_node_input_id: str
-    result_merging_enabled_node_input_id: str
-    external_id_filters_node_input_id: str
-    metadata_filters_node_input_id: str
-
-
-class SearchNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.SEARCH] = WorkflowNodeType.SEARCH
-    data: SearchNodeData
-
-
-class FinalOutputNodeData(UniversalBaseModel):
-    label: str
-    name: str
-    target_handle_id: str
-    output_id: str
-    output_type: VellumVariableType
-    node_input_id: str
-
-
-class FinalOutputNode(BaseWorkflowNode):
-    type: Literal[WorkflowNodeType.TERMINAL] = WorkflowNodeType.TERMINAL
-    data: FinalOutputNodeData
-
-
-WorkflowNode = Union[
-    EntrypointNode,
-    PromptNode,
-    SearchNode,
-    FinalOutputNode,
-]
-
-
-class WorkflowEdge(UniversalBaseModel):
-    id: str
-    source_node_id: str
-    source_handle_id: str
-    target_node_id: str
-    target_handle_id: str
-
-
-class WorkflowRawData(UniversalBaseModel):
-    nodes: List[WorkflowNode]
-    edges: List[WorkflowEdge]
-    display_data: Optional[WorkflowDisplayData] = None
-
-
-class WorkflowVersionExecConfig(UniversalBaseModel):
-    workflow_raw_data: WorkflowRawData
-    input_variables: List[VellumVariable]
-    output_variables: List[VellumVariable]

vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -429,7 +429,7 @@ class BaseWorkflowDisplay(
                     )
                     subworkflow_display_context = subworkflow_display.get_event_display_context()
 
-            node_event_displays[str(node_id)] = NodeEventDisplayContext(
+            node_event_displays[node_id] = NodeEventDisplayContext(
                 input_display=input_display,
                 output_display=output_display,
                 port_display=port_display_meta,

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

@@ -114,8 +114,8 @@ def test_get_event_display_context__node_display_filled_without_base_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
-    assert str(StartNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
 
     # AND so should their output ids
     assert StartNode.__output_ids__ == node_event_display.output_display
@@ -137,8 +137,8 @@ def test_get_event_display_context__node_display_filled_without_output_display()
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the node display should be included
-    assert str(StartNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(StartNode.__id__)]
+    assert StartNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[StartNode.__id__]
 
     # AND so should their output ids
     assert node_event_display.output_display.keys() == {"foo"}
@@ -163,11 +163,11 @@ def test_get_event_display_context__node_display_to_include_subworkflow_display(
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(SubworkflowNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(SubworkflowNode.__id__)]
+    assert SubworkflowNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[SubworkflowNode.__id__]
 
     assert node_event_display.subworkflow_display is not None
-    assert str(InnerNode.__id__) in node_event_display.subworkflow_display.node_displays
+    assert InnerNode.__id__ in node_event_display.subworkflow_display.node_displays
 
 
 @pytest.mark.parametrize(
@@ -204,13 +204,13 @@ def test_get_event_display_context__node_display_for_adornment_nodes(
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(adornment_node_id) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(adornment_node_id)]
+    assert adornment_node_id in display_context.node_displays
+    node_event_display = display_context.node_displays[adornment_node_id]
     assert node_event_display.subworkflow_display is not None
-    assert str(inner_node_id) in node_event_display.subworkflow_display.node_displays
+    assert inner_node_id in node_event_display.subworkflow_display.node_displays
 
     # AND the inner node should have the correct outputs
-    inner_node_display = node_event_display.subworkflow_display.node_displays[str(inner_node_id)]
+    inner_node_display = node_event_display.subworkflow_display.node_displays[inner_node_id]
     assert inner_node_display.output_display.keys() == {"foo"}
     assert node_event_display.output_display.keys() == expected_adornment_output_names
 
@@ -231,7 +231,7 @@ def test_get_event_display_context__templating_node_input_display():
     display_context = VellumWorkflowDisplay(MyWorkflow).get_event_display_context()
 
     # THEN the subworkflow display should be included
-    assert str(MyNode.__id__) in display_context.node_displays
-    node_event_display = display_context.node_displays[str(MyNode.__id__)]
+    assert MyNode.__id__ in display_context.node_displays
+    node_event_display = display_context.node_displays[MyNode.__id__]
 
     assert node_event_display.input_display.keys() == {"inputs.foo"}

vellum_ee/workflows/tests/test_display_meta.py

@@ -1,7 +1,7 @@
 import pytest
 import os
 import sys
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from vellum.workflows import BaseWorkflow
 from vellum_ee.workflows.display.workflows import BaseWorkflowDisplay
@@ -35,22 +35,22 @@ def test_base_class_dynamic_import(files):
         "node_displays": {
             "533c6bd8-6088-4abc-a168-8c1758abcd33": {
                 "input_display": {
-                    "example_var_1": UUID("a0d1d7cf-242a-4bd9-a437-d308a7ced9b3"),
-                    "template": UUID("f97d721a-e685-498e-90c3-9c3d9358fdad"),
+                    "example_var_1": "a0d1d7cf-242a-4bd9-a437-d308a7ced9b3",
+                    "template": "f97d721a-e685-498e-90c3-9c3d9358fdad",
                 },
-                "output_display": {"result": UUID("423bc529-1a1a-4f72-af4d-cbdb5f0a5929")},
-                "port_display": {"default": UUID("afda9a19-0618-42e1-9b63-5d0db2a88f62")},
+                "output_display": {"result": "423bc529-1a1a-4f72-af4d-cbdb5f0a5929"},
+                "port_display": {"default": "afda9a19-0618-42e1-9b63-5d0db2a88f62"},
                 "subworkflow_display": None,
             },
             "f3ef4b2b-fec9-4026-9cc6-e5eac295307f": {
-                "input_display": {"node_input": UUID("fe6cba85-2423-4b5e-8f85-06311a8be5fb")},
-                "output_display": {"value": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+                "input_display": {"node_input": "fe6cba85-2423-4b5e-8f85-06311a8be5fb"},
+                "output_display": {"value": "5469b810-6ea6-4362-9e79-e360d44a1405"},
                 "port_display": {},
                 "subworkflow_display": None,
             },
         },
-        "workflow_inputs": {"input_value": UUID("2268a996-bd17-4832-b3ff-f5662d54b306")},
-        "workflow_outputs": {"final_output": UUID("5469b810-6ea6-4362-9e79-e360d44a1405")},
+        "workflow_inputs": {"input_value": "2268a996-bd17-4832-b3ff-f5662d54b306"},
+        "workflow_outputs": {"final_output": "5469b810-6ea6-4362-9e79-e360d44a1405"},
     }
     assert display_meta
-    assert display_meta.dict() == expected_result
+    assert display_meta.model_dump(mode="json") == expected_result