vellum-ai 0.14.55__py3-none-any.whl → 0.14.57__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.55",
+            "X-Fern-SDK-Version": "0.14.57",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

vellum/workflows/nodes/bases/base.py

@@ -8,7 +8,6 @@ from typing import Any, Dict, Generic, Iterator, Optional, Set, Tuple, Type, Typ
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
 from vellum.workflows.descriptors.utils import is_unresolved, resolve_value
-from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.graph import Graph
@@ -317,8 +316,15 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
                 """
                 # Check if all dependencies have invoked this node
                 dependencies_invoked = state.meta.node_execution_cache._dependencies_invoked.get(node_span_id, set())
-                all_deps_invoked = all(dep in dependencies_invoked for dep in dependencies)
-
+                node_classes_invoked = {
+                    state.meta.node_execution_cache.__node_execution_lookup__[dep]
+                    for dep in dependencies_invoked
+                    if dep in state.meta.node_execution_cache.__node_execution_lookup__
+                }
+                if len(node_classes_invoked) != len(dependencies):
+                    return False
+
+                all_deps_invoked = all(dep in node_classes_invoked for dep in dependencies)
                 return all_deps_invoked
 
             raise NodeException(
@@ -328,7 +334,7 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
 
         @classmethod
         def _queue_node_execution(
-            cls, state: StateType, dependencies: Set["Type[BaseNode]"], invoked_by: Optional[Edge] = None
+            cls, state: StateType, dependencies: Set["Type[BaseNode]"], invoked_by: Optional[UUID] = None
         ) -> UUID:
             """
             Queues a future execution of a node, returning the span id of the execution.
@@ -341,19 +347,21 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
             if not invoked_by:
                 return execution_id
 
+            if invoked_by not in state.meta.node_execution_cache.__node_execution_lookup__:
+                return execution_id
+
             if cls.merge_behavior not in {MergeBehavior.AWAIT_ANY, MergeBehavior.AWAIT_ALL}:
                 return execution_id
 
-            source_node = invoked_by.from_port.node_class
             for queued_node_execution_id in state.meta.node_execution_cache._node_executions_queued[cls.node_class]:
-                if source_node not in state.meta.node_execution_cache._dependencies_invoked[queued_node_execution_id]:
+                if invoked_by not in state.meta.node_execution_cache._dependencies_invoked[queued_node_execution_id]:
                     state.meta.node_execution_cache._invoke_dependency(
-                        queued_node_execution_id, cls.node_class, source_node, dependencies
+                        queued_node_execution_id, cls.node_class, invoked_by, dependencies
                     )
                     return queued_node_execution_id
 
             state.meta.node_execution_cache._node_executions_queued[cls.node_class].append(execution_id)
-            state.meta.node_execution_cache._invoke_dependency(execution_id, cls.node_class, source_node, dependencies)
+            state.meta.node_execution_cache._invoke_dependency(execution_id, cls.node_class, invoked_by, dependencies)
             return execution_id
 
     class Execution(metaclass=_BaseNodeExecutionMeta):

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

@@ -70,6 +70,12 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
 
                     for output_descriptor, output_value in event.outputs:
                         setattr(node_outputs, output_descriptor.name, output_value)
+
+                        if self.__wrapped_node__:
+                            inner_desc = getattr(self.__wrapped_node__.Outputs, output_descriptor.name, None)
+                            if inner_desc:
+                                self.state.meta.node_outputs[inner_desc] = output_value
+
                 elif event.name == "workflow.execution.paused":
                     exception = NodeException(
                         code=WorkflowErrorCode.INVALID_OUTPUTS,

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

@@ -47,11 +47,18 @@ class APINode(BaseAPINode):
             self.bearer_token_value, VellumSecret
         ):
             bearer_token = self.bearer_token_value
+
+        final_headers = {**headers, **header_overrides}
+
+        vellum_client_wrapper = self._context.vellum_client._client_wrapper
+        if self.url.startswith(vellum_client_wrapper._environment.default) and "X-API-Key" not in final_headers:
+            final_headers["X-API-Key"] = vellum_client_wrapper.api_key
+
         return self._run(
             method=self.method,
             url=self.url,
             data=self.data,
             json=self.json,
-            headers={**headers, **header_overrides},
+            headers=final_headers,
             bearer_token=bearer_token,
         )

vellum/workflows/nodes/displayable/api_node/tests/test_api_node.py

@@ -19,7 +19,7 @@ def test_run_workflow__secrets(vellum_client):
     class SimpleBaseAPINode(APINode):
         method = APIRequestMethod.POST
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         json = {
             "key": "value",
         }
@@ -32,7 +32,9 @@ def test_run_workflow__secrets(vellum_client):
     terminal = node.run()
 
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
+    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
     bearer_token = vellum_client.execute_api.call_args.kwargs["bearer_token"]
     assert bearer_token == ClientVellumSecret(name="secret")
     assert terminal.headers == {"X-Response-Header": "bar"}
@@ -45,7 +47,7 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
     class SimpleAPINode(APINode):
         method = APIRequestMethod.GET
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         json = {
             "key": "value",
         }
@@ -65,7 +67,9 @@ def test_api_node_raises_error_when_api_call_fails(vellum_client):
 
     # AND the API call should have been made
     assert vellum_client.execute_api.call_count == 1
+    assert vellum_client.execute_api.call_args.kwargs["url"] == "https://example.vellum.ai"
     assert vellum_client.execute_api.call_args.kwargs["body"] == {"key": "value"}
+    assert vellum_client.execute_api.call_args.kwargs["headers"] == {"X-Test-Header": "foo"}
 
 
 def test_api_node_defaults_to_get_method(vellum_client):
@@ -80,7 +84,7 @@ def test_api_node_defaults_to_get_method(vellum_client):
     # AND an API node without a method specified
     class SimpleAPINodeWithoutMethod(APINode):
         authorization_type = AuthorizationType.BEARER_TOKEN
-        url = "https://api.vellum.ai"
+        url = "https://example.vellum.ai"
         headers = {
             "X-Test-Header": "foo",
         }
@@ -95,3 +99,62 @@ def test_api_node_defaults_to_get_method(vellum_client):
     assert vellum_client.execute_api.call_count == 1
     method = vellum_client.execute_api.call_args.kwargs["method"]
     assert method == APIRequestMethod.GET.value
+
+
+def test_api_node__detects_client_environment_urls__adds_token(mock_httpx_transport, mock_requests, monkeypatch):
+    # GIVEN an API node with a URL pointing back to Vellum
+    class SimpleAPINodeToVellum(APINode):
+        url = "https://api.vellum.ai"
+
+    # AND a mock request sent to the Vellum API would return a 200
+    mock_response = mock_requests.get(
+        "https://api.vellum.ai",
+        status_code=200,
+        json={"data": [1, 2, 3]},
+    )
+
+    # AND an api key is set
+    monkeypatch.setenv("VELLUM_API_KEY", "vellum-api-key-1234")
+
+    # WHEN we run the node
+    node = SimpleAPINodeToVellum()
+    node.run()
+
+    # THEN the execute_api method should not have been called
+    mock_httpx_transport.handle_request.assert_not_called()
+
+    # AND the vellum API should have been called with the correct headers
+    assert mock_response.last_request
+    assert mock_response.last_request.headers["X-API-Key"] == "vellum-api-key-1234"
+
+
+def test_api_node__detects_client_environment_urls__does_not_override_headers(
+    mock_httpx_transport, mock_requests, monkeypatch
+):
+    # GIVEN an API node with a URL pointing back to Vellum
+    class SimpleAPINodeToVellum(APINode):
+        url = "https://api.vellum.ai"
+        headers = {
+            "X-API-Key": "vellum-api-key-5678",
+        }
+
+    # AND a mock request sent to the Vellum API would return a 200
+    mock_response = mock_requests.get(
+        "https://api.vellum.ai",
+        status_code=200,
+        json={"data": [1, 2, 3]},
+    )
+
+    # AND an api key is set
+    monkeypatch.setenv("VELLUM_API_KEY", "vellum-api-key-1234")
+
+    # WHEN we run the node
+    node = SimpleAPINodeToVellum()
+    node.run()
+
+    # THEN the execute_api method should not have been called
+    mock_httpx_transport.handle_request.assert_not_called()
+
+    # AND the vellum API should have been called with the correct headers
+    assert mock_response.last_request
+    assert mock_response.last_request.headers["X-API-Key"] == "vellum-api-key-5678"

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

@@ -30,8 +30,8 @@ from vellum.workflows.nodes.displayable.bases.base_prompt_node import BasePrompt
 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
+from vellum.workflows.types.generics import StateType, is_workflow_class
+from vellum.workflows.utils.functions import compile_function_definition, compile_workflow_function_definition
 
 
 class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
@@ -97,14 +97,18 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             "execution_context": execution_context.model_dump(mode="json"),
             **request_options.get("additional_body_parameters", {}),
         }
-        normalized_functions = (
-            [
-                function if isinstance(function, FunctionDefinition) else compile_function_definition(function)
-                for function in self.functions
-            ]
-            if self.functions
-            else None
-        )
+
+        normalized_functions: Optional[List[FunctionDefinition]] = None
+
+        if self.functions:
+            normalized_functions = []
+            for function in self.functions:
+                if isinstance(function, FunctionDefinition):
+                    normalized_functions.append(function)
+                elif is_workflow_class(function):
+                    normalized_functions.append(compile_workflow_function_definition(function))
+                else:
+                    normalized_functions.append(compile_function_definition(function))
 
         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

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

@@ -1,6 +1,8 @@
 from collections.abc import Callable
 from typing import Any, ClassVar, List, Optional
 
+from pydash import snake_case
+
 from vellum import ChatMessage, PromptBlock
 from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.errors.types import WorkflowErrorCode
@@ -101,7 +103,7 @@ class ToolCallingNode(BaseNode):
         )
 
         self._function_nodes = {
-            function.__name__: create_function_node(
+            snake_case(function.__name__): create_function_node(
                 function=function,
                 tool_router_node=self.tool_router_node,
             )

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

@@ -2,22 +2,30 @@ from collections.abc import Callable
 import json
 from typing import Any, Iterator, List, Optional, Type, cast
 
-from vellum import ChatMessage, FunctionDefinition, PromptBlock
+from pydash import snake_case
+
+from vellum import ChatMessage, 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.string_chat_message_content import StringChatMessageContent
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
+from vellum.workflows.errors.types import WorkflowErrorCode
+from vellum.workflows.exceptions import NodeException
+from vellum.workflows.inputs.base import BaseInputs
 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.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
+from vellum.workflows.types.generics import is_workflow_class
 
 
 class FunctionNode(BaseNode):
     """Node that executes a specific function."""
 
-    function: FunctionDefinition
+    pass
 
 
 class ToolRouterNode(InlinePromptNode):
@@ -57,26 +65,32 @@ def create_tool_router_node(
     functions: List[Callable[..., Any]],
     prompt_inputs: Optional[EntityInputsInterface],
 ) -> Type[ToolRouterNode]:
-    Ports = type("Ports", (), {})
-    for function in functions:
-        function_name = function.__name__
-
-        # Avoid using lambda to capture function_name
-        # lambda will capture the function_name by reference,
-        # and if the function_name is changed, the port_condition will also change.
-        def create_port_condition(fn_name):
-            return LazyReference(
-                lambda: (
-                    node.Outputs.results[0]["type"].equals("FUNCTION_CALL")
-                    & node.Outputs.results[0]["value"]["name"].equals(fn_name)
+    if functions and len(functions) > 0:
+        # If we have functions, create dynamic ports for each function
+        Ports = type("Ports", (), {})
+        for function in functions:
+            function_name = snake_case(function.__name__)
+
+            # Avoid using lambda to capture function_name
+            # lambda will capture the function_name by reference,
+            # and if the function_name is changed, the port_condition will also change.
+            def create_port_condition(fn_name):
+                return LazyReference(
+                    lambda: (
+                        node.Outputs.results[0]["type"].equals("FUNCTION_CALL")
+                        & node.Outputs.results[0]["value"]["name"].equals(fn_name)
+                    )
                 )
-            )
 
-        port_condition = create_port_condition(function_name)
-        port = Port.on_if(port_condition)
-        setattr(Ports, function_name, port)
+            port_condition = create_port_condition(function_name)
+            port = Port.on_if(port_condition)
+            setattr(Ports, function_name, port)
 
-    setattr(Ports, "default", Port.on_else())
+        # Add the else port for when no function conditions match
+        setattr(Ports, "default", Port.on_else())
+    else:
+        # If no functions exist, create a simple Ports class with just a default port
+        Ports = type("Ports", (), {"default": Port(default=True)})
 
     # Add a chat history block to blocks
     blocks.append(
@@ -121,10 +135,41 @@ def create_function_node(function: Callable[..., Any], tool_router_node: Type[To
         outputs = json.loads(outputs)
         arguments = outputs["arguments"]
 
-        # Call the original function directly with the arguments
-        result = function(**arguments)
+        # Call the function based on its type
+        if is_workflow_class(function):
+            # Dynamically define an Inputs subclass of BaseInputs
+            Inputs = type(
+                "Inputs",
+                (BaseInputs,),
+                {"__annotations__": {k: type(v) for k, v in arguments.items()}},
+            )
+
+            # Create an instance with arguments
+            inputs_instance = Inputs(**arguments)
 
-        self.state.chat_history.append(ChatMessage(role="FUNCTION", text=result))
+            workflow = function()
+            terminal_event = workflow.run(
+                inputs=inputs_instance,
+            )
+            if terminal_event.name == "workflow.execution.paused":
+                raise NodeException(
+                    code=WorkflowErrorCode.INVALID_OUTPUTS,
+                    message="Subworkflow unexpectedly paused",
+                )
+            elif terminal_event.name == "workflow.execution.fulfilled":
+                result = terminal_event.outputs
+            elif terminal_event.name == "workflow.execution.rejected":
+                raise Exception(f"Workflow execution rejected: {terminal_event.error}")
+        else:
+            # If it's a regular callable, call it directly
+            result = function(**arguments)
+
+        self.state.chat_history.append(
+            ChatMessage(
+                role="FUNCTION",
+                content=StringChatMessageContent(value=json.dumps(result, cls=DefaultStateEncoder)),
+            )
+        )
 
         return self.Outputs()
 
@@ -132,7 +177,6 @@ def create_function_node(function: Callable[..., Any], tool_router_node: Type[To
         f"FunctionNode_{function.__name__}",
         (FunctionNode,),
         {
-            "function": function,
             "run": execute_function,
             "__module__": __name__,
         },

vellum/workflows/runner/runner.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Dict, Generic, Iterable, Iterator, Option
 from vellum.workflows.constants import undefined
 from vellum.workflows.context import ExecutionContext, execution_context, get_execution_context
 from vellum.workflows.descriptors.base import BaseDescriptor
-from vellum.workflows.edges.edge import Edge
 from vellum.workflows.errors import WorkflowError, WorkflowErrorCode
 from vellum.workflows.events import (
     NodeExecutionFulfilledEvent,
@@ -143,7 +142,7 @@ class WorkflowRunner(Generic[StateType]):
         self._workflow_event_inner_queue: Queue[WorkflowEvent] = Queue()
 
         self._max_concurrency = max_concurrency
-        self._concurrency_queue: Queue[Tuple[StateType, Type[BaseNode], Optional[Edge]]] = Queue()
+        self._concurrency_queue: Queue[Tuple[StateType, Type[BaseNode], Optional[UUID]]] = Queue()
 
         # This queue is responsible for sending events from WorkflowRunner to the background thread
         # for user defined emitters
@@ -389,7 +388,12 @@ class WorkflowRunner(Generic[StateType]):
         ):
             self._run_work_item(node, span_id)
 
-    def _handle_invoked_ports(self, state: StateType, ports: Optional[Iterable[Port]]) -> None:
+    def _handle_invoked_ports(
+        self,
+        state: StateType,
+        ports: Optional[Iterable[Port]],
+        invoked_by: Optional[UUID],
+    ) -> None:
         if not ports:
             return
 
@@ -402,9 +406,9 @@ class WorkflowRunner(Generic[StateType]):
                     next_state = state
 
                 if self._max_concurrency:
-                    self._concurrency_queue.put((next_state, edge.to_node, edge))
+                    self._concurrency_queue.put((next_state, edge.to_node, invoked_by))
                 else:
-                    self._run_node_if_ready(next_state, edge.to_node, edge)
+                    self._run_node_if_ready(next_state, edge.to_node, invoked_by)
 
         if self._max_concurrency:
             num_nodes_to_run = self._max_concurrency - len(self._active_nodes_by_execution_id)
@@ -412,14 +416,14 @@ class WorkflowRunner(Generic[StateType]):
                 if self._concurrency_queue.empty():
                     break
 
-                next_state, node_class, invoked_edge = self._concurrency_queue.get()
-                self._run_node_if_ready(next_state, node_class, invoked_edge)
+                next_state, node_class, invoked_by = self._concurrency_queue.get()
+                self._run_node_if_ready(next_state, node_class, invoked_by)
 
     def _run_node_if_ready(
         self,
         state: StateType,
         node_class: Type[BaseNode],
-        invoked_by: Optional[Edge] = None,
+        invoked_by: Optional[UUID] = None,
     ) -> None:
         with state.__lock__:
             for descriptor in node_class.ExternalInputs:
@@ -482,7 +486,7 @@ class WorkflowRunner(Generic[StateType]):
                     )
                 )
 
-            self._handle_invoked_ports(node.state, event.invoked_ports)
+            self._handle_invoked_ports(node.state, event.invoked_ports, event.span_id)
 
             return None
 
@@ -508,7 +512,7 @@ class WorkflowRunner(Generic[StateType]):
                             )
                         )
 
-            self._handle_invoked_ports(node.state, event.invoked_ports)
+            self._handle_invoked_ports(node.state, event.invoked_ports, event.span_id)
 
             return None
 

vellum/workflows/state/base.py

@@ -93,7 +93,8 @@ def _make_snapshottable(value: Any, snapshot_callback: Callable[[], None]) -> An
 NodeExecutionsFulfilled = Dict[Type["BaseNode"], Stack[UUID]]
 NodeExecutionsInitiated = Dict[Type["BaseNode"], Set[UUID]]
 NodeExecutionsQueued = Dict[Type["BaseNode"], List[UUID]]
-DependenciesInvoked = Dict[UUID, Set[Type["BaseNode"]]]
+NodeExecutionLookup = Dict[UUID, Type["BaseNode"]]
+DependenciesInvoked = Dict[UUID, Set[UUID]]
 
 
 class NodeExecutionCache:
@@ -102,11 +103,15 @@ class NodeExecutionCache:
     _node_executions_queued: NodeExecutionsQueued
     _dependencies_invoked: DependenciesInvoked
 
+    # Derived fields - no need to serialize
+    __node_execution_lookup__: NodeExecutionLookup
+
     def __init__(self) -> None:
         self._dependencies_invoked = defaultdict(set)
         self._node_executions_fulfilled = defaultdict(Stack[UUID])
         self._node_executions_initiated = defaultdict(set)
         self._node_executions_queued = defaultdict(list)
+        self.__node_execution_lookup__ = {}
 
     @classmethod
     def deserialize(cls, raw_data: dict, nodes: Dict[Union[str, UUID], Type["BaseNode"]]):
@@ -124,10 +129,8 @@ class NodeExecutionCache:
         dependencies_invoked = raw_data.get("dependencies_invoked")
         if isinstance(dependencies_invoked, dict):
             for execution_id, dependencies in dependencies_invoked.items():
-                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
-                }
+                dependency_execution_ids = {UUID(dep) for dep in dependencies if is_valid_uuid(dep)}
+                cache._dependencies_invoked[UUID(execution_id)] = dependency_execution_ids
 
         node_executions_fulfilled = raw_data.get("node_executions_fulfilled")
         if isinstance(node_executions_fulfilled, dict):
@@ -151,6 +154,10 @@ class NodeExecutionCache:
                     {UUID(execution_id) for execution_id in execution_ids}
                 )
 
+            for node_class, execution_ids in cache._node_executions_initiated.items():
+                for execution_id in execution_ids:
+                    cache.__node_execution_lookup__[execution_id] = node_class
+
         node_executions_queued = raw_data.get("node_executions_queued")
         if isinstance(node_executions_queued, dict):
             for node, execution_ids in node_executions_queued.items():
@@ -166,18 +173,29 @@ class NodeExecutionCache:
         self,
         execution_id: UUID,
         node: Type["BaseNode"],
-        dependency: Type["BaseNode"],
+        invoked_by: UUID,
         dependencies: Set["Type[BaseNode]"],
     ) -> None:
-        self._dependencies_invoked[execution_id].add(dependency)
-        if all(dep in self._dependencies_invoked[execution_id] for dep in dependencies):
-            self._node_executions_queued[node].remove(execution_id)
+        self._dependencies_invoked[execution_id].add(invoked_by)
+        invoked_node_classes = {
+            self.__node_execution_lookup__[dep]
+            for dep in self._dependencies_invoked[execution_id]
+            if dep in self.__node_execution_lookup__
+        }
+        if len(invoked_node_classes) != len(dependencies):
+            return
+
+        if any(dep not in invoked_node_classes for dep in dependencies):
+            return
+
+        self._node_executions_queued[node].remove(execution_id)
 
     def is_node_execution_initiated(self, node: Type["BaseNode"], execution_id: UUID) -> bool:
         return execution_id in self._node_executions_initiated[node]
 
     def initiate_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
         self._node_executions_initiated[node].add(execution_id)
+        self.__node_execution_lookup__[execution_id] = node
 
     def fulfill_node_execution(self, node: Type["BaseNode"], execution_id: UUID) -> None:
         self._node_executions_fulfilled[node].push(execution_id)
@@ -188,7 +206,7 @@ class NodeExecutionCache:
     def dump(self) -> Dict[str, Any]:
         return {
             "dependencies_invoked": {
-                str(execution_id): [str(dep.__id__) for dep in dependencies]
+                str(execution_id): [str(dep) for dep in dependencies]
                 for execution_id, dependencies in self._dependencies_invoked.items()
             },
             "node_executions_initiated": {

vellum/workflows/state/encoder.py

@@ -1,6 +1,7 @@
 from dataclasses import asdict, is_dataclass
 from datetime import datetime
 import enum
+import inspect
 from json import JSONEncoder
 from queue import Queue
 from uuid import UUID
@@ -59,7 +60,17 @@ class DefaultStateEncoder(JSONEncoder):
             return str(obj)
 
         if callable(obj):
-            return compile_function_definition(obj)
+            function_definition = compile_function_definition(obj)
+            try:
+                source_code = inspect.getsource(obj)
+            except Exception:
+                source_code = f"<source code not available for {obj.__name__}>"
+
+            return {
+                "type": "CODE_EXECUTION",
+                "definition": function_definition,
+                "src": source_code,
+            }
 
         if obj.__class__ in self.encoders:
             return self.encoders[obj.__class__](obj)

vellum/workflows/utils/functions.py

@@ -1,12 +1,16 @@
 import dataclasses
 import inspect
-from typing import Any, Callable, Optional, Union, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type, Union, get_args, get_origin
 
 from pydantic import BaseModel
 from pydantic_core import PydanticUndefined
+from pydash import snake_case
 
 from vellum.client.types.function_definition import FunctionDefinition
 
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
 type_map = {
     str: "string",
     int: "integer",
@@ -108,5 +112,42 @@ def compile_function_definition(function: Callable) -> FunctionDefinition:
 
     return FunctionDefinition(
         name=function.__name__,
+        description=function.__doc__,
+        parameters=parameters,
+    )
+
+
+def compile_workflow_function_definition(workflow_class: Type["BaseWorkflow"]) -> FunctionDefinition:
+    """
+    Converts a base workflow class into our Vellum-native FunctionDefinition type.
+    """
+    # Get the inputs class for the workflow
+    inputs_class = workflow_class.get_inputs_class()
+    vars_inputs_class = vars(inputs_class)
+
+    properties = {}
+    required = []
+    defs: dict[str, Any] = {}
+
+    for name, field_type in inputs_class.__annotations__.items():
+        if name.startswith("__"):
+            continue
+
+        properties[name] = _compile_annotation(field_type, defs)
+
+        # Check if the field has a default value
+        if name not in vars_inputs_class:
+            required.append(name)
+        else:
+            # Field has a default value
+            properties[name]["default"] = vars_inputs_class[name]
+
+    parameters = {"type": "object", "properties": properties, "required": required}
+    if defs:
+        parameters["$defs"] = defs
+
+    return FunctionDefinition(
+        name=snake_case(workflow_class.__name__),
+        description=workflow_class.__doc__,
         parameters=parameters,
     )