vellum-ai 0.14.62__tar.gz → 0.14.64__tar.gz

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 0.14.62
+Version: 0.14.64
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -20,22 +20,21 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.8
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Typing :: Typed
-Requires-Dist: Jinja2 (>=3.1.2,<4.0.0)
-Requires-Dist: cdktf (>=0.20.5,<0.21.0)
-Requires-Dist: click (==8.1.7)
-Requires-Dist: docker (==7.1.0)
+Requires-Dist: Jinja2 (>=3.1.0,<4.0.0)
+Requires-Dist: click (>=8.1.0,<9.0.0)
+Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx (>=0.21.2)
-Requires-Dist: openai (>=1.0.0)
+Requires-Dist: openai (>=1.0.0,<2.0.0)
 Requires-Dist: orderly-set (>=5.2.2,<6.0.0)
 Requires-Dist: publication (==0.0.3)
 Requires-Dist: pydantic (>=1.9.2)
 Requires-Dist: pydantic-core (>=2.18.2,<3.0.0)
-Requires-Dist: pydash (==7.0.6)
-Requires-Dist: python-dotenv (==1.0.1)
-Requires-Dist: pytz (==2025.1)
-Requires-Dist: pyyaml (==6.0.2)
-Requires-Dist: requests (==2.32.3)
-Requires-Dist: tomli (==2.0.2)
+Requires-Dist: pydash (>=7.0.0,<8.0.0)
+Requires-Dist: python-dotenv (>=1.0.0,<2.0.0)
+Requires-Dist: pytz (>=2022.0,<2026.0)
+Requires-Dist: pyyaml (>=6.0.0,<7.0.0)
+Requires-Dist: requests (>=2.31.0,<3.0.0)
+Requires-Dist: tomli (>=2.0.0,<3.0.0)
 Requires-Dist: typing_extensions (>=4.0.0)
 Project-URL: Repository, https://github.com/vellum-ai/vellum-python-sdks
 Description-Content-Type: text/markdown

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/ee/vellum_ee/workflows/display/nodes/vellum/map_node.py

@@ -3,6 +3,7 @@ from typing import Generic, Optional, TypeVar, cast
 
 from vellum.workflows.nodes import MapNode
 from vellum.workflows.types.core import JsonObject
+from vellum.workflows.workflows.base import BaseWorkflow
 from vellum_ee.workflows.display.nodes.utils import raise_if_descriptor
 from vellum_ee.workflows.display.nodes.vellum.base_adornment_node import BaseAdornmentNodeDisplay
 from vellum_ee.workflows.display.nodes.vellum.utils import create_node_input
@@ -21,7 +22,7 @@ class BaseMapNodeDisplay(BaseAdornmentNodeDisplay[_MapNodeType], Generic[_MapNod
         node = self._node
         node_id = self.node_id
 
-        subworkflow = raise_if_descriptor(node.subworkflow)
+        subworkflow = cast(type[BaseWorkflow], raise_if_descriptor(node.subworkflow))
 
         items_node_input = create_node_input(
             node_id=node_id,

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/ee/vellum_ee/workflows/display/nodes/vellum/tests/test_tool_calling_node.py

@@ -1,3 +1,4 @@
+from vellum.client.types.code_execution_package import CodeExecutionPackage
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs import BaseInputs
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
@@ -122,3 +123,102 @@ def test_serialize_node__prompt_inputs__mixed_values():
             ],
         },
     }
+
+
+def test_serialize_node__function_configs():
+    # GIVEN a tool calling node with function packages
+    def foo():
+        pass
+
+    def bar():
+        pass
+
+    class MyToolCallingNode(ToolCallingNode):
+        functions = [foo, bar]
+        function_configs = {
+            "foo": {
+                "runtime": "PYTHON_3_11_6",
+                "packages": [
+                    CodeExecutionPackage(name="first_package", version="1.0.0"),
+                    CodeExecutionPackage(name="second_package", version="2.0.0"),
+                ],
+            },
+            "bar": {
+                "runtime": "PYTHON_3_11_6",
+                "packages": [CodeExecutionPackage(name="third_package", version="3.0.0")],
+            },
+        }
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the function packages
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    function_configs_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "function_configs"
+    )
+
+    assert function_configs_attribute == {
+        "id": "90cc5fc7-9fb3-450f-be5a-e90e7412a601",
+        "name": "function_configs",
+        "value": {
+            "type": "CONSTANT_VALUE",
+            "value": {
+                "type": "JSON",
+                "value": {
+                    "foo": {
+                        "runtime": "PYTHON_3_11_6",
+                        "packages": [
+                            {"version": "1.0.0", "name": "first_package"},
+                            {"version": "2.0.0", "name": "second_package"},
+                        ],
+                    },
+                    "bar": {"runtime": "PYTHON_3_11_6", "packages": [{"version": "3.0.0", "name": "third_package"}]},
+                },
+            },
+        },
+    }
+
+
+def test_serialize_node__function_configs__none():
+    # GIVEN a tool calling node with no function configs
+    def foo():
+        pass
+
+    class MyToolCallingNode(ToolCallingNode):
+        functions = [foo]
+
+    # AND a workflow with the tool calling node
+    class Workflow(BaseWorkflow):
+        graph = MyToolCallingNode
+
+    # WHEN the workflow is serialized
+    workflow_display = get_workflow_display(workflow_class=Workflow)
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN the node should properly serialize the functions
+    my_tool_calling_node = next(
+        node
+        for node in serialized_workflow["workflow_raw_data"]["nodes"]
+        if node["id"] == str(MyToolCallingNode.__id__)
+    )
+
+    function_configs_attribute = next(
+        attribute for attribute in my_tool_calling_node["attributes"] if attribute["name"] == "function_configs"
+    )
+
+    assert function_configs_attribute == {
+        "id": "5f63f2aa-72d7-47ad-a552-630753418b7d",
+        "name": "function_configs",
+        "value": {"type": "CONSTANT_VALUE", "value": {"type": "JSON", "value": None}},
+    }

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_serialization.py

@@ -150,6 +150,22 @@ def test_serialize_workflow():
                     },
                 },
             },
+            {
+                "id": "a4e3bc9f-7112-4d2f-94fb-7362a85db27a",
+                "name": "function_configs",
+                "value": {
+                    "type": "CONSTANT_VALUE",
+                    "value": {
+                        "type": "JSON",
+                        "value": {
+                            "get_current_weather": {
+                                "runtime": "PYTHON_3_11_6",
+                                "packages": [{"version": "2.26.0", "name": "requests"}],
+                            }
+                        },
+                    },
+                },
+            },
             {
                 "id": "0f6dc102-3460-4963-91fa-7ba85d65ef7a",
                 "name": "prompt_inputs",
@@ -167,11 +183,6 @@ def test_serialize_workflow():
                     ],
                 },
             },
-            {
-                "id": "5c041b7d-732c-4773-a93a-32211f2af0b3",
-                "name": "max_tool_calls",
-                "value": {"type": "CONSTANT_VALUE", "value": {"type": "NUMBER", "value": 1.0}},
-            },
         ],
         "outputs": [
             {"id": "e62bc785-a914-4066-b79e-8c89a5d0ec6c", "name": "text", "type": "STRING", "value": None},

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/ee/vellum_ee/workflows/display/utils/expressions.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, cast
+from typing import TYPE_CHECKING, Any, Dict, List, cast
 
 from vellum.client.types.logical_operator import LogicalOperator
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -190,6 +190,16 @@ def _serialize_condition(display_context: "WorkflowDisplayContext", condition: B
     raise UnsupportedSerializationException(f"Unsupported condition type: {condition.__class__.__name__}")
 
 
+def serialize_key(key: Any) -> str:
+    """Serialize dictionary keys to strings, handling function objects properly."""
+    if callable(key):
+        return key.__name__
+    elif isinstance(key, str):
+        return key
+    else:
+        return str(key)
+
+
 def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> JsonObject:
     if isinstance(value, ConstantValueReference):
         return serialize_value(display_context, value._value)
@@ -264,10 +274,10 @@ def serialize_value(display_context: "WorkflowDisplayContext", value: Any) -> Js
             }
 
     if isinstance(value, dict):
-        serialized_entries = [
+        serialized_entries: List[Dict[str, Any]] = [
             {
                 "id": str(uuid4_from_hash(f"{key}|{val}")),
-                "key": key,
+                "key": serialize_key(key),
                 "value": serialize_value(display_context, val),
             }
             for key, val in value.items()

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "0.14.62"
+version = "0.14.64"
 description = ""
 readme = "README.md"
 authors = []
@@ -36,22 +36,21 @@ packages = [
 
 [tool.poetry.dependencies]
 python = "^3.9"
-Jinja2 = "^3.1.2"
-cdktf = "^0.20.5"
-click = "8.1.7"
-docker = "7.1.0"
+Jinja2 = ">=3.1.0,<4.0.0"
+click = ">=8.1.0,<9.0.0"
+docker = ">=7.1.0,<8.0.0"
 httpx = ">=0.21.2"
-openai = ">= 1.0.0"
+openai = ">=1.0.0,<2.0.0"
 orderly-set = "^5.2.2"
 publication = "0.0.3"
 pydantic = ">= 1.9.2"
 pydantic-core = "^2.18.2"
-pydash = "7.0.6"
-python-dotenv = "1.0.1"
-pytz = "2025.1"
-pyyaml = "6.0.2"
-requests = "2.32.3"
-tomli = "2.0.2"
+pydash = ">=7.0.0,<8.0.0"
+python-dotenv = ">=1.0.0,<2.0.0"
+pytz = ">=2022.0,<2026.0"
+pyyaml = ">=6.0.0,<7.0.0"
+requests = ">=2.31.0,<3.0.0"
+tomli = ">=2.0.0,<3.0.0"
 typing_extensions = ">= 4.0.0"
 
 [tool.poetry.dev-dependencies]

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/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.62",
+            "X-Fern-SDK-Version": "0.14.64",
         }
         headers["X-API-KEY"] = self.api_key
         return headers

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/errors/types.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
+import logging
 from typing import Any, Dict
 
 from vellum.client.types.vellum_error import VellumError
@@ -7,6 +8,8 @@ from vellum.client.types.vellum_error_code_enum import VellumErrorCodeEnum
 from vellum.client.types.workflow_event_error import WorkflowEventError
 from vellum.client.types.workflow_execution_event_error_code import WorkflowExecutionEventErrorCode
 
+logger = logging.getLogger(__name__)
+
 
 class WorkflowErrorCode(Enum):
     INVALID_WORKFLOW = "INVALID_WORKFLOW"
@@ -43,9 +46,15 @@ _VELLUM_ERROR_CODE_TO_WORKFLOW_ERROR_CODE: Dict[VellumErrorCodeEnum, WorkflowErr
 
 
 def vellum_error_to_workflow_error(error: VellumError) -> WorkflowError:
+    workflow_error_code = _VELLUM_ERROR_CODE_TO_WORKFLOW_ERROR_CODE.get(error.code)
+
+    if workflow_error_code is None:
+        logger.warning(f"Unknown Vellum error code '{error.code}' encountered. Falling back to INTERNAL_ERROR.")
+        workflow_error_code = WorkflowErrorCode.INTERNAL_ERROR
+
     return WorkflowError(
         message=error.message,
-        code=_VELLUM_ERROR_CODE_TO_WORKFLOW_ERROR_CODE.get(error.code, WorkflowErrorCode.INTERNAL_ERROR),
+        code=workflow_error_code,
     )
 
 
@@ -57,6 +66,7 @@ _WORKFLOW_EVENT_ERROR_CODE_TO_WORKFLOW_ERROR_CODE: Dict[WorkflowExecutionEventEr
     "NODE_EXECUTION": WorkflowErrorCode.NODE_EXECUTION,
     "LLM_PROVIDER": WorkflowErrorCode.PROVIDER_ERROR,
     "INVALID_CODE": WorkflowErrorCode.INVALID_CODE,
+    "INVALID_INPUTS": WorkflowErrorCode.INVALID_INPUTS,
     "INVALID_TEMPLATE": WorkflowErrorCode.INVALID_TEMPLATE,
     "USER_DEFINED_ERROR": WorkflowErrorCode.USER_DEFINED_ERROR,
 }

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/bases/base.py

@@ -1,3 +1,4 @@
+from abc import ABC, ABCMeta
 from dataclasses import field
 from functools import cached_property, reduce
 import inspect
@@ -50,7 +51,7 @@ def _is_annotated(cls: Type, name: str) -> bool:
     return False
 
 
-class BaseNodeMeta(type):
+class BaseNodeMeta(ABCMeta):
     def __new__(mcs, name: str, bases: Tuple[Type, ...], dct: Dict[str, Any]) -> Any:
         if "Outputs" in dct:
             outputs_class = dct["Outputs"]
@@ -258,7 +259,7 @@ class _BaseNodeExecutionMeta(type):
 NodeRunResponse = Union[BaseOutputs, Iterator[BaseOutput]]
 
 
-class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
+class BaseNode(Generic[StateType], ABC, metaclass=BaseNodeMeta):
     __id__: UUID = uuid4_from_hash(__qualname__)
     __output_ids__: Dict[str, UUID] = {}
     state: StateType
@@ -490,7 +491,6 @@ class BaseNode(Generic[StateType], metaclass=BaseNodeMeta):
                         node_output_descriptor.name,
                         node_output_descriptor.instance.resolve(self.state),
                     )
-            delattr(self.Outputs, "_outputs_post_init")
 
         setattr(self.Outputs, "_outputs_post_init", _outputs_post_init)
 

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/bases/base_adornment_node.py

@@ -1,3 +1,4 @@
+from abc import ABC
 from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Tuple, Type
 
 from vellum.workflows.inputs.base import BaseInputs
@@ -67,6 +68,7 @@ class _BaseAdornmentNodeMeta(BaseNodeMeta):
 class BaseAdornmentNode(
     BaseNode[StateType],
     Generic[StateType],
+    ABC,
     metaclass=_BaseAdornmentNodeMeta,
 ):
     """

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/bases/inline_prompt_node/node.py

@@ -193,6 +193,10 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
             return input_variables, input_values
 
         for input_name, input_value in self.prompt_inputs.items():
+            # Exclude inputs that resolved to be null. This ensure that we don't pass input values
+            # to optional prompt inputs whose values were unresolved.
+            if input_value is None:
+                continue
             if isinstance(input_value, str):
                 input_variables.append(
                     VellumVariable(

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/bases/prompt_deployment_node.py

@@ -171,6 +171,10 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
             return compiled_inputs
 
         for input_name, input_value in self.prompt_inputs.items():
+            # Exclude inputs that resolved to be null. This ensure that we don't pass input values
+            # to optional prompt inputs whose values were unresolved.
+            if input_value is None:
+                continue
             if isinstance(input_value, str):
                 compiled_inputs.append(
                     StringInputRequest(

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/code_execution_node/node.py

@@ -27,6 +27,7 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.bases.base import BaseNodeMeta
+from vellum.workflows.nodes.displayable.bases.utils import primitive_to_vellum_value
 from vellum.workflows.nodes.displayable.code_execution_node.utils import read_file_from_path, run_code_inline
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, VellumSecret
@@ -169,10 +170,13 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
                         )
                     )
                 else:
+                    # Convert primitive values to VellumValue objects
+                    vellum_values: List[VellumValue] = [primitive_to_vellum_value(item) for item in input_value]
+
                     compiled_inputs.append(
                         ArrayInput(
                             name=input_name,
-                            value=cast(List[VellumValue], input_value),
+                            value=vellum_values,
                         )
                     )
             elif isinstance(input_value, dict):

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/code_execution_node/tests/test_node.py

@@ -5,7 +5,7 @@ from typing import Any, List, Union
 
 from pydantic import BaseModel
 
-from vellum import CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
+from vellum import ArrayInput, CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
 from vellum.client.errors.bad_request_error import BadRequestError
 from vellum.client.types.chat_message import ChatMessage
 from vellum.client.types.code_execution_package import CodeExecutionPackage
@@ -1099,7 +1099,7 @@ def test_run_node__iter_dict():
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, list[str]]):
         code = """\
 def main(input_dict: dict) -> list[str]:
-    return [item.value for item in input_dict]
+    return [item.value for item in input_dict.values()]
 """
         runtime = "PYTHON_3_11_6"
         code_inputs = {
@@ -1167,3 +1167,69 @@ def main(input: str) -> str:
 
     # THEN the node should return default function call
     assert outputs == {"result": FunctionCall(name="", arguments={}), "log": ""}
+
+
+def test_run_node__array_input_with_primitive_values(vellum_client):
+    # GIVEN a node that subclasses CodeExecutionNode that processes an array of primitive values
+    class State(BaseState):
+        pass
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[State, str]):
+        code = """\
+from typing import List
+def main(numbers: List[str]) -> str:
+    return " ".join(numbers)
+"""
+        runtime = "PYTHON_3_11_6"
+        packages = [
+            CodeExecutionPackage(
+                name="openai",
+                version="1.0.0",
+            )
+        ]
+
+        code_inputs = {
+            "numbers": ["6", "4", "2"],
+        }
+
+    # AND we know what the Code Execution Node will respond with
+    mock_code_execution = CodeExecutorResponse(
+        log="",
+        output=StringVellumValue(value="6 4 2"),
+    )
+    vellum_client.execute_code.return_value = mock_code_execution
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(state=State())
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {"result": "6 4 2", "log": ""}
+
+    # AND we should have invoked the Code with the expected inputs
+    vellum_client.execute_code.assert_called_once_with(
+        input_values=[
+            ArrayInput(
+                name="numbers",
+                value=[
+                    StringVellumValue(value="6"),
+                    StringVellumValue(value="4"),
+                    StringVellumValue(value="2"),
+                ],
+            )
+        ],
+        code="""\
+from typing import List
+def main(numbers: List[str]) -> str:
+    return " ".join(numbers)
+""",
+        runtime="PYTHON_3_11_6",
+        output_type="STRING",
+        packages=[
+            CodeExecutionPackage(
+                name="openai",
+                version="1.0.0",
+            )
+        ],
+        request_options=None,
+    )

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -63,6 +63,10 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
         compiled_inputs: List[WorkflowRequestInputRequest] = []
 
         for input_name, input_value in self.subworkflow_inputs.items():
+            # Exclude inputs that resolved to be null. This ensure that we don't pass input values
+            # to optional subworkflow inputs whose values were unresolved.
+            if input_value is None:
+                continue
             if isinstance(input_value, str):
                 compiled_inputs.append(
                     WorkflowRequestStringInputRequest(

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/displayable/subworkflow_deployment_node/tests/test_node.py

@@ -468,3 +468,28 @@ def test_prompt_deployment_node__parent_context_serialization(mock_httpx_transpo
     request_execution_context = json.loads(call_request_args.read().decode("utf-8"))["execution_context"]
     assert request_execution_context["trace_id"] == str(trace_id)
     assert request_execution_context["parent_context"]
+
+
+def test_run_workflow__missing_required_input(vellum_client):
+    """Confirm that we get an error when a required input is not provided to a Subworkflow Deployment Node"""
+
+    # GIVEN a Subworkflow Deployment Node missing a required input
+    class ExampleSubworkflowDeploymentNode(SubworkflowDeploymentNode):
+        deployment = "example_subworkflow_deployment"
+        subworkflow_inputs = {}
+
+    # AND the Subworkflow Deployment API call fails due to missing required input
+    vellum_client.execute_workflow_stream.side_effect = ApiError(
+        status_code=400, body={"detail": "Missing required input for 'my_var_1'"}
+    )
+
+    # WHEN we run the node
+    node = ExampleSubworkflowDeploymentNode()
+
+    # THEN we should get a NodeException for invalid inputs
+    with pytest.raises(NodeException) as exc_info:
+        list(node.run())
+
+    # AND the error should indicate the missing required input
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert exc_info.value.message == "Missing required input for 'my_var_1'"

{vellum_ai-0.14.62 → vellum_ai-0.14.64}/src/vellum/workflows/nodes/experimental/tool_calling_node/node.py

@@ -1,9 +1,11 @@
-from collections.abc import Callable
-from typing import Any, ClassVar, List, Optional
+from collections.abc import Callable, Sequence
+from typing import Any, ClassVar, Dict, List, Optional, cast
 
 from pydash import snake_case
 
 from vellum import ChatMessage, PromptBlock
+from vellum.client.types.code_execution_package import CodeExecutionPackage
+from vellum.client.types.code_execution_runtime import CodeExecutionRuntime
 from vellum.workflows.context import execution_context, get_parent_context
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
@@ -28,14 +30,14 @@ class ToolCallingNode(BaseNode):
         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
+        function_configs: Optional[Dict[str, Dict[str, Any]]] - Mapping of function names to their configuration
     """
 
     ml_model: ClassVar[str] = "gpt-4o-mini"
     blocks: ClassVar[List[PromptBlock]] = []
     functions: ClassVar[List[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
+    function_configs: ClassVar[Optional[Dict[str, Dict[str, Any]]]] = None
 
     class Outputs(BaseOutputs):
         """
@@ -102,13 +104,28 @@ class ToolCallingNode(BaseNode):
             prompt_inputs=self.prompt_inputs,
         )
 
-        self._function_nodes = {
-            snake_case(function.__name__): create_function_node(
+        self._function_nodes = {}
+        for function in self.functions:
+            function_name = snake_case(function.__name__)
+
+            # Get configuration for this function
+            config = {}
+            if self.function_configs and function.__name__ in self.function_configs:
+                config = self.function_configs[function.__name__]
+
+            packages = config.get("packages", None)
+            if packages is not None:
+                packages = cast(Sequence[CodeExecutionPackage], packages)
+
+            runtime_raw = config.get("runtime", "PYTHON_3_11_6")
+            runtime = cast(CodeExecutionRuntime, runtime_raw)
+
+            self._function_nodes[function_name] = create_function_node(
                 function=function,
                 tool_router_node=self.tool_router_node,
+                packages=packages,
+                runtime=runtime,
             )
-            for function in self.functions
-        }
 
         graph_set = set()