vellum-ai 0.14.6__py3-none-any.whl → 0.14.7__py3-none-any.whl

vellum/workflows/nodes/displayable/inline_prompt_node/tests/test_node.py

@@ -146,22 +146,22 @@ def test_inline_prompt_node__function_definitions(vellum_adhoc_prompt_client):
         (
             ApiError(status_code=404, body={"message": "Model not found"}),
             WorkflowErrorCode.INVALID_INPUTS,
-            "Failed to execute prompt",
+            "Failed to execute Prompt",
         ),
         (
             ApiError(status_code=404, body="Model not found"),
             WorkflowErrorCode.INTERNAL_ERROR,
-            "Failed to execute prompt",
+            "Failed to execute Prompt",
         ),
         (
             ApiError(status_code=None, body={"detail": "Model not found"}),
             WorkflowErrorCode.INTERNAL_ERROR,
-            "Failed to execute prompt",
+            "Failed to execute Prompt",
         ),
         (
             ApiError(status_code=500, body={"detail": "Model not found"}),
             WorkflowErrorCode.INTERNAL_ERROR,
-            "Failed to execute prompt",
+            "Failed to execute Prompt",
         ),
     ],
     ids=["404", "invalid_dict", "invalid_body", "no_status_code", "500"],

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

@@ -12,10 +12,11 @@ from vellum import (
     WorkflowRequestNumberInputRequest,
     WorkflowRequestStringInputRequest,
 )
+from vellum.client.core.api_error import ApiError
 from vellum.client.types.chat_message_request import ChatMessageRequest
 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.context import get_execution_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
 from vellum.workflows.events.types import default_serializer
@@ -120,11 +121,13 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
         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
+        current_context = get_execution_context()
+        parent_context = (
+            current_context.parent_context.model_dump(mode="json") if current_context.parent_context else None
+        )
         request_options = self.request_options or RequestOptions()
         request_options["additional_body_parameters"] = {
-            "execution_context": {"parent_context": parent_context},
+            "execution_context": {"parent_context": parent_context, "trace_id": current_context.trace_id},
             **request_options.get("additional_body_parameters", {}),
         }
 
@@ -137,17 +140,26 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                 message="Expected subworkflow deployment attribute to be either a UUID or STR, got None instead",
             )
 
-        subworkflow_stream = self._context.vellum_client.execute_workflow_stream(
-            inputs=self._compile_subworkflow_inputs(),
-            workflow_deployment_id=deployment_id,
-            workflow_deployment_name=deployment_name,
-            release_tag=self.release_tag,
-            external_id=self.external_id,
-            event_types=["WORKFLOW"],
-            metadata=self.metadata,
-            request_options=request_options,
-        )
-        # for some reason execution context isn't showing as an option? ^ failing mypy
+        try:
+            subworkflow_stream = self._context.vellum_client.execute_workflow_stream(
+                inputs=self._compile_subworkflow_inputs(),
+                workflow_deployment_id=deployment_id,
+                workflow_deployment_name=deployment_name,
+                release_tag=self.release_tag,
+                external_id=self.external_id,
+                event_types=["WORKFLOW"],
+                metadata=self.metadata,
+                request_options=request_options,
+            )
+        except ApiError as e:
+            self._handle_api_error(e)
+
+        # We don't use the INITIATED event anyway, so we can just skip it
+        # and use the exception handling to catch other api level errors
+        try:
+            next(subworkflow_stream)
+        except ApiError as e:
+            self._handle_api_error(e)
 
         outputs: Optional[List[WorkflowOutput]] = None
         fulfilled_output_names: Set[str] = set()
@@ -195,3 +207,15 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                     name=output.name,
                     value=output.value,
                 )
+
+    def _handle_api_error(self, e: ApiError):
+        if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
+            raise NodeException(
+                message=e.body.get("detail", "Failed to execute Subworkflow Deployment"),
+                code=WorkflowErrorCode.INVALID_INPUTS,
+            ) from e
+
+        raise NodeException(
+            message="Failed to execute Subworkflow Deployment",
+            code=WorkflowErrorCode.INTERNAL_ERROR,
+        ) from e

vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py

@@ -3,6 +3,7 @@ from datetime import datetime
 from uuid import uuid4
 from typing import Any, Iterator, List
 
+from vellum.client.core.api_error import ApiError
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.client.types.workflow_execution_workflow_result_event import WorkflowExecutionWorkflowResultEvent
@@ -152,3 +153,144 @@ def test_run_workflow__no_deployment():
     assert "Expected subworkflow deployment attribute to be either a UUID or STR, got None instead" in str(
         exc_info.value
     )
+
+
+def test_run_workflow__hyphenated_output(vellum_client):
+    """Confirm that we can successfully handle subworkflow outputs with hyphenated names"""
+
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "test_input": "test_value",
+        }
+
+        class Outputs(SubworkflowDeploymentNode.Outputs):
+            final_output_copy: str
+
+    # AND we know what the Subworkflow Deployment will respond with
+    def generate_subworkflow_events(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        execution_id = str(uuid4())
+        expected_events: List[WorkflowStreamEvent] = [
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            ),
+            WorkflowExecutionWorkflowResultEvent(
+                execution_id=execution_id,
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="FULFILLED",
+                    ts=datetime.now(),
+                    outputs=[
+                        WorkflowOutputString(
+                            id=str(uuid4()),
+                            name="final-output_copy",  # Note the hyphen here
+                            value="test success",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        yield from expected_events
+
+    vellum_client.execute_workflow_stream.side_effect = generate_subworkflow_events
+
+    # WHEN we run the node
+    node = ExampleSubworkflowDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].name == "final_output_copy"  # Note the underscore here
+    assert events[-1].value == "test success"
+
+
+@pytest.mark.parametrize(
+    ["exception", "expected_code", "expected_message"],
+    [
+        (
+            ApiError(status_code=400, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Missing required input variable: 'foo'",
+        ),
+        (
+            ApiError(status_code=400, body={"message": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INVALID_INPUTS,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=400, body="Missing required input variable: 'foo'"),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=None, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+        (
+            ApiError(status_code=500, body={"detail": "Missing required input variable: 'foo'"}),
+            WorkflowErrorCode.INTERNAL_ERROR,
+            "Failed to execute Subworkflow Deployment",
+        ),
+    ],
+    ids=["400", "invalid_dict", "invalid_body", "no_status_code", "500"],
+)
+def test_subworkflow_deployment_node__api_error__invalid_inputs_node_exception(
+    vellum_client, exception, expected_code, expected_message
+):
+    # GIVEN a prompt node with an invalid model name
+    class MyNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "not_foo": "bar",
+        }
+
+    # AND the Subworkflow Deployment API call fails
+    def _side_effect(*args: Any, **kwargs: Any) -> Iterator[WorkflowStreamEvent]:
+        if kwargs.get("_mock_condition_to_induce_an_error"):
+            yield WorkflowExecutionWorkflowResultEvent(
+                execution_id=str(uuid4()),
+                data=WorkflowResultEvent(
+                    id=str(uuid4()),
+                    state="INITIATED",
+                    ts=datetime.now(),
+                ),
+            )
+        else:
+            raise exception
+
+    # AND the vellum client execute workflow stream raises a 4xx error
+    vellum_client.execute_workflow_stream.side_effect = _side_effect
+
+    # WHEN the node is run
+    with pytest.raises(NodeException) as e:
+        list(MyNode().run())
+
+    # THEN the node raises the correct NodeException
+    assert e.value.code == expected_code
+    assert e.value.message == expected_message
+
+
+def test_subworkflow_deployment_node__immediate_api_error__node_exception(vellum_client):
+    # GIVEN a prompt node with an invalid model name
+    class MyNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "not_foo": "bar",
+        }
+
+    # AND the vellum client execute workflow stream raises a 4xx error
+    vellum_client.execute_workflow_stream.side_effect = ApiError(status_code=404, body={"detail": "Not found"})
+
+    # WHEN the node is run
+    with pytest.raises(NodeException) as e:
+        list(MyNode().run())
+
+    # THEN the node raises the correct NodeException
+    assert e.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert e.value.message == "Not found"

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

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

vellum/workflows/outputs/base.py

@@ -32,7 +32,7 @@ class BaseOutput(Generic[_Delta, _Accumulated]):
         if value is not undefined and delta is not undefined:
             raise ValueError("Cannot set both value and delta")
 
-        self._name = name
+        self._name = name.replace("-", "_")  # Convert hyphens to underscores for valid python variable names
         self._value = value
         self._delta = delta
 

vellum/workflows/runner/runner.py

@@ -7,7 +7,7 @@ from uuid import UUID
 from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
 
 from vellum.workflows.constants import undefined
-from vellum.workflows.context import execution_context, get_parent_context
+from vellum.workflows.context import ExecutionContext, execution_context, get_execution_context, get_parent_context
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
@@ -29,7 +29,7 @@ from vellum.workflows.events.node import (
     NodeExecutionRejectedBody,
     NodeExecutionStreamingBody,
 )
-from vellum.workflows.events.types import BaseEvent, NodeParentContext, ParentContext, WorkflowParentContext
+from vellum.workflows.events.types import BaseEvent, NodeParentContext, WorkflowParentContext
 from vellum.workflows.events.workflow import (
     WorkflowExecutionFulfilledBody,
     WorkflowExecutionInitiatedBody,
@@ -75,8 +75,8 @@ class WorkflowRunner(Generic[StateType]):
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
         node_output_mocks: Optional[MockNodeExecutionArg] = None,
-        parent_context: Optional[ParentContext] = None,
         max_concurrency: Optional[int] = None,
+        init_execution_context: Optional[ExecutionContext] = None,
     ):
         if state and external_inputs:
             raise ValueError("Can only run a Workflow providing one of state or external inputs, not both")
@@ -98,6 +98,11 @@ class WorkflowRunner(Generic[StateType]):
         elif external_inputs:
             self._initial_state = self.workflow.get_most_recent_state()
             for descriptor, value in external_inputs.items():
+                if not any(isinstance(value, type_) for type_ in descriptor.types):
+                    raise NodeException(
+                        f"Invalid external input type for {descriptor.name}",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
                 self._initial_state.meta.external_inputs[descriptor] = value
 
             self._entrypoints = [
@@ -133,7 +138,8 @@ class WorkflowRunner(Generic[StateType]):
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
-        self._parent_context = get_parent_context() or parent_context
+        self._execution_context = init_execution_context or get_execution_context()
+        self._parent_context = self._execution_context.parent_context
 
         setattr(
             self._initial_state,
@@ -196,7 +202,7 @@ class WorkflowRunner(Generic[StateType]):
                     break
 
             if not was_mocked:
-                with execution_context(parent_context=updated_parent_context):
+                with execution_context(parent_context=updated_parent_context, trace_id=node.state.meta.trace_id):
                     node_run_response = node.run()
 
             ports = node.Ports()
@@ -243,7 +249,7 @@ class WorkflowRunner(Generic[StateType]):
                         ),
                     )
 
-                with execution_context(parent_context=updated_parent_context):
+                with execution_context(parent_context=updated_parent_context, trace_id=node.state.meta.trace_id):
                     for output in node_run_response:
                         invoked_ports = output > ports
                         if output.is_initiated:
@@ -346,7 +352,7 @@ class WorkflowRunner(Generic[StateType]):
         if parent_context is None:
             parent_context = get_parent_context() or self._parent_context
 
-        with execution_context(parent_context=parent_context):
+        with execution_context(parent_context=parent_context, trace_id=node.state.meta.trace_id):
             self._run_work_item(node, span_id)
 
     def _handle_invoked_ports(self, state: StateType, ports: Optional[Iterable[Port]]) -> None:
@@ -524,7 +530,7 @@ class WorkflowRunner(Generic[StateType]):
         for node_cls in self._entrypoints:
             try:
                 if not self._max_concurrency or len(self._active_nodes_by_execution_id) < self._max_concurrency:
-                    with execution_context(parent_context=current_parent):
+                    with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                         self._run_node_if_ready(self._initial_state, node_cls)
                 else:
                     self._concurrency_queue.put((self._initial_state, node_cls, None))
@@ -551,7 +557,7 @@ class WorkflowRunner(Generic[StateType]):
 
             self._workflow_event_outer_queue.put(event)
 
-            with execution_context(parent_context=current_parent):
+            with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                 rejection_error = self._handle_work_item_event(event)
 
             if rejection_error:
@@ -562,7 +568,7 @@ class WorkflowRunner(Generic[StateType]):
             while event := self._workflow_event_inner_queue.get_nowait():
                 self._workflow_event_outer_queue.put(event)
 
-                with execution_context(parent_context=current_parent):
+                with execution_context(parent_context=current_parent, trace_id=self._initial_state.meta.trace_id):
                     rejection_error = self._handle_work_item_event(event)
 
                 if rejection_error:

vellum/workflows/state/context.py

@@ -3,7 +3,7 @@ from queue import Queue
 from typing import TYPE_CHECKING, Dict, List, Optional, Type
 
 from vellum import Vellum
-from vellum.workflows.events.types import ParentContext
+from vellum.workflows.context import ExecutionContext, get_execution_context
 from vellum.workflows.nodes.mocks import MockNodeExecution, MockNodeExecutionArg
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.constant import ConstantValueReference
@@ -18,12 +18,14 @@ class WorkflowContext:
         self,
         *,
         vellum_client: Optional[Vellum] = None,
-        parent_context: Optional[ParentContext] = None,
+        execution_context: Optional[ExecutionContext] = None,
     ):
         self._vellum_client = vellum_client
-        self._parent_context = parent_context
         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
 
     @cached_property
     def vellum_client(self) -> Vellum:
@@ -33,10 +35,8 @@ class WorkflowContext:
         return create_vellum_client()
 
     @cached_property
-    def parent_context(self) -> Optional[ParentContext]:
-        if self._parent_context:
-            return self._parent_context
-        return None
+    def execution_context(self) -> ExecutionContext:
+        return self._execution_context
 
     @cached_property
     def node_output_mocks_map(self) -> Dict[Type[BaseOutputs], List[MockNodeExecution]]:

vellum/workflows/workflows/base.py

@@ -24,6 +24,7 @@ from typing import (
     get_args,
 )
 
+from vellum.workflows.context import get_execution_context
 from vellum.workflows.edges import Edge
 from vellum.workflows.emitters.base import BaseWorkflowEmitter
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
@@ -171,6 +172,7 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         self.resolvers = resolvers or (self.resolvers if hasattr(self, "resolvers") else [])
         self._context = context or WorkflowContext()
         self._store = Store()
+        self._execution_context = self._context.execution_context
 
         self.validate()
 
@@ -320,8 +322,8 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
-            parent_context=self._context.parent_context,
             max_concurrency=max_concurrency,
+            init_execution_context=self._execution_context,
         ).stream()
         first_event: Optional[Union[WorkflowExecutionInitiatedEvent, WorkflowExecutionResumedEvent]] = None
         last_event = None
@@ -431,8 +433,8 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
             external_inputs=external_inputs,
             cancel_signal=cancel_signal,
             node_output_mocks=node_output_mocks,
-            parent_context=self.context.parent_context,
             max_concurrency=max_concurrency,
+            init_execution_context=self._execution_context,
         ).stream():
             if should_yield(self.__class__, event):
                 yield event
@@ -479,10 +481,19 @@ class BaseWorkflow(Generic[InputsType, StateType], metaclass=_BaseWorkflowMeta):
         return self.get_inputs_class()()
 
     def get_default_state(self, workflow_inputs: Optional[InputsType] = None) -> StateType:
+        execution_context = get_execution_context()
         return self.get_state_class()(
-            meta=StateMeta(
-                parent=self._parent_state,
-                workflow_inputs=workflow_inputs or self.get_default_inputs(),
+            meta=(
+                StateMeta(
+                    parent=self._parent_state,
+                    workflow_inputs=workflow_inputs or self.get_default_inputs(),
+                    trace_id=execution_context.trace_id,
+                )
+                if execution_context and execution_context.trace_id
+                else StateMeta(
+                    parent=self._parent_state,
+                    workflow_inputs=workflow_inputs or self.get_default_inputs(),
+                )
             )
         )