vellum-ai 0.11.0__py3-none-any.whl → 0.11.3__py3-none-any.whl

vellum/workflows/nodes/core/map_node/node.py

@@ -3,14 +3,18 @@ from queue import Empty, Queue
 from threading import Thread
 from typing import TYPE_CHECKING, Callable, Dict, Generic, List, Optional, Tuple, Type, TypeVar, Union, overload
 
+from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.errors.types import VellumErrorCode
+from vellum.workflows.events.types import ParentContext
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.generics import NodeType, StateType
+from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 if TYPE_CHECKING:
     from vellum.workflows import BaseWorkflow
@@ -53,7 +57,15 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
         fulfilled_iterations: List[bool] = []
         for index, item in enumerate(self.items):
             fulfilled_iterations.append(False)
-            thread = Thread(target=self._run_subworkflow, kwargs={"item": item, "index": index})
+            parent_context = get_parent_context() or self._context.parent_context
+            thread = Thread(
+                target=self._context_run_subworkflow,
+                kwargs={
+                    "item": item,
+                    "index": index,
+                    "parent_context": parent_context,
+                },
+            )
             thread.start()
 
         try:
@@ -62,6 +74,7 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
             while map_node_event := self._event_queue.get():
                 index = map_node_event[0]
                 terminal_event = map_node_event[1]
+                self._context._emit_subworkflow_event(terminal_event)
 
                 if terminal_event.name == "workflow.execution.fulfilled":
                     workflow_output_vars = vars(terminal_event.outputs)
@@ -85,12 +98,22 @@ class MapNode(BaseNode, Generic[StateType, MapNodeItemType]):
                     )
         except Empty:
             pass
-
         return self.Outputs(**mapped_items)
 
+    def _context_run_subworkflow(
+        self, *, item: MapNodeItemType, index: int, parent_context: Optional[ParentContext] = None
+    ) -> None:
+        parent_context = parent_context or self._context.parent_context
+        with execution_context(parent_context=parent_context):
+            self._run_subworkflow(item=item, index=index)
+
     def _run_subworkflow(self, *, item: MapNodeItemType, index: int) -> None:
-        subworkflow = self.subworkflow(parent_state=self.state, context=self._context)
-        events = subworkflow.stream(inputs=self.SubworkflowInputs(index=index, item=item, all_items=self.items))
+        context = WorkflowContext(_vellum_client=self._context._vellum_client)
+        subworkflow = self.subworkflow(parent_state=self.state, context=context)
+        events = subworkflow.stream(
+            inputs=self.SubworkflowInputs(index=index, item=item, all_items=self.items),
+            event_filter=all_workflow_event_filter,
+        )
 
         for event in events:
             self._event_queue.put((index, event))

vellum/workflows/nodes/core/templating_node/tests/test_templating_node.py

@@ -1,5 +1,6 @@
 import json
 
+from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.nodes.core.templating_node.node import TemplatingNode
 
 
@@ -19,3 +20,23 @@ def test_templating_node__dict_output():
 
     # THEN the output is json serializable
     assert json.loads(outputs.result) == {"key": "value"}
+
+
+def test_templating_node__execution_count_reference():
+    # GIVEN a random node
+    class OtherNode(BaseNode):
+        pass
+
+    # AND a templating node that references the execution count of the random node
+    class TemplateNode(TemplatingNode):
+        template = "{{ total }}"
+        inputs = {
+            "total": OtherNode.Execution.count,
+        }
+
+    # WHEN the node is run
+    node = TemplateNode()
+    outputs = node.run()
+
+    # THEN the output is just the total
+    assert outputs.result == "0"

vellum/workflows/nodes/displayable/api_node/node.py

@@ -18,13 +18,14 @@ class APINode(BaseAPINode):
 
     authorization_type: Optional[AuthorizationType] = None - The type of authorization to use for the API call.
     api_key_header_key: Optional[str] = None - The header key to use for the API key authorization.
-    bearer_token_value: Optional[str] = None - The bearer token value to use for the bearer token authorization.
+    bearer_token_value: Optional[Union[str, VellumSecretReference]] = None - The bearer token value to use
+    for the bearer token authorization.
     """
 
     authorization_type: Optional[AuthorizationType] = None
     api_key_header_key: Optional[str] = None
     api_key_header_value: Optional[Union[str, VellumSecretReference]] = None
-    bearer_token_value: Optional[str] = None
+    bearer_token_value: Optional[Union[str, VellumSecretReference]] = None
 
     def run(self) -> BaseAPINode.Outputs:
         headers = self.headers or {}

vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py

@@ -14,7 +14,9 @@ from vellum import (
     PromptRequestStringInput,
     VellumVariable,
 )
+from vellum.client import RequestOptions
 from vellum.workflows.constants import OMIT
+from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
@@ -48,6 +50,13 @@ class BaseInlinePromptNode(BasePromptNode, Generic[StateType]):
 
     def _get_prompt_event_stream(self) -> Iterator[AdHocExecutePromptEvent]:
         input_variables, input_values = self._compile_prompt_inputs()
+        current_parent_context = get_parent_context()
+        parent_context = current_parent_context.model_dump_json() if current_parent_context else None
+        request_options = self.request_options or RequestOptions()
+        request_options["additional_body_parameters"] = {
+            "execution_context": {"parent_context": parent_context},
+            **request_options.get("additional_body_parameters", {}),
+        }
 
         return self._context.vellum_client.ad_hoc.adhoc_execute_prompt_stream(
             ml_model=self.ml_model,

vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py

@@ -11,7 +11,9 @@ from vellum import (
     RawPromptExecutionOverridesRequest,
     StringInputRequest,
 )
+from vellum.client import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
+from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
@@ -46,6 +48,13 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
     metadata: Optional[Dict[str, Optional[Any]]] = OMIT
 
     def _get_prompt_event_stream(self) -> Iterator[ExecutePromptEvent]:
+        current_parent_context = get_parent_context()
+        parent_context = current_parent_context.model_dump() if current_parent_context else None
+        request_options = self.request_options or RequestOptions()
+        request_options["additional_body_parameters"] = {
+            "execution_context": {"parent_context": parent_context},
+            **request_options.get("additional_body_parameters", {}),
+        }
         return self._context.vellum_client.execute_prompt_stream(
             inputs=self._compile_prompt_inputs(),
             prompt_deployment_id=str(self.deployment) if isinstance(self.deployment, UUID) else None,
@@ -56,7 +65,7 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
             raw_overrides=self.raw_overrides,
             expand_raw=self.expand_raw,
             metadata=self.metadata,
-            request_options=self.request_options,
+            request_options=request_options,
         )
 
     def _compile_prompt_inputs(self) -> List[PromptDeploymentInputRequest]:

vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -13,6 +13,7 @@ from vellum import (
 )
 from vellum.core import RequestOptions
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
+from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import VellumErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases.base_subworkflow_node.node import BaseSubworkflowNode
@@ -89,6 +90,13 @@ class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateTyp
         return compiled_inputs
 
     def run(self) -> Iterator[BaseOutput]:
+        current_parent_context = get_parent_context()
+        parent_context = current_parent_context.model_dump(mode="json") if current_parent_context else None
+        request_options = self.request_options or RequestOptions()
+        request_options["additional_body_parameters"] = {
+            "execution_context": {"parent_context": parent_context},
+            **request_options.get("additional_body_parameters", {}),
+        }
         subworkflow_stream = self._context.vellum_client.execute_workflow_stream(
             inputs=self._compile_subworkflow_inputs(),
             workflow_deployment_id=str(self.deployment) if isinstance(self.deployment, UUID) else None,
@@ -97,8 +105,9 @@ class SubworkflowDeploymentNode(BaseSubworkflowNode[StateType], Generic[StateTyp
             external_id=self.external_id,
             event_types=["WORKFLOW"],
             metadata=self.metadata,
-            request_options=self.request_options,
+            request_options=request_options,
         )
+        # for some reason execution context isn't showing as an option? ^ failing mypy
 
         outputs: Optional[List[WorkflowOutput]] = None
         fulfilled_output_names: Set[str] = set()

vellum/workflows/nodes/displayable/tests/test_text_prompt_deployment_node.py

@@ -75,5 +75,5 @@ def test_text_prompt_deployment_node__basic(vellum_client):
         prompt_deployment_name="my-deployment",
         raw_overrides=OMIT,
         release_tag="LATEST",
-        request_options=None,
+        request_options={"additional_body_parameters": {"execution_context": {"parent_context": None}}},
     )

vellum/workflows/runner/runner.py

@@ -7,6 +7,7 @@ from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Type, Union
 
 from vellum.workflows.constants import UNDEF
+from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors import VellumError, VellumErrorCode
@@ -28,7 +29,7 @@ from vellum.workflows.events.node import (
     NodeExecutionRejectedBody,
     NodeExecutionStreamingBody,
 )
-from vellum.workflows.events.types import BaseEvent, ParentContext, WorkflowParentContext
+from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, WorkflowParentContext
 from vellum.workflows.events.workflow import (
     WorkflowExecutionFulfilledBody,
     WorkflowExecutionInitiatedBody,
@@ -125,7 +126,7 @@ class WorkflowRunner(Generic[StateType]):
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
-        self._parent_context = parent_context
+        self._parent_context = get_parent_context() or parent_context
 
         setattr(
             self._initial_state,
@@ -156,6 +157,7 @@ class WorkflowRunner(Generic[StateType]):
         return event
 
     def _run_work_item(self, node: BaseNode[StateType], span_id: UUID) -> None:
+        parent_context = get_parent_context()
         self._workflow_event_inner_queue.put(
             NodeExecutionInitiatedEvent(
                 trace_id=node.state.meta.trace_id,
@@ -164,19 +166,20 @@ class WorkflowRunner(Generic[StateType]):
                     node_definition=node.__class__,
                     inputs=node._inputs,
                 ),
-                parent=WorkflowParentContext(
-                    span_id=span_id,
-                    workflow_definition=self.workflow.__class__,
-                    parent=self._parent_context,
-                    type="WORKFLOW",
-                ),
+                parent=parent_context,
             )
         )
 
         logger.debug(f"Started running node: {node.__class__.__name__}")
 
         try:
-            node_run_response = node.run()
+            updated_parent_context = NodeParentContext(
+                span_id=span_id,
+                node_definition=node.__class__,
+                parent=parent_context,
+            )
+            with execution_context(parent_context=updated_parent_context):
+                node_run_response = node.run()
             ports = node.Ports()
             if not isinstance(node_run_response, (BaseOutputs, Iterator)):
                 raise NodeException(
@@ -197,6 +200,7 @@ class WorkflowRunner(Generic[StateType]):
                 outputs = node.Outputs()
 
                 def initiate_node_streaming_output(output: BaseOutput) -> None:
+                    parent_context = get_parent_context()
                     streaming_output_queues[output.name] = Queue()
                     output_descriptor = OutputReference(
                         name=output.name,
@@ -216,60 +220,49 @@ class WorkflowRunner(Generic[StateType]):
                                 output=initiated_output,
                                 invoked_ports=initiated_ports,
                             ),
-                            parent=WorkflowParentContext(
-                                span_id=span_id,
-                                workflow_definition=self.workflow.__class__,
-                                parent=self._parent_context,
-                            ),
+                            parent=parent_context,
                         ),
                     )
 
-                for output in node_run_response:
-                    invoked_ports = output > ports
-                    if output.is_initiated:
-                        initiate_node_streaming_output(output)
-                    elif output.is_streaming:
-                        if output.name not in streaming_output_queues:
+                with execution_context(parent_context=updated_parent_context):
+                    for output in node_run_response:
+                        invoked_ports = output > ports
+                        if output.is_initiated:
                             initiate_node_streaming_output(output)
-
-                        streaming_output_queues[output.name].put(output.delta)
-                        self._workflow_event_inner_queue.put(
-                            NodeExecutionStreamingEvent(
-                                trace_id=node.state.meta.trace_id,
-                                span_id=span_id,
-                                body=NodeExecutionStreamingBody(
-                                    node_definition=node.__class__,
-                                    output=output,
-                                    invoked_ports=invoked_ports,
-                                ),
-                                parent=WorkflowParentContext(
+                        elif output.is_streaming:
+                            if output.name not in streaming_output_queues:
+                                initiate_node_streaming_output(output)
+
+                            streaming_output_queues[output.name].put(output.delta)
+                            self._workflow_event_inner_queue.put(
+                                NodeExecutionStreamingEvent(
+                                    trace_id=node.state.meta.trace_id,
                                     span_id=span_id,
-                                    workflow_definition=self.workflow.__class__,
-                                    parent=self._parent_context,
+                                    body=NodeExecutionStreamingBody(
+                                        node_definition=node.__class__,
+                                        output=output,
+                                        invoked_ports=invoked_ports,
+                                    ),
+                                    parent=parent_context,
                                 ),
-                            ),
-                        )
-                    elif output.is_fulfilled:
-                        if output.name in streaming_output_queues:
-                            streaming_output_queues[output.name].put(UNDEF)
-
-                        setattr(outputs, output.name, output.value)
-                        self._workflow_event_inner_queue.put(
-                            NodeExecutionStreamingEvent(
-                                trace_id=node.state.meta.trace_id,
-                                span_id=span_id,
-                                body=NodeExecutionStreamingBody(
-                                    node_definition=node.__class__,
-                                    output=output,
-                                    invoked_ports=invoked_ports,
-                                ),
-                                parent=WorkflowParentContext(
+                            )
+                        elif output.is_fulfilled:
+                            if output.name in streaming_output_queues:
+                                streaming_output_queues[output.name].put(UNDEF)
+
+                            setattr(outputs, output.name, output.value)
+                            self._workflow_event_inner_queue.put(
+                                NodeExecutionStreamingEvent(
+                                    trace_id=node.state.meta.trace_id,
                                     span_id=span_id,
-                                    workflow_definition=self.workflow.__class__,
-                                    parent=self._parent_context,
-                                ),
+                                    body=NodeExecutionStreamingBody(
+                                        node_definition=node.__class__,
+                                        output=output,
+                                        invoked_ports=invoked_ports,
+                                    ),
+                                    parent=parent_context,
+                                )
                             )
-                        )
 
             invoked_ports = ports(outputs, node.state)
             node.state.meta.node_execution_cache.fulfill_node_execution(node.__class__, span_id)
@@ -291,11 +284,7 @@ class WorkflowRunner(Generic[StateType]):
                         outputs=outputs,
                         invoked_ports=invoked_ports,
                     ),
-                    parent=WorkflowParentContext(
-                        span_id=span_id,
-                        workflow_definition=self.workflow.__class__,
-                        parent=self._parent_context,
-                    ),
+                    parent=parent_context,
                 )
             )
         except NodeException as e:
@@ -328,16 +317,19 @@ class WorkflowRunner(Generic[StateType]):
                             code=VellumErrorCode.INTERNAL_ERROR,
                         ),
                     ),
-                    parent=WorkflowParentContext(
-                        span_id=span_id,
-                        workflow_definition=self.workflow.__class__,
-                        parent=self._parent_context,
-                    ),
+                    parent=parent_context,
                 ),
             )
 
         logger.debug(f"Finished running node: {node.__class__.__name__}")
 
+    def _context_run_work_item(self, node: BaseNode[StateType], span_id: UUID, parent_context=None) -> None:
+        if parent_context is None:
+            parent_context = get_parent_context() or self._parent_context
+
+        with execution_context(parent_context=parent_context):
+            self._run_work_item(node, span_id)
+
     def _handle_invoked_ports(self, state: StateType, ports: Optional[Iterable[Port]]) -> None:
         if not ports:
             return
@@ -372,13 +364,14 @@ class WorkflowRunner(Generic[StateType]):
             if not node_class.Trigger.should_initiate(state, all_deps, node_span_id):
                 return
 
+            current_parent = get_parent_context()
             node = node_class(state=state, context=self.workflow.context)
             state.meta.node_execution_cache.initiate_node_execution(node_class, node_span_id)
             self._active_nodes_by_execution_id[node_span_id] = node
 
             worker_thread = Thread(
-                target=self._run_work_item,
-                kwargs={"node": node, "span_id": node_span_id},
+                target=self._context_run_work_item,
+                kwargs={"node": node, "span_id": node_span_id, "parent_context": current_parent},
             )
             worker_thread.start()
 
@@ -504,9 +497,16 @@ class WorkflowRunner(Generic[StateType]):
         for edge in self.workflow.get_edges():
             self._dependencies[edge.to_node].add(edge.from_port.node_class)
 
+        current_parent = WorkflowParentContext(
+            span_id=self._initial_state.meta.span_id,
+            workflow_definition=self.workflow.__class__,
+            parent=self._parent_context,
+            type="WORKFLOW",
+        )
         for node_cls in self._entrypoints:
             try:
-                self._run_node_if_ready(self._initial_state, node_cls)
+                with execution_context(parent_context=current_parent):
+                    self._run_node_if_ready(self._initial_state, node_cls)
             except NodeException as e:
                 self._workflow_event_outer_queue.put(self._reject_workflow_event(e.error))
                 return
@@ -530,7 +530,9 @@ class WorkflowRunner(Generic[StateType]):
 
             self._workflow_event_outer_queue.put(event)
 
-            rejection_error = self._handle_work_item_event(event)
+            with execution_context(parent_context=current_parent):
+                rejection_error = self._handle_work_item_event(event)
+
             if rejection_error:
                 break
 
@@ -539,7 +541,9 @@ class WorkflowRunner(Generic[StateType]):
             while event := self._workflow_event_inner_queue.get_nowait():
                 self._workflow_event_outer_queue.put(event)
 
-                rejection_error = self._handle_work_item_event(event)
+                with execution_context(parent_context=current_parent):
+                    rejection_error = self._handle_work_item_event(event)
+
                 if rejection_error:
                     break
         except Empty:

vellum/workflows/workflows/event_filters.py

@@ -46,7 +46,10 @@ def root_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event:
     if event.parent.type != "WORKFLOW":
         return False
 
-    return event.parent.workflow_definition == CodeResourceDefinition.encode(workflow_definition)
+    event_parent_definition = event.parent.workflow_definition
+    current_workflow_definition = CodeResourceDefinition.encode(workflow_definition)
+
+    return event_parent_definition.model_dump() == current_workflow_definition.model_dump()
 
 
 def all_workflow_event_filter(workflow_definition: Type["BaseWorkflow"], event: "WorkflowEvent") -> bool:

{vellum_ai-0.11.0.dist-info → vellum_ai-0.11.3.dist-info}/METADATA

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