vellum-ai 0.14.39__py3-none-any.whl → 0.14.40__py3-none-any.whl

vellum/client/core/client_wrapper.py

@@ -18,7 +18,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "0.14.39",
+            "X-Fern-SDK-Version": "0.14.40",
         }
         headers["X_API_KEY"] = self.api_key
         return headers

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

@@ -4,11 +4,13 @@ from typing import Optional
 
 from vellum.client.types.string_vellum_value_request import StringVellumValueRequest
 from vellum.core.pydantic_utilities import UniversalBaseModel
+from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.tests.test_utils import FixtureState
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes import FinalOutputNode
 from vellum.workflows.nodes.bases.base import BaseNode
 from vellum.workflows.outputs.base import BaseOutputs
+from vellum.workflows.references.output import OutputReference
 from vellum.workflows.state.base import BaseState, StateMeta
 
 
@@ -259,3 +261,25 @@ def test_resolve_value__for_falsy_values(falsy_value, expected_type):
 
     # THEN the output has the correct value
     assert falsy_output.value == falsy_value
+
+
+def test_node_outputs__inherits_instance():
+    # GIVEN a node with two outputs, one with and one without a default instance
+    class MyNode(BaseNode):
+        class Outputs:
+            foo: str
+            bar = "hello"
+
+    # AND a node that inherits from MyNode
+    class InheritedNode(MyNode):
+        pass
+
+    # WHEN we reference each output
+    foo_output = InheritedNode.Outputs.foo
+    bar_output = InheritedNode.Outputs.bar
+
+    # THEN the output reference instances are correct
+    assert isinstance(foo_output, OutputReference)
+    assert foo_output.instance is undefined
+    assert isinstance(bar_output, OutputReference)
+    assert bar_output.instance == "hello"

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

@@ -4,7 +4,6 @@ from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.errors.types import WorkflowError, WorkflowErrorCode
 from vellum.workflows.events.workflow import is_workflow_event
 from vellum.workflows.exceptions import NodeException
-from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base_adornment_node import BaseAdornmentNode
 from vellum.workflows.nodes.utils import create_adornment
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
@@ -24,7 +23,7 @@ class TryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
     on_error_code: Optional[WorkflowErrorCode] = None
 
-    class Outputs(BaseNode.Outputs):
+    class Outputs(BaseAdornmentNode.Outputs):
         error: Optional[WorkflowError] = None
 
     def run(self) -> Iterator[BaseOutput]:

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

@@ -0,0 +1,3 @@
+from vellum.workflows.nodes.experimental.tool_calling_node.node import ToolCallingNode
+
+__all__ = ["ToolCallingNode"]

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

@@ -0,0 +1,147 @@
+from collections.abc import Callable
+from typing import Any, ClassVar, Dict, List, Optional, cast
+
+from vellum import ChatMessage, FunctionDefinition, PromptBlock
+from vellum.client.types.chat_message_request import ChatMessageRequest
+from vellum.workflows.context import execution_context, get_parent_context
+from vellum.workflows.errors.types import WorkflowErrorCode
+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.outputs.base import BaseOutputs
+from vellum.workflows.state.base import BaseState
+from vellum.workflows.state.context import WorkflowContext
+from vellum.workflows.types.core import EntityInputsInterface
+from vellum.workflows.workflows.base import BaseWorkflow
+
+
+class ToolCallingNode(BaseNode):
+    """
+    A Node that dynamically invokes the provided functions to the underlying Prompt
+
+    Attributes:
+        ml_model: str - The model to use for tool calling (e.g., "gpt-4o-mini")
+        blocks: List[PromptBlock] - The prompt blocks to use (same format as InlinePromptNode)
+        functions: List[FunctionDefinition] - The functions that can be called
+        function_callables: List[Callable] - The callables that can be called
+        prompt_inputs: Optional[EntityInputsInterface] - Mapping of input variable names to values
+    """
+
+    ml_model: ClassVar[str] = "gpt-4o-mini"
+    blocks: ClassVar[List[PromptBlock]] = []
+    functions: ClassVar[List[FunctionDefinition]] = []
+    function_callables: ClassVar[Dict[str, Callable[..., Any]]] = {}
+    prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
+    # TODO: https://linear.app/vellum/issue/APO-342/support-tool-call-max-retries
+    max_tool_calls: ClassVar[int] = 1
+
+    class Outputs(BaseOutputs):
+        """
+        The outputs of the ToolCallingNode.
+
+        text: The final text response after tool calling
+        chat_history: The complete chat history including tool calls
+        """
+
+        text: str = ""
+        chat_history: List[ChatMessage] = []
+
+    def run(self) -> Outputs:
+        """
+        Run the tool calling workflow.
+
+        This dynamically builds a graph with router and function nodes,
+        then executes the workflow.
+        """
+        self._validate_functions()
+
+        initial_chat_history = []
+
+        # Extract chat history from prompt inputs if available
+        if self.prompt_inputs and "chat_history" in self.prompt_inputs:
+            chat_history_input = self.prompt_inputs["chat_history"]
+            if isinstance(chat_history_input, list) and all(
+                isinstance(msg, (ChatMessage, ChatMessageRequest)) for msg in chat_history_input
+            ):
+                initial_chat_history = [
+                    msg if isinstance(msg, ChatMessage) else ChatMessage.model_validate(msg.model_dump())
+                    for msg in chat_history_input
+                ]
+
+        self._build_graph()
+
+        with execution_context(parent_context=get_parent_context()):
+
+            class ToolCallingState(BaseState):
+                chat_history: List[ChatMessage] = initial_chat_history
+
+            class ToolCallingWorkflow(BaseWorkflow[BaseInputs, ToolCallingState]):
+                graph = self._graph
+
+                class Outputs(BaseWorkflow.Outputs):
+                    text: str = ToolRouterNode.Outputs.text
+                    chat_history: List[ChatMessage] = ToolCallingState.chat_history
+
+            subworkflow = ToolCallingWorkflow(
+                parent_state=self.state,
+                context=WorkflowContext.create_from(self._context),
+            )
+
+            terminal_event = subworkflow.run()
+
+            if terminal_event.name == "workflow.execution.paused":
+                raise NodeException(
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    message="Subworkflow unexpectedly paused",
+                )
+            elif terminal_event.name == "workflow.execution.fulfilled":
+                node_outputs = self.Outputs()
+
+                for output_descriptor, output_value in terminal_event.outputs:
+                    setattr(node_outputs, output_descriptor.name, output_value)
+
+                return node_outputs
+            elif terminal_event.name == "workflow.execution.rejected":
+                raise Exception(f"Workflow execution rejected: {terminal_event.error}")
+
+            raise Exception(f"Unexpected workflow event: {terminal_event.name}")
+
+    def _build_graph(self) -> None:
+        self.tool_router_node = create_tool_router_node(
+            ml_model=self.ml_model,
+            blocks=self.blocks,
+            functions=self.functions,
+            prompt_inputs=self.prompt_inputs,
+        )
+
+        self._function_nodes = {
+            function.name: create_function_node(
+                function=function,
+                function_callable=cast(Callable[..., Any], self.function_callables[function.name]),  # type: ignore
+            )
+            for function in self.functions
+        }
+
+        graph_set = set()
+
+        # Add connections from ports of router to function nodes and back to router
+        for function_name, FunctionNodeClass in self._function_nodes.items():
+            router_port = getattr(self.tool_router_node.Ports, function_name)  # type: ignore  # mypy thinks name is still optional # noqa: E501
+            edge_graph = router_port >> FunctionNodeClass >> self.tool_router_node
+            graph_set.add(edge_graph)
+
+        default_port = getattr(self.tool_router_node.Ports, "default")
+        graph_set.add(default_port)
+
+        self._graph = Graph.from_set(graph_set)
+
+    def _validate_functions(self) -> None:
+        for function in self.functions:
+            if function.name is None:
+                raise ValueError("Function name is required")

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

@@ -0,0 +1,132 @@
+from collections.abc import Callable
+import json
+from typing import Any, Iterator, List, Optional, Type, cast
+
+from vellum import ChatMessage, FunctionDefinition, PromptBlock
+from vellum.client.types.function_call_chat_message_content import FunctionCallChatMessageContent
+from vellum.client.types.function_call_chat_message_content_value import FunctionCallChatMessageContentValue
+from vellum.client.types.variable_prompt_block import VariablePromptBlock
+from vellum.workflows.nodes.bases import BaseNode
+from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
+from vellum.workflows.outputs.base import BaseOutput
+from vellum.workflows.ports.port import Port
+from vellum.workflows.references.lazy import LazyReference
+from vellum.workflows.types.core import EntityInputsInterface
+
+
+class FunctionNode(BaseNode):
+    """Node that executes a specific function."""
+
+    function: FunctionDefinition
+
+
+class ToolRouterNode(InlinePromptNode):
+    def run(self) -> Iterator[BaseOutput]:
+        self.prompt_inputs = {**self.prompt_inputs, "chat_history": self.state.chat_history}  # type: ignore
+        generator = super().run()
+        for output in generator:
+            if output.name == "results" and output.value:
+                values = cast(List[Any], output.value)
+                if values and len(values) > 0:
+                    if values[0].type == "STRING":
+                        self.state.chat_history.append(ChatMessage(role="ASSISTANT", text=values[0].value))
+                    elif values[0].type == "FUNCTION_CALL":
+                        function_call = values[0].value
+                        if function_call is not None:
+                            self.state.chat_history.append(
+                                ChatMessage(
+                                    role="FUNCTION",
+                                    content=FunctionCallChatMessageContent(
+                                        value=FunctionCallChatMessageContentValue(
+                                            name=function_call.name,
+                                            arguments=function_call.arguments,
+                                            id=function_call.id,
+                                        ),
+                                    ),
+                                )
+                            )
+            yield output
+
+
+def create_tool_router_node(
+    ml_model: str,
+    blocks: List[PromptBlock],
+    functions: List[FunctionDefinition],
+    prompt_inputs: Optional[EntityInputsInterface],
+) -> Type[ToolRouterNode]:
+    Ports = type("Ports", (), {})
+    for function in functions:
+        if function.name is None:
+            # We should not raise an error here since we filter out functions without names
+            raise ValueError("Function name is required")
+
+        function_name = function.name
+        port_condition = LazyReference(
+            lambda: (
+                ToolRouterNode.Outputs.results[0]["type"].equals("FUNCTION_CALL")
+                & ToolRouterNode.Outputs.results[0]["value"]["name"].equals(function_name)
+            )
+        )
+        port = Port.on_if(port_condition)
+        setattr(Ports, function_name, port)
+
+    setattr(Ports, "default", Port.on_else())
+
+    # Add a chat history block to blocks
+    blocks.append(
+        VariablePromptBlock(
+            block_type="VARIABLE",
+            input_variable="chat_history",
+            state=None,
+            cache_config=None,
+        )
+    )
+
+    node = type(
+        "ToolRouterNode",
+        (ToolRouterNode,),
+        {
+            "ml_model": ml_model,
+            "blocks": blocks,
+            "functions": functions,
+            "prompt_inputs": prompt_inputs,
+            "Ports": Ports,
+            "__module__": __name__,
+        },
+    )
+
+    return node
+
+
+def create_function_node(function: FunctionDefinition, function_callable: Callable[..., Any]) -> Type[FunctionNode]:
+    """
+    Create a FunctionNode class for a given function.
+
+    This ensures the callable is properly registered and can be called with the expected arguments.
+    """
+
+    # Create a class-level wrapper that calls the original function
+    def execute_function(self) -> BaseNode.Outputs:
+        outputs = self.state.meta.node_outputs.get(ToolRouterNode.Outputs.text)
+        # first parse into json
+        outputs = json.loads(outputs)
+        arguments = outputs["arguments"]
+
+        # Call the original function directly with the arguments
+        result = function_callable(**arguments)
+
+        self.state.chat_history.append(ChatMessage(role="FUNCTION", text=result))
+
+        return self.Outputs()
+
+    node = type(
+        f"FunctionNode_{function.name}",
+        (FunctionNode,),
+        {
+            "function": function,
+            "run": execute_function,
+            "__module__": __name__,
+        },
+    )
+
+    return node

vellum/workflows/nodes/utils.py

@@ -57,10 +57,12 @@ def create_adornment(
         class Subworkflow(BaseWorkflow):
             graph = inner_cls
 
-            # mypy is wrong here, this works and is defined
-            class Outputs(inner_cls.Outputs):  # type: ignore[name-defined]
+            class Outputs(BaseWorkflow.Outputs):
                 pass
 
+        for output_reference in inner_cls.Outputs:
+            setattr(Subworkflow.Outputs, output_reference.name, output_reference)
+
         dynamic_module = f"{inner_cls.__module__}.{inner_cls.__name__}.{ADORNMENT_MODULE_NAME}"
         # This dynamic module allows calls to `type_hints` to work
         sys.modules[dynamic_module] = ModuleType(dynamic_module)

vellum/workflows/outputs/base.py

@@ -147,8 +147,9 @@ class _BaseOutputsMeta(type):
             instance = vars(cls).get(name, undefined)
             if instance is undefined:
                 for base in cls.__mro__[1:]:
-                    if hasattr(base, name):
-                        instance = getattr(base, name)
+                    inherited_output_reference = getattr(base, name, undefined)
+                    if isinstance(inherited_output_reference, OutputReference):
+                        instance = inherited_output_reference.instance
                         break
 
             types = infer_types(cls, name)

vellum/workflows/references/output.py

@@ -1,4 +1,6 @@
+from functools import cached_property
 from queue import Queue
+from uuid import UUID, uuid4
 from typing import TYPE_CHECKING, Any, Generator, Generic, Optional, Tuple, Type, TypeVar, cast
 
 from pydantic import GetCoreSchemaHandler
@@ -31,6 +33,24 @@ class OutputReference(BaseDescriptor[_OutputType], Generic[_OutputType]):
     def outputs_class(self) -> Type["BaseOutputs"]:
         return self._outputs_class
 
+    @cached_property
+    def id(self) -> UUID:
+        self._outputs_class = self._outputs_class
+
+        node_class = getattr(self._outputs_class, "_node_class", None)
+        if not node_class:
+            return uuid4()
+
+        output_ids = getattr(node_class, "__output_ids__", {})
+        if not isinstance(output_ids, dict):
+            return uuid4()
+
+        output_id = output_ids.get(self.name)
+        if not isinstance(output_id, UUID):
+            return uuid4()
+
+        return output_id
+
     def resolve(self, state: "BaseState") -> _OutputType:
         node_output = state.meta.node_outputs.get(self, undefined)
         if isinstance(node_output, Queue):

vellum/workflows/state/base.py

@@ -7,13 +7,14 @@ import logging
 from queue import Queue
 from threading import Lock
 from uuid import UUID, uuid4
-from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, cast
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional, Set, Tuple, Type, Union, cast
 from typing_extensions import dataclass_transform
 
 from pydantic import GetCoreSchemaHandler, ValidationInfo, field_serializer, field_validator
 from pydantic_core import core_schema
 
 from vellum.core.pydantic_utilities import UniversalBaseModel
+from vellum.utils.uuid import is_valid_uuid
 from vellum.workflows.constants import undefined
 from vellum.workflows.edges.edge import Edge
 from vellum.workflows.inputs.base import BaseInputs
@@ -109,18 +110,30 @@ class NodeExecutionCache:
         self._node_executions_queued = defaultdict(list)
 
     @classmethod
-    def deserialize(cls, raw_data: dict, nodes: Dict[str, Type["BaseNode"]]):
+    def deserialize(cls, raw_data: dict, nodes: Dict[Union[str, UUID], Type["BaseNode"]]):
         cache = cls()
 
+        def get_node_class(node_id: Any) -> Optional[Type["BaseNode"]]:
+            if not isinstance(node_id, str):
+                return None
+
+            if is_valid_uuid(node_id):
+                return nodes.get(UUID(node_id))
+
+            return nodes.get(node_id)
+
         dependencies_invoked = raw_data.get("dependencies_invoked")
         if isinstance(dependencies_invoked, dict):
             for execution_id, dependencies in dependencies_invoked.items():
-                cache._dependencies_invoked[UUID(execution_id)] = {nodes[dep] for dep in dependencies if dep in nodes}
+                dependency_classes = {get_node_class(dep) for dep in dependencies}
+                cache._dependencies_invoked[UUID(execution_id)] = {
+                    dep_class for dep_class in dependency_classes if dep_class is not None
+                }
 
         node_executions_fulfilled = raw_data.get("node_executions_fulfilled")
         if isinstance(node_executions_fulfilled, dict):
             for node, execution_ids in node_executions_fulfilled.items():
-                node_class = nodes.get(node)
+                node_class = get_node_class(node)
                 if not node_class:
                     continue
 
@@ -131,7 +144,7 @@ class NodeExecutionCache:
         node_executions_initiated = raw_data.get("node_executions_initiated")
         if isinstance(node_executions_initiated, dict):
             for node, execution_ids in node_executions_initiated.items():
-                node_class = nodes.get(node)
+                node_class = get_node_class(node)
                 if not node_class:
                     continue
 
@@ -142,7 +155,7 @@ class NodeExecutionCache:
         node_executions_queued = raw_data.get("node_executions_queued")
         if isinstance(node_executions_queued, dict):
             for node, execution_ids in node_executions_queued.items():
-                node_class = nodes.get(node)
+                node_class = get_node_class(node)
                 if not node_class:
                     continue
 
@@ -193,17 +206,18 @@ class NodeExecutionCache:
     def dump(self) -> Dict[str, Any]:
         return {
             "dependencies_invoked": {
-                str(execution_id): [str(dep) for dep in dependencies]
+                str(execution_id): [str(dep.__id__) for dep in dependencies]
                 for execution_id, dependencies in self._dependencies_invoked.items()
             },
             "node_executions_initiated": {
-                str(node): list(execution_ids) for node, execution_ids in self._node_executions_initiated.items()
+                str(node.__id__): list(execution_ids) for node, execution_ids in self._node_executions_initiated.items()
             },
             "node_executions_fulfilled": {
-                str(node): execution_ids.dump() for node, execution_ids in self._node_executions_fulfilled.items()
+                str(node.__id__): execution_ids.dump()
+                for node, execution_ids in self._node_executions_fulfilled.items()
             },
             "node_executions_queued": {
-                str(node): execution_ids for node, execution_ids in self._node_executions_queued.items()
+                str(node.__id__): execution_ids for node, execution_ids in self._node_executions_queued.items()
             },
         }
 
@@ -279,7 +293,7 @@ class StateMeta(UniversalBaseModel):
 
     @field_serializer("node_outputs")
     def serialize_node_outputs(self, node_outputs: Dict[OutputReference, Any], _info: Any) -> Dict[str, Any]:
-        return {str(descriptor): value for descriptor, value in node_outputs.items()}
+        return {str(descriptor.id): value for descriptor, value in node_outputs.items()}
 
     @field_validator("node_outputs", mode="before")
     @classmethod
@@ -290,15 +304,22 @@ class StateMeta(UniversalBaseModel):
                 return node_outputs
 
             raw_workflow_nodes = workflow_definition.get_nodes()
-            workflow_node_outputs = {}
+            workflow_node_outputs: Dict[Union[str, UUID], OutputReference] = {}
             for node in raw_workflow_nodes:
                 for output in node.Outputs:
                     workflow_node_outputs[str(output)] = output
+                    output_id = node.__output_ids__.get(output.name)
+                    if output_id:
+                        workflow_node_outputs[output_id] = output
 
             node_output_keys = list(node_outputs.keys())
             deserialized_node_outputs = {}
             for node_output_key in node_output_keys:
-                output_reference = workflow_node_outputs.get(node_output_key)
+                if is_valid_uuid(node_output_key):
+                    output_reference = workflow_node_outputs.get(UUID(node_output_key))
+                else:
+                    output_reference = workflow_node_outputs.get(node_output_key)
+
                 if not output_reference:
                     continue
 
@@ -316,10 +337,11 @@ class StateMeta(UniversalBaseModel):
             if not workflow_definition:
                 return node_execution_cache
 
-            nodes_cache: Dict[str, Type["BaseNode"]] = {}
+            nodes_cache: Dict[Union[str, UUID], Type["BaseNode"]] = {}
             raw_workflow_nodes = workflow_definition.get_nodes()
             for node in raw_workflow_nodes:
                 nodes_cache[str(node)] = node
+                nodes_cache[node.__id__] = node
 
             return NodeExecutionCache.deserialize(node_execution_cache, nodes_cache)
 

vellum/workflows/state/tests/test_state.py

@@ -47,6 +47,9 @@ class MockNode(BaseNode):
         baz: str
 
 
+MOCK_NODE_OUTPUT_ID = "e4dc3136-0c27-4bda-b3ab-ea355d5219d6"
+
+
 def test_state_snapshot__node_attribute_edit():
     # GIVEN an initial state instance
     state = MockState(foo="bar")
@@ -144,7 +147,7 @@ def test_state_json_serialization__with_node_output_updates():
     json_state = json.loads(json.dumps(state, cls=DefaultStateEncoder))
 
     # THEN the state is serialized correctly
-    assert json_state["meta"]["node_outputs"] == {"MockNode.Outputs.baz": "hello"}
+    assert json_state["meta"]["node_outputs"] == {MOCK_NODE_OUTPUT_ID: "hello"}
 
 
 def test_state_deepcopy__with_external_input_updates():
@@ -185,7 +188,7 @@ def test_state_json_serialization__with_queue():
     json_state = json.loads(json.dumps(state, cls=DefaultStateEncoder))
 
     # THEN the state is serialized correctly with the queue turned into a list
-    assert json_state["meta"]["node_outputs"] == {"MockNode.Outputs.baz": ["test1", "test2"]}
+    assert json_state["meta"]["node_outputs"] == {MOCK_NODE_OUTPUT_ID: ["test1", "test2"]}
 
 
 def test_state_snapshot__deepcopy_fails__logs_error(mock_deepcopy, mock_logger):

vellum/workflows/types/stack.py

@@ -37,3 +37,14 @@ class Stack(Generic[_T]):
 
     def dump(self) -> List[_T]:
         return [item for item in self._items][::-1]
+
+    @classmethod
+    def from_list(cls, items: List[_T]) -> "Stack[_T]":
+        stack = cls()
+        stack.extend(items)
+        return stack
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Stack):
+            return False
+        return self._items == other._items

vellum/workflows/workflows/base.py

@@ -80,6 +80,11 @@ class _BaseWorkflowMeta(type):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         if "graph" not in dct:
             dct["graph"] = set()
+            for base in bases:
+                base_graph = getattr(base, "graph", None)
+                if base_graph:
+                    dct["graph"] = base_graph
+                    break
 
         if "Outputs" in dct:
             outputs_class = dct["Outputs"]