vellum-ai 0.14.41__py3-none-any.whl → 0.14.43__py3-none-any.whl

vellum/client/types/workflow_execution_initiated_body.py

@@ -1,19 +1,14 @@
 # This file was auto-generated by Fern from our API Definition.
 
-from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
-from .node_event_display_context import NodeEventDisplayContext
-from .workflow_event_display_context import WorkflowEventDisplayContext
 from .vellum_code_resource_definition import VellumCodeResourceDefinition
 import typing
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
-from ..core.pydantic_utilities import update_forward_refs
 
 
 class WorkflowExecutionInitiatedBody(UniversalBaseModel):
     workflow_definition: VellumCodeResourceDefinition
-    display_context: typing.Optional[WorkflowEventDisplayContext] = None
     inputs: typing.Dict[str, typing.Optional[typing.Any]]
 
     if IS_PYDANTIC_V2:
@@ -24,7 +19,3 @@ class WorkflowExecutionInitiatedBody(UniversalBaseModel):
             frozen = True
             smart_union = True
             extra = pydantic.Extra.allow
-
-
-update_forward_refs(NodeEventDisplayContext, WorkflowExecutionInitiatedBody=WorkflowExecutionInitiatedBody)
-update_forward_refs(WorkflowEventDisplayContext, WorkflowExecutionInitiatedBody=WorkflowExecutionInitiatedBody)

vellum/client/types/workflow_execution_initiated_event.py

@@ -9,8 +9,6 @@ from .span_link import SpanLink
 from .workflow_deployment_parent_context import WorkflowDeploymentParentContext
 from .workflow_parent_context import WorkflowParentContext
 from .workflow_sandbox_parent_context import WorkflowSandboxParentContext
-from .node_event_display_context import NodeEventDisplayContext
-from .workflow_event_display_context import WorkflowEventDisplayContext
 import typing
 from .parent_context import ParentContext
 from .workflow_execution_initiated_body import WorkflowExecutionInitiatedBody
@@ -49,5 +47,3 @@ update_forward_refs(SpanLink, WorkflowExecutionInitiatedEvent=WorkflowExecutionI
 update_forward_refs(WorkflowDeploymentParentContext, WorkflowExecutionInitiatedEvent=WorkflowExecutionInitiatedEvent)
 update_forward_refs(WorkflowParentContext, WorkflowExecutionInitiatedEvent=WorkflowExecutionInitiatedEvent)
 update_forward_refs(WorkflowSandboxParentContext, WorkflowExecutionInitiatedEvent=WorkflowExecutionInitiatedEvent)
-update_forward_refs(NodeEventDisplayContext, WorkflowExecutionInitiatedEvent=WorkflowExecutionInitiatedEvent)
-update_forward_refs(WorkflowEventDisplayContext, WorkflowExecutionInitiatedEvent=WorkflowExecutionInitiatedEvent)

vellum/client/types/workflow_execution_span.py

@@ -3,12 +3,10 @@
 from __future__ import annotations
 from ..core.pydantic_utilities import UniversalBaseModel
 from .api_request_parent_context import ApiRequestParentContext
-from .node_event_display_context import NodeEventDisplayContext
 from .node_parent_context import NodeParentContext
 from .prompt_deployment_parent_context import PromptDeploymentParentContext
 from .span_link import SpanLink
 from .workflow_deployment_parent_context import WorkflowDeploymentParentContext
-from .workflow_event_display_context import WorkflowEventDisplayContext
 from .workflow_parent_context import WorkflowParentContext
 from .workflow_sandbox_parent_context import WorkflowSandboxParentContext
 import typing
@@ -40,11 +38,9 @@ class WorkflowExecutionSpan(UniversalBaseModel):
 
 
 update_forward_refs(ApiRequestParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
-update_forward_refs(NodeEventDisplayContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(NodeParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(PromptDeploymentParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(SpanLink, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(WorkflowDeploymentParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
-update_forward_refs(WorkflowEventDisplayContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(WorkflowParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)
 update_forward_refs(WorkflowSandboxParentContext, WorkflowExecutionSpan=WorkflowExecutionSpan)

vellum/types/{node_event_display_context.py → execute_api_response_json.py}

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

vellum/types/{workflow_event_display_context.py → prompt_push_response.py}

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

vellum/workflows/inputs/base.py

@@ -1,4 +1,4 @@
-from typing import Any, Iterator, Tuple, Type, Union, get_args, get_origin
+from typing import Any, Dict, Iterator, Set, Tuple, Type, Union, get_args, get_origin
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler
@@ -14,11 +14,28 @@ from vellum.workflows.types.utils import get_class_attr_names, infer_types
 
 @dataclass_transform(kw_only_default=True)
 class _BaseInputsMeta(type):
+    def __new__(cls, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
+        dct["__parent_class__"] = type(None)
+        return super().__new__(cls, name, bases, dct)
+
     def __getattribute__(cls, name: str) -> Any:
-        if not name.startswith("_") and name in cls.__annotations__ and issubclass(cls, BaseInputs):
+        if name.startswith("_") or not issubclass(cls, BaseInputs):
+            return super().__getattribute__(name)
+
+        attr_names = get_class_attr_names(cls)
+        if name in attr_names:
+            # We first try to resolve the instance that this class attribute name is mapped to. If it's not found,
+            # we iterate through its inheritance hierarchy to find the first base class that has this attribute
+            # and use its mapping.
             instance = vars(cls).get(name, undefined)
-            types = infer_types(cls, name)
+            if instance is undefined:
+                for base in cls.__mro__[1:]:
+                    inherited_input_reference = getattr(base, name, undefined)
+                    if isinstance(inherited_input_reference, (ExternalInputReference, WorkflowInputReference)):
+                        instance = inherited_input_reference.instance
+                        break
 
+            types = infer_types(cls, name)
             if getattr(cls, "__descriptor_class__", None) is ExternalInputReference:
                 return ExternalInputReference(name=name, types=types, instance=instance, inputs_class=cls)
             else:
@@ -29,14 +46,20 @@ class _BaseInputsMeta(type):
     def __iter__(cls) -> Iterator[InputReference]:
         # We iterate through the inheritance hierarchy to find all the WorkflowInputReference attached to this
         # Inputs class. __mro__ is the method resolution order, which is the order in which base classes are resolved.
+        yielded_attr_names: Set[str] = set()
+
         for resolved_cls in cls.__mro__:
             attr_names = get_class_attr_names(resolved_cls)
             for attr_name in attr_names:
+                if attr_name in yielded_attr_names:
+                    continue
+
                 attr_value = getattr(resolved_cls, attr_name)
                 if not isinstance(attr_value, (WorkflowInputReference, ExternalInputReference)):
                     continue
 
                 yield attr_value
+                yielded_attr_names.add(attr_name)
 
 
 class BaseInputs(metaclass=_BaseInputsMeta):

vellum/workflows/inputs/tests/test_inputs.py

@@ -47,3 +47,18 @@ def test_base_inputs_with_default():
 
     # THEN it should use the default value
     assert inputs.string_with_default == "default_value"
+
+
+def test_base_inputs__supports_inherited_inputs():
+    # GIVEN an inputs class
+    class TopInputs(BaseInputs):
+        first: str
+
+    # WHEN we inherit from the base inputs class
+    class BottomInputs(TopInputs):
+        second: int
+
+    # THEN both references should be available
+    assert BottomInputs.first.name == "first"
+    assert BottomInputs.second.name == "second"
+    assert len([ref for ref in BottomInputs]) == 2

vellum/workflows/nodes/bases/base.py

@@ -41,9 +41,6 @@ def is_nested_class(nested: Any, parent: Type) -> bool:
 
 class BaseNodeMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
-        # TODO: Inherit the inner Output classes from every base class.
-        #   https://app.shortcut.com/vellum/story/4007/support-auto-inheriting-parent-node-outputs
-
         if "Outputs" in dct:
             outputs_class = dct["Outputs"]
             if not any(issubclass(base, BaseOutputs) for base in outputs_class.__bases__):

vellum/workflows/nodes/bases/base_adornment_node.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Tuple, Type
 
+from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode, BaseNodeMeta
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.references.output import OutputReference
@@ -13,6 +14,14 @@ class _BaseAdornmentNodeMeta(BaseNodeMeta):
     def __new__(cls, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         node_class = super().__new__(cls, name, bases, dct)
 
+        SubworkflowInputs = dct.get("SubworkflowInputs")
+        if (
+            isinstance(SubworkflowInputs, type)
+            and issubclass(SubworkflowInputs, BaseInputs)
+            and SubworkflowInputs.__parent_class__ is type(None)
+        ):
+            SubworkflowInputs.__parent_class__ = node_class
+
         subworkflow_attribute = dct.get("subworkflow")
         if not subworkflow_attribute:
             return node_class

vellum/workflows/nodes/bases/tests/test_base_adornment_node.py

@@ -0,0 +1,31 @@
+from typing import Callable, Type
+
+from vellum.workflows.nodes.bases.base import BaseNode
+from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
+from vellum.workflows.nodes.utils import create_adornment
+
+
+def test_base_adornment_node__output_references_of_same_name():
+    # GIVEN a custom adornment node
+    class CustomAdornmentNode(BaseAdornmentNode):
+        @classmethod
+        def wrap(cls) -> Callable[..., Type["CustomAdornmentNode"]]:
+            return create_adornment(cls)
+
+    # AND two nodes wrapped by the adornment with the same output
+    @CustomAdornmentNode.wrap()
+    class AppleNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            fruit: str
+
+    @CustomAdornmentNode.wrap()
+    class OrangeNode(BaseNode):
+        class Outputs(BaseNode.Outputs):
+            fruit: str
+
+    # WHEN get output references of these outputs
+    apple_output_reference = AppleNode.Outputs.fruit
+    orange_output_reference = OrangeNode.Outputs.fruit
+
+    # THEN the output references should not be equal
+    assert apple_output_reference != orange_output_reference, "Output references should not be equal"

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

@@ -62,7 +62,7 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
 
         item: MapNodeItemType  # type: ignore[valid-type]
         index: int
-        all_items: List[MapNodeItemType]  # type: ignore[valid-type]
+        items: List[MapNodeItemType]  # type: ignore[valid-type]
 
     def run(self) -> Iterator[BaseOutput]:
         mapped_items: Dict[str, List] = defaultdict(list)
@@ -176,8 +176,9 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
             parent_state=self.state,
             context=context,
         )
+        SubworkflowInputsClass = self.subworkflow.get_inputs_class()
         events = subworkflow.stream(
-            inputs=self.SubworkflowInputs(index=index, item=item, all_items=self.items),
+            inputs=SubworkflowInputsClass(index=index, item=item, items=self.items),
             node_output_mocks=self._context._get_all_node_output_mocks(),
             event_filter=all_workflow_event_filter,
         )

vellum/workflows/nodes/core/map_node/tests/test_node.py

@@ -116,3 +116,59 @@ def test_map_node__inner_try():
     # THEN the workflow should succeed
     assert outputs[-1].name == "final_output"
     assert len(outputs[-1].value) == 2
+
+
+def test_map_node__nested_map_node():
+    # GIVEN the inner map node's inputs
+    class VegetableMapNodeInputs(MapNode.SubworkflowInputs):
+        item: str
+
+    # AND the outer map node's inputs
+    class FruitMapNodeInputs(MapNode.SubworkflowInputs):
+        item: str
+
+    # AND a simple node that concats both attributes
+    class SimpleConcatNode(BaseNode):
+        fruit = FruitMapNodeInputs.item
+        vegetable = VegetableMapNodeInputs.item
+
+        class Outputs(BaseNode.Outputs):
+            medley: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(medley=f"{self.fruit} {self.vegetable}")
+
+    # AND a workflow using that node
+    class VegetableMapNodeWorkflow(BaseWorkflow[VegetableMapNodeInputs, BaseState]):
+        graph = SimpleConcatNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_output = SimpleConcatNode.Outputs.medley
+
+    # AND an inner map node referencing that workflow
+    class VegetableMapNode(MapNode):
+        items = ["carrot", "potato"]
+        subworkflow = VegetableMapNodeWorkflow
+
+    # AND an outer subworkflow referencing the inner map node
+    class FruitMapNodeWorkflow(BaseWorkflow[FruitMapNodeInputs, BaseState]):
+        graph = VegetableMapNode
+
+        class Outputs(BaseWorkflow.Outputs):
+            final_output = VegetableMapNode.Outputs.final_output
+
+    # AND an outer map node referencing the outer subworkflow
+    class FruitMapNode(MapNode):
+        items = ["apple", "banana"]
+        subworkflow = FruitMapNodeWorkflow
+
+    # WHEN we run the workflow
+    stream = FruitMapNode().run()
+    outputs = list(stream)
+
+    # THEN the workflow should succeed
+    assert outputs[-1].name == "final_output"
+    assert outputs[-1].value == [
+        ["apple carrot", "apple potato"],
+        ["banana carrot", "banana potato"],
+    ]

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

@@ -47,8 +47,9 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
                     parent_state=self.state,
                     context=WorkflowContext.create_from(self._context),
                 )
+                inputs_class = subworkflow.get_inputs_class()
                 subworkflow_stream = subworkflow.stream(
-                    inputs=self.SubworkflowInputs(attempt_number=attempt_number),
+                    inputs=inputs_class(attempt_number=attempt_number),
                     event_filter=all_workflow_event_filter,
                     node_output_mocks=self._context._get_all_node_output_mocks(),
                 )

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, Set, Tuple, Union
+from typing import Callable, ClassVar, Generator, Generic, Iterator, List, Optional, Set, Tuple, Union
 
 from vellum import (
     AdHocExecutePromptEvent,
@@ -8,6 +8,7 @@ from vellum import (
     ChatMessage,
     FunctionDefinition,
     PromptBlock,
+    PromptOutput,
     PromptParameters,
     PromptRequestChatHistoryInput,
     PromptRequestInput,
@@ -15,17 +16,19 @@ from vellum import (
     PromptRequestStringInput,
     VellumVariable,
 )
-from vellum.client import RequestOptions
+from vellum.client import ApiError, RequestOptions
 from vellum.client.types.chat_message_request import ChatMessageRequest
 from vellum.client.types.prompt_settings import PromptSettings
 from vellum.client.types.rich_text_child_block import RichTextChildBlock
 from vellum.workflows.constants import OMIT
 from vellum.workflows.context import get_execution_context
 from vellum.workflows.errors import WorkflowErrorCode
+from vellum.workflows.errors.types import vellum_error_to_workflow_error
 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.nodes.displayable.bases.inline_prompt_node.constants import DEFAULT_PROMPT_PARAMETERS
+from vellum.workflows.outputs import BaseOutput
 from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 from vellum.workflows.utils.functions import compile_function_definition
@@ -103,17 +106,63 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             else None
         )
 
-        return self._context.vellum_client.ad_hoc.adhoc_execute_prompt_stream(
-            ml_model=self.ml_model,
-            input_values=input_values,
-            input_variables=input_variables,
-            parameters=self.parameters,
-            blocks=self.blocks,
-            settings=self.settings,
-            functions=normalized_functions,
-            expand_meta=self.expand_meta,
-            request_options=request_options,
-        )
+        if self.settings and not self.settings.stream_enabled:
+            # This endpoint is returning a single event, so we need to wrap it in a generator
+            # to match the existing interface.
+            response = self._context.vellum_client.ad_hoc.adhoc_execute_prompt(
+                ml_model=self.ml_model,
+                input_values=input_values,
+                input_variables=input_variables,
+                parameters=self.parameters,
+                blocks=self.blocks,
+                settings=self.settings,
+                functions=normalized_functions,
+                expand_meta=self.expand_meta,
+                request_options=request_options,
+            )
+            return iter([response])
+        else:
+            return self._context.vellum_client.ad_hoc.adhoc_execute_prompt_stream(
+                ml_model=self.ml_model,
+                input_values=input_values,
+                input_variables=input_variables,
+                parameters=self.parameters,
+                blocks=self.blocks,
+                settings=self.settings,
+                functions=normalized_functions,
+                expand_meta=self.expand_meta,
+                request_options=request_options,
+            )
+
+    def _process_prompt_event_stream(self) -> Generator[BaseOutput, None, Optional[List[PromptOutput]]]:
+        self._validate()
+        try:
+            prompt_event_stream = self._get_prompt_event_stream()
+        except ApiError as e:
+            self._handle_api_error(e)
+
+        if not self.settings or (self.settings and self.settings.stream_enabled):
+            # 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(prompt_event_stream)
+            except ApiError as e:
+                self._handle_api_error(e)
+
+        outputs: Optional[List[PromptOutput]] = None
+        for event in prompt_event_stream:
+            if event.state == "INITIATED":
+                continue
+            elif event.state == "STREAMING":
+                yield BaseOutput(name="results", delta=event.output.value)
+            elif event.state == "FULFILLED":
+                outputs = event.outputs
+                yield BaseOutput(name="results", value=event.outputs)
+            elif event.state == "REJECTED":
+                workflow_error = vellum_error_to_workflow_error(event.error)
+                raise NodeException.of(workflow_error)
+
+        return outputs
 
     def _compile_prompt_inputs(self) -> Tuple[List[VellumVariable], List[PromptRequestInput]]:
         input_variables: List[VellumVariable] = []

vellum/workflows/nodes/displayable/bases/inline_prompt_node/tests/test_inline_prompt_node.py

@@ -5,11 +5,14 @@ from uuid import uuid4
 from typing import Any, Iterator, List
 
 from vellum import (
+    AdHocExecutePromptEvent,
     ChatMessagePromptBlock,
+    FulfilledAdHocExecutePromptEvent,
     JinjaPromptBlock,
     PlainTextPromptBlock,
     PromptBlock,
     PromptParameters,
+    PromptSettings,
     RichTextPromptBlock,
     VariablePromptBlock,
 )
@@ -296,3 +299,177 @@ def test_inline_prompt_node__json_output(vellum_adhoc_prompt_client):
         request_options=mock.ANY,
         settings=None,
     )
+
+
+def test_inline_prompt_node__streaming_disabled(vellum_adhoc_prompt_client):
+    # GIVEN an InlinePromptNode
+    class Inputs(BaseInputs):
+        input: str
+
+    class State(BaseState):
+        pass
+
+    # AND it has streaming disabled
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+        parameters = PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={},
+        )
+        settings = PromptSettings(stream_enabled=False)
+
+    # AND a known response from invoking an inline prompt
+    expected_output: list[PromptOutput] = [StringVellumValue(value="Hello, world!")]
+
+    def generate_prompt_event(*args: Any, **kwargs: Any) -> AdHocExecutePromptEvent:
+        execution_id = str(uuid4())
+        return FulfilledAdHocExecutePromptEvent(
+            execution_id=execution_id,
+            outputs=expected_output,
+        )
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.side_effect = generate_prompt_event
+
+    # WHEN the node is run
+    node = MyInlinePromptNode()
+    outputs = [o for o in node.run()]
+
+    # THEN the node should have produced the outputs we expect
+    result_output = outputs[0]
+    assert result_output.name == "results"
+    assert result_output.value == expected_output
+
+    # AND we should have made the expected call to Vellum search
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.assert_called_once_with(
+        blocks=[],
+        expand_meta=Ellipsis,
+        functions=None,
+        input_values=[],
+        input_variables=[],
+        ml_model="gpt-4o",
+        parameters=PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={},
+        ),
+        request_options=mock.ANY,
+        settings=PromptSettings(stream_enabled=False),
+    )
+
+
+def test_inline_prompt_node__json_output_with_streaming_disabled(vellum_adhoc_prompt_client):
+    # GIVEN an InlinePromptNode
+    class Inputs(BaseInputs):
+        input: str
+
+    class State(BaseState):
+        pass
+
+    class MyInlinePromptNode(InlinePromptNode):
+        ml_model = "gpt-4o"
+        blocks = []
+        parameters = PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        )
+        settings = PromptSettings(stream_enabled=False)
+
+    # AND a known JSON response from invoking an inline prompt
+    expected_json = {"result": "Hello, world!"}
+    expected_outputs: List[PromptOutput] = [
+        StringVellumValue(value=json.dumps(expected_json)),
+    ]
+
+    def generate_prompt_event(*args: Any, **kwargs: Any) -> AdHocExecutePromptEvent:
+        execution_id = str(uuid4())
+        return FulfilledAdHocExecutePromptEvent(
+            execution_id=execution_id,
+            outputs=expected_outputs,
+        )
+
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.side_effect = generate_prompt_event
+
+    # WHEN the node is run
+    node = MyInlinePromptNode(
+        state=State(
+            meta=StateMeta(workflow_inputs=Inputs(input="Generate JSON.")),
+        )
+    )
+    outputs = [o for o in node.run()]
+
+    # THEN the node should have produced the outputs we expect
+    results_output = outputs[0]
+    assert results_output.name == "results"
+    assert results_output.value == expected_outputs
+
+    text_output = outputs[1]
+    assert text_output.name == "text"
+    assert text_output.value == '{"result": "Hello, world!"}'
+
+    json_output = outputs[2]
+    assert json_output.name == "json"
+    assert json_output.value == expected_json
+
+    # AND we should have made the expected call to Vellum search
+    vellum_adhoc_prompt_client.adhoc_execute_prompt.assert_called_once_with(
+        blocks=[],
+        expand_meta=Ellipsis,
+        functions=None,
+        input_values=[],
+        input_variables=[],
+        ml_model="gpt-4o",
+        parameters=PromptParameters(
+            stop=[],
+            temperature=0.0,
+            max_tokens=4096,
+            top_p=1.0,
+            top_k=0,
+            frequency_penalty=0.0,
+            presence_penalty=0.0,
+            logit_bias=None,
+            custom_parameters={
+                "json_mode": False,
+                "json_schema": {
+                    "name": "get_result",
+                    "schema": {
+                        "type": "object",
+                        "required": ["result"],
+                        "properties": {"result": {"type": "string", "description": ""}},
+                    },
+                },
+            },
+        ),
+        request_options=mock.ANY,
+        settings=PromptSettings(stream_enabled=False),
+    )

vellum/workflows/nodes/experimental/tool_calling_node/node.py

@@ -8,11 +8,7 @@ from vellum.workflows.exceptions import NodeException
 from vellum.workflows.graph.graph import Graph
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.experimental.tool_calling_node.utils import (
-    ToolRouterNode,
-    create_function_node,
-    create_tool_router_node,
-)
+from vellum.workflows.nodes.experimental.tool_calling_node.utils import create_function_node, create_tool_router_node
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.state.context import WorkflowContext
@@ -69,7 +65,7 @@ class ToolCallingNode(BaseNode):
                 graph = self._graph
 
                 class Outputs(BaseWorkflow.Outputs):
-                    text: str = ToolRouterNode.Outputs.text
+                    text: str = self.tool_router_node.Outputs.text
                     chat_history: List[ChatMessage] = ToolCallingState.chat_history
 
             subworkflow = ToolCallingWorkflow(
@@ -107,6 +103,7 @@ class ToolCallingNode(BaseNode):
         self._function_nodes = {
             function.__name__: create_function_node(
                 function=function,
+                tool_router_node=self.tool_router_node,
             )
             for function in self.functions
         }