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

vellum/workflows/state/context.py

@@ -1,16 +1,37 @@
+from dataclasses import dataclass
 from functools import cached_property
+import json
 from queue import Queue
-from typing import TYPE_CHECKING, Dict, List, Optional, Type
+from uuid import UUID, uuid4
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type
 
-from vellum import Vellum
-from vellum.workflows.context import ExecutionContext, get_execution_context
+from vellum import Vellum, __version__
+from vellum.workflows.context import ExecutionContext, get_execution_context, set_execution_context
+from vellum.workflows.events.types import ExternalParentContext
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
+from vellum.workflows.state.store import Store
+from vellum.workflows.utils.uuids import generate_workflow_deployment_prefix
+from vellum.workflows.utils.zip import extract_zip_files
 from vellum.workflows.vellum_client import create_vellum_client
 
 if TYPE_CHECKING:
     from vellum.workflows.events.workflow import WorkflowEvent
+    from vellum.workflows.state.base import BaseState
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+@dataclass
+class WorkflowDeploymentMetadata:
+    """Metadata about a workflow deployment needed for parent context construction."""
+
+    deployment_id: UUID
+    deployment_name: str
+    deployment_history_item_id: UUID
+    release_tag_id: UUID
+    release_tag_name: str
+    workflow_version_id: UUID
 
 
 class WorkflowContext:
@@ -20,13 +41,29 @@ class WorkflowContext:
         vellum_client: Optional[Vellum] = None,
         execution_context: Optional[ExecutionContext] = None,
         generated_files: Optional[dict[str, str]] = None,
+        namespace: Optional[str] = None,
+        store_class: Optional[Type[Store]] = None,
     ):
         self._vellum_client = vellum_client
         self._event_queue: Optional[Queue["WorkflowEvent"]] = None
         self._node_output_mocks_map: Dict[Type[BaseOutputs], List[MockNodeExecution]] = {}
         self._execution_context = get_execution_context()
-        if not self._execution_context.parent_context and execution_context:
-            self._execution_context = execution_context
+        self._namespace = namespace
+        self._store_class = store_class if store_class is not None else Store
+
+        if execution_context is not None:
+            self._execution_context.trace_id = execution_context.trace_id
+            if execution_context.parent_context is not None:
+                self._execution_context.parent_context = execution_context.parent_context
+
+        if self._execution_context.parent_context is None:
+            self._execution_context.parent_context = ExternalParentContext(span_id=uuid4())
+            # Only generate a new trace_id if one wasn't explicitly provided (i.e., if it's the default zero UUID)
+            if self._execution_context.trace_id == UUID("00000000-0000-0000-0000-000000000000"):
+                self._execution_context.trace_id = uuid4()
+            # Propagate the updated context back to the global execution context
+            set_execution_context(self._execution_context)
+
         self._generated_files = generated_files
 
     @cached_property
@@ -44,10 +81,65 @@ class WorkflowContext:
     def generated_files(self) -> Optional[dict[str, str]]:
         return self._generated_files
 
+    @cached_property
+    def namespace(self) -> Optional[str]:
+        return self._namespace
+
     @cached_property
     def node_output_mocks_map(self) -> Dict[Type[BaseOutputs], List[MockNodeExecution]]:
         return self._node_output_mocks_map
 
+    @property
+    def store_class(self) -> Type[Store]:
+        return self._store_class
+
+    @property
+    def monitoring_url(self) -> Optional[str]:
+        """
+        Get the base monitoring URL for this workflow context.
+
+        Returns the base URL to view executions in Vellum UI.
+        """
+        # Get UI URL from the Vellum client's API URL
+        api_url = self.vellum_client._client_wrapper.get_environment().default
+        return self._get_ui_url_from_api_url(api_url)
+
+    def get_monitoring_url(self, span_id: str) -> Optional[str]:
+        """
+        Generate the monitoring URL for this workflow execution.
+
+        Args:
+        span_id: The span ID from the workflow execution result.
+
+        Returns:
+        The URL to view execution details in Vellum UI, or None if monitoring is disabled.
+        """
+        base_url = self.monitoring_url
+        if base_url is None:
+            return None
+
+        return f"{base_url}/workflows/executions/{span_id}"
+
+    def _get_ui_url_from_api_url(self, api_url: str) -> str:
+        """
+        Convert an API URL to the corresponding UI URL.
+
+        Args:
+        api_url: The API base URL (e.g., https://api.vellum.ai or http://localhost:8000)
+
+        Returns:
+        The corresponding UI URL (e.g., https://app.vellum.ai or http://localhost:3000)
+        """
+        if "localhost" in api_url:
+            # For local development: localhost:8000 (API) -> localhost:3000 (UI)
+            return api_url.replace(":8000", ":3000")
+        elif "api.vellum.ai" in api_url:
+            # For production: api.vellum.ai -> app.vellum.ai
+            return api_url.replace("api.vellum.ai", "app.vellum.ai")
+        else:
+            # For custom domains, assume the same pattern: api.* -> app.*
+            return api_url.replace("api.", "app.", 1)
+
     def _emit_subworkflow_event(self, event: "WorkflowEvent") -> None:
         if self._event_queue:
             self._event_queue.put(event)
@@ -75,6 +167,133 @@ class WorkflowContext:
     def _get_all_node_output_mocks(self) -> List[MockNodeExecution]:
         return [mock for mocks in self._node_output_mocks_map.values() for mock in mocks]
 
+    def resolve_workflow_deployment(
+        self, deployment_name: str, release_tag: str, state: "BaseState"
+    ) -> Optional[Tuple[Type["BaseWorkflow"], Optional[WorkflowDeploymentMetadata]]]:
+        """
+        Resolve a workflow deployment by name and release tag.
+
+        Args:
+            deployment_name: The name of the workflow deployment
+            release_tag: The release tag to resolve
+            state: The base state to pass to the workflow
+
+        Returns:
+            Tuple of (BaseWorkflow class, deployment metadata) if found
+        """
+        if not self._generated_files or not self._namespace:
+            return None
+
+        expected_prefix = generate_workflow_deployment_prefix(deployment_name, release_tag)
+        WorkflowClass: Optional[Type["BaseWorkflow"]] = None
+
+        try:
+            from vellum.workflows.workflows.base import BaseWorkflow
+
+            WorkflowClass = BaseWorkflow.load_from_module(f"{self.namespace}.{expected_prefix}")
+            WorkflowClass.is_dynamic = True
+        except Exception:
+            pass
+
+        if WorkflowClass is None:
+            try:
+                major_version = __version__.split(".")[0]
+                version_range = f">={major_version}.0.0,<={__version__}"
+
+                response = self.vellum_client.workflows.pull(
+                    deployment_name,
+                    release_tag=release_tag,
+                    version=version_range,
+                    request_options={"additional_headers": {"X-Vellum-Always-Success": "true"}},
+                )
+
+                if isinstance(response, dict) and response.get("success") is False:
+                    return None
+
+                zip_bytes = b"".join(response)
+                pulled_files = extract_zip_files(zip_bytes)
+
+                for file_name, content in pulled_files.items():
+                    prefixed_file_name = f"{expected_prefix}/{file_name}"
+                    self._generated_files[prefixed_file_name] = content
+
+                from vellum.workflows.workflows.base import BaseWorkflow
+
+                WorkflowClass = BaseWorkflow.load_from_module(f"{self.namespace}.{expected_prefix}")
+                WorkflowClass.is_dynamic = True
+
+            except Exception:
+                pass
+
+        # If we successfully loaded the workflow class, fetch metadata from generated_files (important-comment)
+        # This ensures we pick up metadata.json if a pull just populated it
+        if WorkflowClass is not None:
+            deployment_metadata = self._fetch_deployment_metadata(deployment_name, release_tag)
+            return (WorkflowClass, deployment_metadata)
+
+        return None
+
+    def _fetch_deployment_metadata(
+        self, deployment_name: str, release_tag: str
+    ) -> Optional[WorkflowDeploymentMetadata]:
+        """
+        Fetch deployment metadata from metadata.json in generated_files.
+
+        Args:
+            deployment_name: The name of the workflow deployment
+            release_tag: The release tag name
+
+        Returns:
+            WorkflowDeploymentMetadata if successful, None otherwise
+        """
+        if not self._generated_files:
+            return None
+
+        expected_prefix = generate_workflow_deployment_prefix(deployment_name, release_tag)
+        metadata_path = f"{expected_prefix}/metadata.json"
+
+        metadata_content = self._generated_files.get(metadata_path)
+        if metadata_content:
+            try:
+                metadata_json = json.loads(metadata_content)
+
+                deployment_id = metadata_json.get("deployment_id")
+                deployment_name_from_metadata = metadata_json.get("deployment_name")
+                deployment_history_item_id = metadata_json.get("deployment_history_item_id")
+                release_tag_id = metadata_json.get("release_tag_id")
+                release_tag_name_from_metadata = metadata_json.get("release_tag_name")
+                workflow_version_id = metadata_json.get("workflow_version_id")
+
+                if all(
+                    [
+                        deployment_id,
+                        deployment_name_from_metadata,
+                        deployment_history_item_id,
+                        release_tag_id,
+                        release_tag_name_from_metadata,
+                        workflow_version_id,
+                    ]
+                ):
+                    return WorkflowDeploymentMetadata(
+                        deployment_id=UUID(deployment_id),
+                        deployment_name=deployment_name_from_metadata,
+                        deployment_history_item_id=UUID(deployment_history_item_id),
+                        release_tag_id=UUID(release_tag_id),
+                        release_tag_name=release_tag_name_from_metadata,
+                        workflow_version_id=UUID(workflow_version_id),
+                    )
+            except (json.JSONDecodeError, ValueError, KeyError):
+                # If we fail to parse metadata, return None - the workflow can still run
+                # but won't have the full parent context hierarchy
+                pass
+
+        return None
+
     @classmethod
     def create_from(cls, context):
-        return cls(vellum_client=context.vellum_client, generated_files=context.generated_files)
+        return cls(
+            vellum_client=context.vellum_client,
+            generated_files=context.generated_files,
+            namespace=context.namespace,
+            store_class=context.store_class,
+        )

vellum/workflows/state/delta.py

@@ -0,0 +1,20 @@
+from typing import Any, Literal, Union
+
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
+
+
+class BaseStateDelta(UniversalBaseModel):
+    name: str
+
+
+class SetStateDelta(BaseStateDelta):
+    delta: Any
+    delta_type: Literal["SET_STATE"] = "SET_STATE"
+
+
+class AppendStateDelta(BaseStateDelta):
+    delta: Any
+    delta_type: Literal["APPEND_STATE"] = "APPEND_STATE"
+
+
+StateDelta = Union[SetStateDelta, AppendStateDelta]

vellum/workflows/state/encoder.py

@@ -1,78 +1,3 @@
-from dataclasses import asdict, is_dataclass
-from datetime import datetime
-import enum
-import inspect
-from json import JSONEncoder
-from queue import Queue
-from uuid import UUID
-from typing import Any, Callable, Dict, Type
+from vellum.utils.json_encoder import VellumJsonEncoder as DefaultStateEncoder  # noqa: F401
 
-from pydantic import BaseModel
-
-from vellum.workflows.constants import undefined
-from vellum.workflows.inputs.base import BaseInputs
-from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
-from vellum.workflows.ports.port import Port
-from vellum.workflows.state.base import BaseState, NodeExecutionCache
-from vellum.workflows.utils.functions import compile_function_definition
-
-
-class DefaultStateEncoder(JSONEncoder):
-    encoders: Dict[Type, Callable] = {}
-
-    def default(self, obj: Any) -> Any:
-        if isinstance(obj, BaseState):
-            return dict(obj)
-
-        if isinstance(obj, (BaseInputs, BaseOutputs)):
-            return {descriptor.name: value for descriptor, value in obj if value is not undefined}
-
-        if isinstance(obj, (BaseOutput, Port)):
-            return obj.serialize()
-
-        if isinstance(obj, NodeExecutionCache):
-            return obj.dump()
-
-        if isinstance(obj, UUID):
-            return str(obj)
-
-        if isinstance(obj, set):
-            return list(obj)
-
-        if isinstance(obj, BaseModel):
-            return obj.model_dump()
-
-        if isinstance(obj, datetime):
-            return obj.isoformat()
-
-        if isinstance(obj, enum.Enum):
-            return obj.value
-
-        if isinstance(obj, Queue):
-            return list(obj.queue)
-
-        if is_dataclass(obj):
-            # Technically, obj is DataclassInstance | type[DataclassInstance], but asdict expects a DataclassInstance
-            # in practice, we only ever pass the former
-            return asdict(obj)  # type: ignore[call-overload]
-
-        if isinstance(obj, type):
-            return str(obj)
-
-        if callable(obj):
-            function_definition = compile_function_definition(obj)
-            try:
-                source_code = inspect.getsource(obj)
-            except Exception:
-                source_code = f"<source code not available for {obj.__name__}>"
-
-            return {
-                "type": "CODE_EXECUTION",
-                "definition": function_definition,
-                "src": source_code,
-            }
-
-        if obj.__class__ in self.encoders:
-            return self.encoders[obj.__class__](obj)
-
-        return super().default(obj)
+__all__ = ["DefaultStateEncoder"]

vellum/workflows/state/tests/test_state.py

@@ -2,12 +2,15 @@ import pytest
 from copy import deepcopy
 import json
 from queue import Queue
-from typing import Dict, cast
+from typing import Dict, List, cast
 
+from vellum.utils.json_encoder import VellumJsonEncoder
+from vellum.workflows.constants import undefined
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.state.encoder import DefaultStateEncoder
+from vellum.workflows.state.delta import SetStateDelta, StateDelta
+from vellum.workflows.types.code_execution_node_wrappers import DictWrapper
 
 
 @pytest.fixture()
@@ -28,7 +31,7 @@ class MockState(BaseState):
 
     def __init__(self, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
-        self.__snapshot_callback__ = lambda _: self.__mock_snapshot__()
+        self.__snapshot_callback__ = lambda _, __: self.__mock_snapshot__()
 
     def __mock_snapshot__(self) -> None:
         self.__snapshot_count__ += 1
@@ -47,7 +50,7 @@ class MockNode(BaseNode):
         baz: str
 
 
-MOCK_NODE_OUTPUT_ID = "e4dc3136-0c27-4bda-b3ab-ea355d5219d6"
+MOCK_NODE_OUTPUT_ID = "a864331d-9b1c-43db-90e7-6b1998d8179d"
 
 
 def test_state_snapshot__node_attribute_edit():
@@ -144,7 +147,7 @@ def test_state_json_serialization__with_node_output_updates():
     state.meta.node_outputs[MockNode.Outputs.baz] = "hello"
 
     # WHEN we serialize the state
-    json_state = json.loads(json.dumps(state, cls=DefaultStateEncoder))
+    json_state = json.loads(json.dumps(state, cls=VellumJsonEncoder))
 
     # THEN the state is serialized correctly
     assert json_state["meta"]["node_outputs"] == {MOCK_NODE_OUTPUT_ID: "hello"}
@@ -185,7 +188,7 @@ def test_state_json_serialization__with_queue():
     state.meta.node_outputs[MockNode.Outputs.baz] = queue
 
     # WHEN we serialize the state
-    json_state = json.loads(json.dumps(state, cls=DefaultStateEncoder))
+    json_state = json.loads(json.dumps(state, cls=VellumJsonEncoder))
 
     # THEN the state is serialized correctly with the queue turned into a list
     assert json_state["meta"]["node_outputs"] == {MOCK_NODE_OUTPUT_ID: ["test1", "test2"]}
@@ -201,7 +204,7 @@ def test_state_snapshot__deepcopy_fails__logs_error(mock_deepcopy, mock_logger):
     # AND we have a snapshot callback that we are tracking
     snapshot_count = 0
 
-    def __snapshot_callback__(state: "BaseState") -> None:
+    def __snapshot_callback__(state: "BaseState", deltas: List[StateDelta]) -> None:
         nonlocal snapshot_count
         snapshot_count += 1
 
@@ -220,11 +223,23 @@ def test_state_snapshot__deepcopy_fails__logs_error(mock_deepcopy, mock_logger):
     mock_deepcopy.side_effect = side_effect
 
     # WHEN we snapshot the state twice
-    state.__snapshot__()
-    state.__snapshot__()
+    state.__snapshot__(SetStateDelta(name="foo", delta="bar"))
+    state.__snapshot__(SetStateDelta(name="foo", delta="baz"))
 
     # THEN we were able to invoke the callback once
     assert snapshot_count == 1
 
     # AND alert sentry once
     assert mock_logger.exception.call_count == 1
+
+
+def test_state_deepcopy_handles_undefined_values():
+    # GIVEN a state with undefined values in node outputs
+    state = MockState(foo="bar")
+    state.meta.node_outputs[MockNode.Outputs.baz] = DictWrapper({"foo": undefined})
+
+    # WHEN we deepcopy the state
+    deepcopied_state = deepcopy(state)
+
+    # THEN the undefined values are preserved
+    assert deepcopied_state.meta.node_outputs[MockNode.Outputs.baz] == {"foo": undefined}

vellum/workflows/tests/test_dataset_row.py

@@ -0,0 +1,190 @@
+from typing import Optional
+
+from vellum.client.types.chat_message import ChatMessage
+from vellum.workflows.inputs.base import BaseInputs
+from vellum.workflows.inputs.dataset_row import DatasetRow
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.outputs.base import BaseOutputs
+
+
+def test_dataset_row_serialization():
+    """
+    Test that DatasetRow can be properly serialized to JSON and back.
+    """
+
+    class TestInputs(BaseInputs):
+        message: str
+        count: int
+        optional_field: Optional[str] = None
+        chat_history: list[ChatMessage]
+
+    test_inputs = TestInputs(
+        message="Hello World",
+        count=42,
+        optional_field="test",
+        chat_history=[ChatMessage(text="Hello", role="USER"), ChatMessage(text="Hi there!", role="ASSISTANT")],
+    )
+    dataset_row = DatasetRow(label="test_label", inputs=test_inputs)
+
+    serialized_dict = dataset_row.model_dump()
+
+    assert "label" in serialized_dict
+    assert "inputs" in serialized_dict
+    assert serialized_dict["label"] == "test_label"
+
+    inputs_data = serialized_dict["inputs"]
+    assert inputs_data["message"] == "Hello World"
+    assert inputs_data["count"] == 42
+    assert inputs_data["optional_field"] == "test"
+    assert "chat_history" in inputs_data
+    assert len(inputs_data["chat_history"]) == 2
+    assert inputs_data["chat_history"][0]["text"] == "Hello"
+    assert inputs_data["chat_history"][0]["role"] == "USER"
+    assert inputs_data["chat_history"][1]["text"] == "Hi there!"
+    assert inputs_data["chat_history"][1]["role"] == "ASSISTANT"
+
+
+def test_dataset_row_dict_serialization():
+    """
+    Test that DatasetRow can be properly converted to dict.
+    """
+
+    class SimpleInputs(BaseInputs):
+        text: str
+
+    simple_inputs = SimpleInputs(text="sample text")
+    dataset_row = DatasetRow(label="simple_label", inputs=simple_inputs)
+
+    result_dict = dataset_row.model_dump()
+
+    assert result_dict["label"] == "simple_label"
+    assert result_dict["inputs"]["text"] == "sample text"
+
+
+def test_dataset_row_with_multiple_fields():
+    """
+    Test that DatasetRow works with BaseInputs that have multiple fields.
+    """
+
+    class MultiFieldInputs(BaseInputs):
+        text_field: str
+        number_field: int
+        optional_field: Optional[str] = None
+
+    multi_inputs = MultiFieldInputs(text_field="test_text", number_field=456, optional_field="optional_value")
+    dataset_row = DatasetRow(label="multi_field_test", inputs=multi_inputs)
+
+    result_dict = dataset_row.model_dump()
+
+    assert result_dict["label"] == "multi_field_test"
+    assert result_dict["inputs"]["text_field"] == "test_text"
+    assert result_dict["inputs"]["number_field"] == 456
+    assert result_dict["inputs"]["optional_field"] == "optional_value"
+
+
+def test_dataset_row_with_default_inputs():
+    """
+    Test that DatasetRow works with BaseInputs that have default values.
+    """
+
+    class InputsWithDefaults(BaseInputs):
+        required_field: str
+        optional_with_default: str = "default_value"
+
+    inputs_with_defaults = InputsWithDefaults(required_field="required_value")
+    dataset_row = DatasetRow(label="defaults_test", inputs=inputs_with_defaults)
+
+    serialized_dict = dataset_row.model_dump()
+
+    assert serialized_dict["label"] == "defaults_test"
+    assert serialized_dict["inputs"]["required_field"] == "required_value"
+    assert serialized_dict["inputs"]["optional_with_default"] == "default_value"
+
+
+def test_dataset_row_without_inputs():
+    """
+    Test that DatasetRow can be created with only a label (no inputs parameter).
+    """
+
+    dataset_row = DatasetRow(label="test_label_only")
+
+    serialized_dict = dataset_row.model_dump()
+
+    assert serialized_dict["label"] == "test_label_only"
+    assert serialized_dict["inputs"] == {}
+
+    assert isinstance(dataset_row.inputs, BaseInputs)
+
+
+def test_dataset_row_with_empty_inputs():
+    """
+    Test that DatasetRow can be created with explicitly empty BaseInputs.
+    """
+
+    # GIVEN a DatasetRow with explicitly empty BaseInputs
+    dataset_row = DatasetRow(label="test_label", inputs=BaseInputs())
+
+    serialized_dict = dataset_row.model_dump()
+
+    assert serialized_dict["label"] == "test_label"
+    assert serialized_dict["inputs"] == {}
+
+
+def test_dataset_row_with_dict_inputs():
+    """
+    Test that DatasetRow can accept a dict for the inputs parameter.
+    """
+
+    # GIVEN a dict with input data
+    inputs_dict = {"message": "Hello World", "count": 42}
+
+    dataset_row = DatasetRow(label="test_label", inputs=inputs_dict)
+
+    # THEN the inputs should remain as a dict
+    assert isinstance(dataset_row.inputs, dict)
+
+    # AND the serialized dict should contain the correct data
+    serialized_dict = dataset_row.model_dump()
+    assert serialized_dict["label"] == "test_label"
+    assert serialized_dict["inputs"]["message"] == "Hello World"
+    assert serialized_dict["inputs"]["count"] == 42
+
+
+def test_dataset_row_with_node_output_mocks():
+    """
+    Test that DatasetRow can be created with node_output_mocks and properly serialized.
+    """
+
+    # GIVEN a node with outputs
+    class DummyNode(BaseNode):
+        class Outputs(BaseOutputs):
+            result: str
+
+    # AND a DatasetRow with node output mocks
+    mock_output = DummyNode.Outputs(result="mocked output")
+
+    class TestInputs(BaseInputs):
+        message: str
+
+    test_inputs = TestInputs(message="test message")
+
+    dataset_row = DatasetRow(label="test_with_mocks", inputs=test_inputs, node_output_mocks=[mock_output])
+
+    serialized_dict = dataset_row.model_dump()
+
+    # THEN the serialized dict should contain the label and inputs
+    assert serialized_dict["label"] == "test_with_mocks"
+    assert serialized_dict["inputs"]["message"] == "test message"
+
+    # AND the node_output_mocks should be present in the serialized dict
+    assert "node_output_mocks" in serialized_dict
+    assert serialized_dict["node_output_mocks"] is not None
+    assert len(serialized_dict["node_output_mocks"]) == 1
+
+    # AND the mock output should be serialized as a dict with the correct structure
+    mock_data = serialized_dict["node_output_mocks"][0]
+    assert mock_data == {
+        "node_id": str(DummyNode.__id__),
+        "when_condition": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": True}},
+        "then_outputs": {"result": "mocked output"},
+    }