vellum-ai 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

vellum/types/node_input_compiled_video_value.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.node_input_compiled_video_value import *

vellum/types/prompt_request_audio_input.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.prompt_request_audio_input import *

vellum/types/prompt_request_document_input.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.prompt_request_document_input import *

vellum/types/prompt_request_image_input.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.prompt_request_image_input import *

vellum/types/prompt_request_video_input.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.prompt_request_video_input import *

vellum/types/video_input_request.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.video_input_request import *

vellum/workflows/context.py

@@ -1,10 +1,11 @@
 from contextlib import contextmanager
 from dataclasses import field
 import threading
-from uuid import UUID
+from uuid import UUID, uuid4
 from typing import Iterator, Optional, cast
 
 from vellum.client.core import UniversalBaseModel
+from vellum.workflows.events.context import MonitoringContextStore
 from vellum.workflows.events.types import ParentContext
 
 
@@ -17,16 +18,37 @@ _CONTEXT_KEY = "_execution_context"
 
 local = threading.local()
 
+monitoring_context_store = MonitoringContextStore()
+
 
 def get_execution_context() -> ExecutionContext:
-    """Retrieve the current execution context."""
-    return getattr(local, _CONTEXT_KEY, ExecutionContext())
+    """Get the current monitoring execution context, with intelligent fallback."""
+    context = getattr(local, _CONTEXT_KEY, ExecutionContext())
+    if context.parent_context:
+        return context
+
+    # If no thread-local context, try to restore from global store using current trace_id
+    context = monitoring_context_store.retrieve_context()
+    if context and context.parent_context:
+        set_execution_context(context)
+        return context
+    return ExecutionContext()
 
 
 def set_execution_context(context: ExecutionContext) -> None:
-    """Set the current execution context."""
+    """Set the current monitoring execution context and persist it for cross-boundary access."""
     setattr(local, _CONTEXT_KEY, context)
 
+    # Always store in global store for cross-thread access
+    monitoring_context_store.store_context(context)
+
+
+def clear_execution_context() -> None:
+    """Clear the current monitoring execution context."""
+    if hasattr(local, _CONTEXT_KEY):
+        delattr(local, _CONTEXT_KEY)
+    monitoring_context_store.clear_context()
+
 
 def get_parent_context() -> ParentContext:
     return cast(ParentContext, get_execution_context().parent_context)
@@ -38,11 +60,7 @@ def execution_context(
 ) -> Iterator[None]:
     """Context manager for handling execution context."""
     prev_context = get_execution_context()
-    set_trace_id = (
-        prev_context.trace_id
-        if int(prev_context.trace_id)
-        else trace_id or UUID("00000000-0000-0000-0000-000000000000")
-    )
+    set_trace_id = prev_context.trace_id if int(prev_context.trace_id) else trace_id or uuid4()
     set_parent_context = parent_context or prev_context.parent_context
     set_context = ExecutionContext(parent_context=set_parent_context, trace_id=set_trace_id)
     try:

vellum/workflows/events/context.py

@@ -2,110 +2,85 @@
 
 import threading
 from uuid import UUID
-from typing import Dict, Optional
+from typing import TYPE_CHECKING, Dict, Optional
 
-from vellum.workflows.context import ExecutionContext
+if TYPE_CHECKING:
+    from vellum.workflows.context import ExecutionContext
 
 DEFAULT_TRACE_ID = UUID("00000000-0000-0000-0000-000000000000")
 
-# Thread-local storage for monitoring execution context
-_monitoring_execution_context: threading.local = threading.local()
-# Thread-local storage for current span_id
-_current_span_id: threading.local = threading.local()
 
-
-class _MonitoringContextStore:
+class MonitoringContextStore:
     """
     thread-safe storage for monitoring contexts.
     handles context persistence and retrieval across threads.
-    relies on the execution context manager for manual retrieval
     """
 
     def __init__(self):
         self._lock = threading.Lock()
-        self._contexts: Dict[str, ExecutionContext] = {}
-        self._thread_contexts: Dict[int, ExecutionContext] = {}
-        self._current_trace_id: Optional[UUID] = None
-
-    def set_current_trace_id(self, trace_id: UUID) -> None:
-        """Set the current active trace_id that should be used by all threads."""
-        if trace_id != DEFAULT_TRACE_ID:
-            with self._lock:
-                self._current_trace_id = trace_id
+        self._contexts: Dict[str, "ExecutionContext"] = {}
 
     def get_current_trace_id(self) -> Optional[UUID]:
         """Get the current active trace_id that should be used by all threads."""
         with self._lock:
-            return self._current_trace_id
-
-    def set_current_span_id(self, span_id: UUID) -> None:
-        """Set the current active span_id for this thread."""
-        _current_span_id.span_id = span_id
-
-    def get_current_span_id(self) -> Optional[UUID]:
-        """Get the current active span_id for this thread."""
-        return getattr(_current_span_id, "span_id", None)
-
-    def store_context(self, context: Optional[ExecutionContext]) -> None:
-        """Store monitoring parent context using multiple keys for reliable retrieval."""
-        if not context or context.parent_context is None:
-            return
+            current_context = self.retrieve_context()
+            return current_context.trace_id if current_context else None
 
+    def store_context(self, context: "ExecutionContext") -> None:
+        """Store monitoring parent context using trace:span:thread keys."""
         thread_id = threading.get_ident()
-        trace_id = self.get_current_trace_id()
-        if context.trace_id != DEFAULT_TRACE_ID and trace_id is None:
-            self.set_current_trace_id(context.trace_id)
+        current_thread = threading.current_thread()
+        trace_id = context.trace_id
 
         with self._lock:
-            # Use trace:span:thread for unique context storage
-            trace_span_thread_key = (
-                f"trace:{str(trace_id)}:span:{str(context.parent_context.span_id)}:thread:{thread_id}"
-            )
-            self._contexts[trace_span_thread_key] = context
-
-    def retrieve_context(self, trace_id: UUID, span_id: Optional[UUID] = None) -> Optional[ExecutionContext]:
-        """Retrieve monitoring parent context with multiple fallback strategies."""
-        thread_id = threading.get_ident()
-        with self._lock:
-            if not span_id:
-                span_id = getattr(_current_span_id, "span_id", None)
-                if not span_id:
-                    return None
-
-            span_key = f"trace:{str(trace_id)}:span:{str(span_id)}:thread:{thread_id}"
-            if span_key in self._contexts:
-                result = self._contexts[span_key]
-                return result
+            # Get span_id from RelationalThread first
+            span_id = None
 
-        return None
+            if context.parent_context and hasattr(context.parent_context, "span_id"):
+                span_id = context.parent_context.span_id
 
+            if not span_id and hasattr(current_thread, "get_parent_span_id"):
+                span_id = current_thread.get_parent_span_id()
 
-# Global instance for cross-boundary context persistence
-_monitoring_context_store = _MonitoringContextStore()
+            # Always use trace:span:thread format - require span_id
+            if span_id:
+                context_key = f"trace:{str(trace_id)}:span:{str(span_id)}:thread:{thread_id}"
+                self._contexts[context_key] = context
 
+    def retrieve_context(self) -> Optional["ExecutionContext"]:
+        """Retrieve monitoring parent context using trace:span:thread keys."""
+        current_thread = threading.current_thread()
+        current_thread_id = current_thread.ident
 
-def get_monitoring_execution_context() -> ExecutionContext:
-    """Get the current monitoring execution context, with intelligent fallback."""
-    if hasattr(_monitoring_execution_context, "context"):
-        context = _monitoring_execution_context.context
-        if context.trace_id != DEFAULT_TRACE_ID and context.parent_context:
-            return context
+        # Get trace_id and span_id directly from the thread if it's a RelationalThread
+        trace_id = None
+        span_id = None
+        if hasattr(current_thread, "get_trace_id"):
+            trace_id = current_thread.get_trace_id()
+        if hasattr(current_thread, "get_parent_span_id"):
+            span_id = current_thread.get_parent_span_id()
 
-    # If no thread-local context, try to restore from global store using current trace_id
-    trace_id = _monitoring_context_store.get_current_trace_id()
-    span_id = _current_span_id.span_id if hasattr(_current_span_id, "span_id") else None
-    if trace_id:
-        if trace_id != DEFAULT_TRACE_ID:
-            context = _monitoring_context_store.retrieve_context(trace_id, span_id)
-            if context:
-                _monitoring_execution_context.context = context
-                return context
-    return ExecutionContext()
+        # Require both trace_id and span_id - no fallback searching
+        if not trace_id or not span_id:
+            return None
 
+        with self._lock:
+            # Try current thread
+            current_key = f"trace:{str(trace_id)}:span:{str(span_id)}:thread:{current_thread_id}"
+            if current_key in self._contexts:
+                return self._contexts[current_key]
+
+            # Try parent thread with same trace and span
+            if hasattr(current_thread, "get_parent_thread"):
+                parent_thread_id = current_thread.get_parent_thread()
+                if parent_thread_id:
+                    parent_key = f"trace:{str(trace_id)}:span:{str(span_id)}:thread:{parent_thread_id}"
+                    if parent_key in self._contexts:
+                        return self._contexts[parent_key]
 
-def set_monitoring_execution_context(context: ExecutionContext) -> None:
-    """Set the current monitoring execution context and persist it for cross-boundary access."""
-    _monitoring_execution_context.context = context
+        return None
 
-    if context.trace_id and context.parent_context:
-        _monitoring_context_store.store_context(context)
+    def clear_context(self):
+        """Clear all stored contexts."""
+        with self._lock:
+            self._contexts.clear()

vellum/workflows/events/node.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, Generic, List, Literal, Optional, Set, Type, Union
 
-from pydantic import field_serializer, model_serializer
+from pydantic import SerializationInfo, field_serializer, model_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.errors import WorkflowError
@@ -82,8 +82,8 @@ class NodeExecutionStreamingEvent(_BaseNodeEvent):
         return self.body.invoked_ports
 
     @model_serializer(mode="plain", when_used="json")
-    def serialize_model(self) -> Any:
-        serialized = super().serialize_model()
+    def serialize_model(self, info: SerializationInfo) -> Any:
+        serialized = super().serialize_model(info)
         if (
             "body" in serialized
             and isinstance(serialized["body"], dict)
@@ -127,8 +127,8 @@ class NodeExecutionFulfilledEvent(_BaseNodeEvent, Generic[OutputsType]):
         return self.body.mocked
 
     @model_serializer(mode="plain", when_used="json")
-    def serialize_model(self) -> Any:
-        serialized = super().serialize_model()
+    def serialize_model(self, info: SerializationInfo) -> Any:
+        serialized = super().serialize_model(info)
         if (
             "body" in serialized
             and isinstance(serialized["body"], dict)

vellum/workflows/events/relational_threads.py

@@ -0,0 +1,41 @@
+import threading
+from uuid import UUID
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from vellum.workflows.context import ExecutionContext
+
+
+class RelationalThread(threading.Thread):
+    _parent_thread: Optional[int] = None
+    _trace_id: Optional[UUID] = None
+    _parent_span_id: Optional[UUID] = None
+
+    def __init__(self, *args, execution_context: Optional["ExecutionContext"] = None, **kwargs):
+        self._collect_parent_context(execution_context)
+        threading.Thread.__init__(self, *args, **kwargs)
+
+    def _collect_parent_context(self, execution_context: Optional["ExecutionContext"] = None) -> None:
+        """Collect parent thread ID, trace ID, and parent span ID from passed execution context."""
+        self._parent_thread = threading.get_ident()
+
+        # Only use explicitly passed execution context
+        if execution_context:
+            self._trace_id = execution_context.trace_id
+            self._parent_span_id = (
+                execution_context.parent_context.span_id
+                if execution_context.parent_context and hasattr(execution_context.parent_context, "span_id")
+                else None
+            )
+        else:
+            self._trace_id = None
+            self._parent_span_id = None
+
+    def get_parent_thread(self) -> Optional[int]:
+        return self._parent_thread
+
+    def get_trace_id(self) -> Optional[UUID]:
+        return self._trace_id
+
+    def get_parent_span_id(self) -> Optional[UUID]:
+        return self._parent_span_id

vellum/workflows/events/tests/test_basic_workflow.py

@@ -0,0 +1,50 @@
+import logging
+from uuid import uuid4
+
+from vellum.workflows import BaseWorkflow
+from vellum.workflows.context import execution_context
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
+
+
+class StartNode(BaseNode):
+    pass
+
+
+class TrivialWorkflow(BaseWorkflow):
+    graph = StartNode
+
+
+logger = logging.getLogger(__name__)
+
+
+def test_basic_workflow_monitoring_context_flow():
+    """Test that monitoring creates the correct workflow→node context hierarchy using streamed events.
+    What's missing:
+    - span_id from previous event mapping to parent context
+    """
+
+    workflow = TrivialWorkflow()
+
+    with execution_context(trace_id=uuid4()):
+        events = list(workflow.stream(event_filter=all_workflow_event_filter))
+
+    # Verify workflow succeeded
+    assert len(events) >= 2
+    assert events[0].name == "workflow.execution.initiated"
+    assert events[-1].name == "workflow.execution.fulfilled"
+
+    # Collect all events with parent context
+    events_with_parent = [event for event in events if event.parent is not None]
+
+    assert len(events_with_parent) > 0, "Expected at least some events with parent context"
+
+    # Filter for node events
+    node_events = [event for event in events if event.name.startswith("node.")]
+    assert len(node_events) > 0, "Expected at least some node events"
+
+    # Verify each node event has the workflow as its parent context
+    for event in node_events:
+        assert event.parent is not None, "Node event should have parent context"
+        assert event.parent.type == "WORKFLOW", "Node event parent should be workflow"
+        assert event.parent.workflow_definition.name == "TrivialWorkflow", "Parent workflow name mismatch"

vellum/workflows/events/tests/test_event.py

@@ -93,6 +93,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     "workflow_version_exec_config": None,
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -156,6 +157,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     "type": "WORKFLOW_NODE",
                     "span_id": "123e4567-e89b-12d3-a456-426614174000",
                 },
+                "links": None,
             },
         ),
         (
@@ -191,6 +193,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     },
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -224,6 +227,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     },
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -259,6 +263,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     },
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -294,6 +299,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     },
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -334,6 +340,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     "mocked": None,
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -372,6 +379,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     "mocked": None,
                 },
                 "parent": None,
+                "links": None,
             },
         ),
         (
@@ -405,6 +413,7 @@ mock_node_uuid = str(uuid4_from_hash(MockNode.__qualname__))
                     "mocked": True,
                 },
                 "parent": None,
+                "links": None,
             },
         ),
     ],

vellum/workflows/events/types.py

@@ -1,12 +1,13 @@
 from datetime import datetime
 import json
 from uuid import UUID, uuid4
-from typing import Annotated, Any, Literal, Optional, Union, get_args
+from typing import Annotated, Any, List, Literal, Optional, Union, get_args
 
 from pydantic import Field, GetCoreSchemaHandler, Tag, ValidationInfo
 from pydantic_core import CoreSchema, core_schema
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.types.span_link import SpanLink
 from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.definition import VellumCodeResourceDefinition
 from vellum.workflows.types.utils import datetime_now
@@ -164,3 +165,4 @@ class BaseEvent(UniversalBaseModel):
     trace_id: UUID
     span_id: UUID
     parent: Optional[ParentContext] = None
+    links: Optional[List[SpanLink]] = None

vellum/workflows/events/workflow.py

@@ -2,7 +2,7 @@ from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Literal, Optional, Type, Union
 from typing_extensions import TypeGuard
 
-from pydantic import field_serializer
+from pydantic import SerializationInfo, field_serializer
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
 from vellum.workflows.errors import WorkflowError
@@ -101,6 +101,17 @@ class WorkflowExecutionInitiatedEvent(_BaseWorkflowEvent, Generic[InputsType, St
     def initial_state(self) -> Optional[StateType]:
         return self.body.initial_state
 
+    @field_serializer("body")
+    def serialize_body(
+        self, body: WorkflowExecutionInitiatedBody[InputsType, StateType], info: SerializationInfo
+    ) -> WorkflowExecutionInitiatedBody[InputsType, StateType]:
+        context = info.context if info and hasattr(info, "context") else {}
+        if context and "event_enricher" in context and callable(context["event_enricher"]):
+            event = context["event_enricher"](self)
+            return event.body
+        else:
+            return body
+
 
 class WorkflowExecutionStreamingBody(_BaseWorkflowExecutionBody):
     output: BaseOutput

vellum/workflows/expressions/contains.py

@@ -35,4 +35,11 @@ class ContainsExpression(BaseDescriptor[bool], Generic[LHS, RHS]):
             )
 
         rhs = resolve_value(self._rhs, state)
+
+        if isinstance(rhs, dict):
+            raise InvalidExpressionException(
+                "Cannot use dict as right-hand side of contains operation. "
+                "Use dict keys/values or convert to strings for comparison."
+            )
+
         return rhs in lhs