vellum-ai 0.13.21__py3-none-any.whl → 0.13.23__py3-none-any.whl

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

@@ -1,6 +1,6 @@
 import json
 from uuid import uuid4
-from typing import Callable, ClassVar, Generic, Iterator, List, Optional, Tuple, Union, cast
+from typing import Callable, ClassVar, Generic, Iterator, List, Optional, Tuple, Union
 
 from vellum import (
     AdHocExecutePromptEvent,
@@ -16,6 +16,7 @@ from vellum import (
     VellumVariable,
 )
 from vellum.client import RequestOptions
+from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.workflows.constants import OMIT
 from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import WorkflowErrorCode
@@ -108,7 +109,14 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                         value=input_value,
                     )
                 )
-            elif isinstance(input_value, list) and all(isinstance(message, ChatMessage) for message in input_value):
+            elif isinstance(input_value, list) and all(
+                isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value
+            ):
+                chat_history = [
+                    message if isinstance(message, ChatMessage) else ChatMessage.model_validate(message.model_dump())
+                    for message in input_value
+                    if isinstance(message, (ChatMessage, ChatMessageRequest))
+                ]
                 input_variables.append(
                     VellumVariable(
                         # TODO: Determine whether or not we actually need an id here and if we do,
@@ -122,7 +130,7 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                 input_values.append(
                     PromptRequestChatHistoryInput(
                         key=input_name,
-                        value=cast(List[ChatMessage], input_value),
+                        value=chat_history,
                     )
                 )
             else:

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

@@ -1,5 +1,6 @@
+import json
 from uuid import UUID
-from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Sequence, Union, cast
+from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Sequence, Union
 
 from vellum import (
     ChatHistoryInputRequest,
@@ -12,9 +13,11 @@ from vellum import (
     StringInputRequest,
 )
 from vellum.client import RequestOptions
+from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
 from vellum.workflows.context import get_parent_context
 from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.events.types import default_serializer
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePromptNode
 from vellum.workflows.types import MergeBehavior
@@ -89,26 +92,38 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                         value=input_value,
                     )
                 )
-            elif isinstance(input_value, list) and all(isinstance(message, ChatMessage) for message in input_value):
+            elif isinstance(input_value, list) and all(
+                isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value
+            ):
+                chat_history = [
+                    (
+                        message
+                        if isinstance(message, ChatMessageRequest)
+                        else ChatMessageRequest.model_validate(message.model_dump())
+                    )
+                    for message in input_value
+                    if isinstance(message, (ChatMessage, ChatMessageRequest))
+                ]
                 compiled_inputs.append(
                     ChatHistoryInputRequest(
                         name=input_name,
-                        value=cast(List[ChatMessage], input_value),
+                        value=chat_history,
                     )
                 )
-            elif isinstance(input_value, dict):
-                # Note: We may want to fail early here if we know that input_value is not
-                #   JSON serializable.
+            else:
+                try:
+                    input_value = default_serializer(input_value)
+                except json.JSONDecodeError as e:
+                    raise NodeException(
+                        message=f"Failed to serialize input '{input_name}' of type '{input_value.__class__}': {e}",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
+
                 compiled_inputs.append(
                     JsonInputRequest(
                         name=input_name,
-                        value=cast(Dict[str, Any], input_value),
+                        value=input_value,
                     )
                 )
-            else:
-                raise NodeException(
-                    message=f"Unrecognized input type for input '{input_name}': {input_value.__class__}",
-                    code=WorkflowErrorCode.INVALID_INPUTS,
-                )
 
         return compiled_inputs

vellum/workflows/nodes/displayable/code_execution_node/tests/test_code_execution_node.py

@@ -1,5 +1,6 @@
 import pytest
 import os
+from typing import Any
 
 from vellum import CodeExecutorResponse, NumberVellumValue, StringInput
 from vellum.client.types.code_execution_package import CodeExecutionPackage
@@ -413,3 +414,35 @@ name
   Field required [type=missing, input_value={'n': 'hello', 'a': {}}, input_type=dict]\
 """
     )
+
+
+def test_run_workflow__run_inline__valid_dict_to_pydantic_any_type():
+    """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a dict matching Any
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
+        code = """\
+def main(word: str) -> dict:
+    return {
+        "name": "word",
+        "arguments": {},
+    }
+    """
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "word": "hello",
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode()
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {
+        "result": {
+            "name": "word",
+            "arguments": {},
+        },
+        "log": "",
+    }

vellum/workflows/nodes/displayable/code_execution_node/utils.py

@@ -91,19 +91,21 @@ __arg__out = main({", ".join(run_args)})
     logs = log_buffer.getvalue()
     result = exec_globals["__arg__out"]
 
-    if issubclass(output_type, BaseModel) and not isinstance(result, output_type):
-        try:
-            result = output_type.model_validate(result)
-        except ValidationError as e:
+    if output_type != Any:
+        if issubclass(output_type, BaseModel) and not isinstance(result, output_type):
+            try:
+                result = output_type.model_validate(result)
+            except ValidationError as e:
+                raise NodeException(
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    message=re.sub(r"\s+For further information visit [^\s]+", "", str(e)),
+                ) from e
+
+        if not isinstance(result, output_type):
             raise NodeException(
                 code=WorkflowErrorCode.INVALID_OUTPUTS,
-                message=re.sub(r"\s+For further information visit [^\s]+", "", str(e)),
-            ) from e
-
-    if not isinstance(result, output_type):
-        raise NodeException(
-            code=WorkflowErrorCode.INVALID_OUTPUTS,
-            message=f"Expected an output of type '{output_type.__name__}', but received '{result.__class__.__name__}'",
-        )
+                message=f"Expected an output of type '{output_type.__name__}',"
+                f" but received '{result.__class__.__name__}'",
+            )
 
     return logs, result

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

@@ -5,6 +5,8 @@ from typing import Any, Iterator, List
 
 from vellum.client.core.api_error import ApiError
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.client.types.execute_prompt_event import ExecutePromptEvent
 from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
 from vellum.client.types.function_call import FunctionCall
@@ -12,6 +14,7 @@ from vellum.client.types.function_call_vellum_value import FunctionCallVellumVal
 from vellum.client.types.function_definition import FunctionDefinition
 from vellum.client.types.initiated_execute_prompt_event import InitiatedExecutePromptEvent
 from vellum.client.types.prompt_output import PromptOutput
+from vellum.client.types.prompt_request_chat_history_input import PromptRequestChatHistoryInput
 from vellum.client.types.prompt_request_json_input import PromptRequestJsonInput
 from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.workflows.errors.types import WorkflowErrorCode
@@ -181,3 +184,49 @@ def test_inline_prompt_node__api_error__invalid_inputs_node_exception(
     # THEN the node raises the correct NodeException
     assert e.value.code == expected_code
     assert e.value.message == expected_message
+
+
+def test_inline_prompt_node__chat_history_inputs(vellum_adhoc_prompt_client):
+    # GIVEN a prompt node with a chat history input
+    class MyNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+        prompt_inputs = {
+            "chat_history": [ChatMessageRequest(role="USER", text="Hello, how are you?")],
+        }
+
+    # AND a known response from invoking an inline prompt
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value="Great!"),
+    ]
+
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=expected_outputs,
+            ),
+        ]
+        yield from events
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN the node is run
+    events = list(MyNode().run())
+
+    # THEN the prompt is executed with the correct inputs
+    assert events[-1].value == "Great!"
+
+    # AND the prompt is executed with the correct inputs
+    mock_api = vellum_adhoc_prompt_client.adhoc_execute_prompt_stream
+    assert mock_api.call_count == 1
+    assert mock_api.call_args.kwargs["input_values"] == [
+        PromptRequestChatHistoryInput(
+            key="chat_history",
+            type="CHAT_HISTORY",
+            value=[ChatMessage(role="USER", text="Hello, how are you?")],
+        ),
+    ]
+    assert mock_api.call_args.kwargs["input_variables"][0].type == "CHAT_HISTORY"

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

@@ -0,0 +1,96 @@
+import pytest
+from uuid import uuid4
+from typing import Any, Iterator, List
+
+from vellum.client.types.chat_history_input_request import ChatHistoryInputRequest
+from vellum.client.types.chat_message import ChatMessage
+from vellum.client.types.chat_message_request import ChatMessageRequest
+from vellum.client.types.execute_prompt_event import ExecutePromptEvent
+from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
+from vellum.client.types.initiated_execute_prompt_event import InitiatedExecutePromptEvent
+from vellum.client.types.json_input_request import JsonInputRequest
+from vellum.client.types.string_vellum_value import StringVellumValue
+from vellum.workflows.nodes.displayable.prompt_deployment_node.node import PromptDeploymentNode
+
+
+@pytest.mark.parametrize("ChatMessageClass", [ChatMessageRequest, ChatMessage])
+def test_run_node__chat_history_input(vellum_client, ChatMessageClass):
+    """Confirm that we can successfully invoke a Prompt Deployment Node that uses Chat History Inputs"""
+
+    # GIVEN a Prompt Deployment Node
+    class ExamplePromptDeploymentNode(PromptDeploymentNode):
+        deployment = "example_prompt_deployment"
+        prompt_inputs = {
+            "chat_history": [ChatMessageClass(role="USER", text="Hello, how are you?")],
+        }
+
+    # AND we know what the Prompt Deployment will respond with
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=[
+                    StringVellumValue(value="Great!"),
+                ],
+            ),
+        ]
+        yield from events
+
+    vellum_client.execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN we run the node
+    node = ExamplePromptDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Prompt Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_prompt_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        ChatHistoryInputRequest(
+            name="chat_history", value=[ChatMessageRequest(role="USER", text="Hello, how are you?")]
+        ),
+    ]
+
+
+def test_run_node__any_array_input(vellum_client):
+    """Confirm that we can successfully invoke a Prompt Deployment Node that uses any array input"""
+
+    # GIVEN a Prompt Deployment Node
+    class ExamplePromptDeploymentNode(PromptDeploymentNode):
+        deployment = "example_prompt_deployment"
+        prompt_inputs = {
+            "fruits": ["apple", "banana", "cherry"],
+        }
+
+    # AND we know what the Prompt Deployment will respond with
+    def generate_prompt_events(*args: Any, **kwargs: Any) -> Iterator[ExecutePromptEvent]:
+        execution_id = str(uuid4())
+        events: List[ExecutePromptEvent] = [
+            InitiatedExecutePromptEvent(execution_id=execution_id),
+            FulfilledExecutePromptEvent(
+                execution_id=execution_id,
+                outputs=[
+                    StringVellumValue(value="Great!"),
+                ],
+            ),
+        ]
+        yield from events
+
+    vellum_client.execute_prompt_stream.side_effect = generate_prompt_events
+
+    # WHEN we run the node
+    node = ExamplePromptDeploymentNode()
+    events = list(node.run())
+
+    # THEN the node should have completed successfully
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Prompt Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_prompt_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        JsonInputRequest(name="fruits", value=["apple", "banana", "cherry"]),
+    ]

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

@@ -1,3 +1,4 @@
+import json
 from uuid import UUID
 from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Union, cast
 
@@ -11,11 +12,13 @@ from vellum import (
     WorkflowRequestNumberInputRequest,
     WorkflowRequestStringInputRequest,
 )
+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.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
+from vellum.workflows.events.types import default_serializer
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.outputs.base import BaseOutput
@@ -66,11 +69,22 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                         value=input_value,
                     )
                 )
-            elif isinstance(input_value, list) and all(isinstance(message, ChatMessage) for message in input_value):
+            elif isinstance(input_value, list) and all(
+                isinstance(message, (ChatMessage, ChatMessageRequest)) for message in input_value
+            ):
+                chat_history = [
+                    (
+                        message
+                        if isinstance(message, ChatMessageRequest)
+                        else ChatMessageRequest.model_validate(message.model_dump())
+                    )
+                    for message in input_value
+                    if isinstance(message, (ChatMessage, ChatMessageRequest))
+                ]
                 compiled_inputs.append(
                     WorkflowRequestChatHistoryInputRequest(
                         name=input_name,
-                        value=cast(List[ChatMessage], input_value),
+                        value=chat_history,
                     )
                 )
             elif isinstance(input_value, dict):
@@ -88,9 +102,19 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
                     )
                 )
             else:
-                raise NodeException(
-                    message=f"Unrecognized input type for input '{input_name}'",
-                    code=WorkflowErrorCode.INVALID_INPUTS,
+                try:
+                    input_value = default_serializer(input_value)
+                except json.JSONDecodeError as e:
+                    raise NodeException(
+                        message=f"Failed to serialize input '{input_name}' of type '{input_value.__class__}': {e}",
+                        code=WorkflowErrorCode.INVALID_INPUTS,
+                    )
+
+                compiled_inputs.append(
+                    WorkflowRequestJsonInputRequest(
+                        name=input_name,
+                        value=input_value,
+                    )
                 )
 
         return compiled_inputs

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

@@ -0,0 +1,131 @@
+import pytest
+from datetime import datetime
+from uuid import uuid4
+from typing import Any, Iterator, List
+
+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
+from vellum.client.types.workflow_output_string import WorkflowOutputString
+from vellum.client.types.workflow_request_chat_history_input_request import WorkflowRequestChatHistoryInputRequest
+from vellum.client.types.workflow_request_json_input_request import WorkflowRequestJsonInputRequest
+from vellum.client.types.workflow_result_event import WorkflowResultEvent
+from vellum.client.types.workflow_stream_event import WorkflowStreamEvent
+from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
+
+
+@pytest.mark.parametrize("ChatMessageClass", [ChatMessageRequest, ChatMessage])
+def test_run_workflow__chat_history_input(vellum_client, ChatMessageClass):
+    """Confirm that we can successfully invoke a Subworkflow Deployment Node that uses Chat History Inputs"""
+
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "chat_history": [ChatMessageClass(role="USER", text="Hello, how are you?")],
+        }
+
+    # 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="greeting",
+                            value="Great!",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        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 == "greeting"
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Subworkflow Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_workflow_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        WorkflowRequestChatHistoryInputRequest(
+            name="chat_history", value=[ChatMessageRequest(role="USER", text="Hello, how are you?")]
+        ),
+    ]
+
+
+def test_run_workflow__any_array(vellum_client):
+    """Confirm that we can successfully invoke a Subworkflow Deployment Node that uses any array input"""
+
+    # GIVEN a Subworkflow Deployment Node
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {
+            "fruits": ["apple", "banana", "cherry"],
+        }
+
+    # 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="greeting",
+                            value="Great!",
+                        )
+                    ],
+                ),
+            ),
+        ]
+        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 == "greeting"
+    assert events[-1].value == "Great!"
+
+    # AND we should have invoked the Subworkflow Deployment with the expected inputs
+    call_kwargs = vellum_client.execute_workflow_stream.call_args.kwargs
+    assert call_kwargs["inputs"] == [
+        WorkflowRequestJsonInputRequest(name="fruits", value=["apple", "banana", "cherry"]),
+    ]

vellum/workflows/nodes/mocks.py

@@ -0,0 +1,17 @@
+from typing import Sequence, Union
+
+from pydantic import ConfigDict
+
+from vellum.client.core.pydantic_utilities import UniversalBaseModel
+from vellum.workflows.descriptors.base import BaseDescriptor
+from vellum.workflows.outputs.base import BaseOutputs
+
+
+class MockNodeExecution(UniversalBaseModel):
+    when_condition: BaseDescriptor
+    then_outputs: BaseOutputs
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+MockNodeExecutionArg = Sequence[Union[BaseOutputs, MockNodeExecution]]

vellum/workflows/runner/runner.py

@@ -4,21 +4,7 @@ import logging
 from queue import Empty, Queue
 from threading import Event as ThreadingEvent, Thread
 from uuid import UUID
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    Generic,
-    Iterable,
-    Iterator,
-    List,
-    Optional,
-    Sequence,
-    Set,
-    Tuple,
-    Type,
-    Union,
-)
+from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Optional, Sequence, Set, Tuple, Type, Union
 
 from vellum.workflows.constants import UNDEF
 from vellum.workflows.context import execution_context, get_parent_context
@@ -59,12 +45,12 @@ from vellum.workflows.events.workflow import (
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import NodeRunResponse
+from vellum.workflows.nodes.mocks import MockNodeExecutionArg
 from vellum.workflows.outputs import BaseOutputs
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
 from vellum.workflows.references import ExternalInputReference, OutputReference
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.types.cycle_map import CycleMap
 from vellum.workflows.types.generics import InputsType, OutputsType, StateType
 
 if TYPE_CHECKING:
@@ -88,7 +74,7 @@ class WorkflowRunner(Generic[StateType]):
         entrypoint_nodes: Optional[RunFromNodeArg] = None,
         external_inputs: Optional[ExternalInputsArg] = None,
         cancel_signal: Optional[ThreadingEvent] = None,
-        node_output_mocks: Optional[List[BaseOutputs]] = None,
+        node_output_mocks: Optional[MockNodeExecutionArg] = None,
         parent_context: Optional[ParentContext] = None,
         max_concurrency: Optional[int] = None,
     ):
@@ -144,7 +130,6 @@ class WorkflowRunner(Generic[StateType]):
 
         self._dependencies: Dict[Type[BaseNode], Set[Type[BaseNode]]] = defaultdict(set)
         self._state_forks: Set[StateType] = {self._initial_state}
-        self._mocks_by_node_outputs_class = CycleMap(items=node_output_mocks or [], key_by=lambda mock: mock.__class__)
 
         self._active_nodes_by_execution_id: Dict[UUID, BaseNode[StateType]] = {}
         self._cancel_signal = cancel_signal
@@ -156,6 +141,7 @@ class WorkflowRunner(Generic[StateType]):
             lambda s: self._snapshot_state(s),
         )
         self.workflow.context._register_event_queue(self._workflow_event_inner_queue)
+        self.workflow.context._register_node_output_mocks(node_output_mocks or [])
 
     def _snapshot_state(self, state: StateType) -> StateType:
         self._workflow_event_inner_queue.put(
@@ -201,11 +187,18 @@ class WorkflowRunner(Generic[StateType]):
                 parent=parent_context,
             )
             node_run_response: NodeRunResponse
-            if node.Outputs not in self._mocks_by_node_outputs_class:
+            was_mocked = False
+            mock_candidates = self.workflow.context.node_output_mocks_map.get(node.Outputs) or []
+            for mock_candidate in mock_candidates:
+                if mock_candidate.when_condition.resolve(node.state):
+                    node_run_response = mock_candidate.then_outputs
+                    was_mocked = True
+                    break
+
+            if not was_mocked:
                 with execution_context(parent_context=updated_parent_context):
                     node_run_response = node.run()
-            else:
-                node_run_response = self._mocks_by_node_outputs_class[node.Outputs]
+
             ports = node.Ports()
             if not isinstance(node_run_response, (BaseOutputs, Iterator)):
                 raise NodeException(
@@ -519,19 +512,6 @@ class WorkflowRunner(Generic[StateType]):
         )
 
     def _stream(self) -> None:
-        # TODO: We should likely handle this during initialization
-        # https://app.shortcut.com/vellum/story/4327
-        if not self._entrypoints:
-            self._workflow_event_outer_queue.put(
-                self._reject_workflow_event(
-                    WorkflowError(
-                        message="No entrypoints defined",
-                        code=WorkflowErrorCode.INVALID_WORKFLOW,
-                    )
-                )
-            )
-            return
-
         for edge in self.workflow.get_edges():
             self._dependencies[edge.to_node].add(edge.from_port.node_class)