vellum-ai 1.11.2__py3-none-any.whl → 1.13.5__py3-none-any.whl

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

@@ -69,23 +69,24 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
 
     def run(self) -> Iterator[BaseOutput]:
         mapped_items: Dict[str, List] = defaultdict(list)
+        items = self.items or []
         for output_descripter in self.subworkflow.Outputs:
-            mapped_items[output_descripter.name] = [None] * len(self.items)
+            mapped_items[output_descripter.name] = [None] * len(items)
 
-        if not self.items:
+        if not items:
             for output_name, output_list in mapped_items.items():
                 yield BaseOutput(name=output_name, value=output_list)
             return
 
         self._event_queue: Queue[Tuple[int, WorkflowEvent]] = Queue()
-        fulfilled_iterations: List[bool] = [False] * len(self.items)
+        fulfilled_iterations: List[bool] = [False] * len(items)
 
-        max_workers = self.max_concurrency if self.max_concurrency is not None else len(self.items)
+        max_workers = self.max_concurrency if self.max_concurrency is not None else len(items)
 
         with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
             futures = []
             current_execution_context = get_execution_context()
-            for index, item in enumerate(self.items):
+            for index, item in enumerate(items):
                 future = executor.submit(
                     self._context_run_subworkflow,
                     item=item,
@@ -180,6 +181,7 @@ class MapNode(BaseAdornmentNode[StateType], Generic[StateType, MapNodeItemType])
             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,
+            event_max_size=self._context.event_max_size,
         )
 
         for event in events:

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

@@ -71,6 +71,8 @@ def test_map_node__use_parallelism():
 
 
 def test_map_node__empty_list():
+    """Tests that a map node with an empty list returns an empty output."""
+
     # GIVEN a map node that is configured to use the parent's inputs and state
     @MapNode.wrap(items=[])
     class TestNode(BaseNode):
@@ -92,6 +94,37 @@ def test_map_node__empty_list():
     assert fulfilled_output == BaseOutput(name="value", value=[])
 
 
+def test_map_node__none_items():
+    """Tests that a map node with None items treats it as an empty array and fulfills successfully."""
+
+    # GIVEN a map node with items set to None
+    class TestNode(BaseNode):
+        item = MapNode.SubworkflowInputs.item
+
+        class Outputs(BaseOutputs):
+            value: str
+
+        def run(self) -> Outputs:
+            return self.Outputs(value=str(self.item))
+
+    class TestSubworkflow(BaseWorkflow[MapNode.SubworkflowInputs, BaseState]):
+        graph = TestNode
+
+        class Outputs(BaseOutputs):
+            value = TestNode.Outputs.value
+
+    class TestMapNode(MapNode):
+        subworkflow = TestSubworkflow
+
+    # WHEN the node is run
+    node = TestMapNode()
+    outputs = list(node.run())
+
+    # THEN the node should return an empty output without error
+    fulfilled_output = outputs[-1]
+    assert fulfilled_output == BaseOutput(name="value", value=[])
+
+
 def test_map_node__inner_try():
     # GIVEN a try wrapped node
     @TryNode.wrap()

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

@@ -52,6 +52,7 @@ class RetryNode(BaseAdornmentNode[StateType], Generic[StateType]):
                     inputs=inputs_class(attempt_number=attempt_number),
                     event_filter=all_workflow_event_filter,
                     node_output_mocks=self._context._get_all_node_output_mocks(),
+                    event_max_size=self._context.event_max_size,
                 )
 
             node_outputs: Optional[BaseNode.Outputs] = None

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

@@ -36,6 +36,7 @@ class TryNode(BaseAdornmentNode[StateType], Generic[StateType]):
             subworkflow_stream = subworkflow.stream(
                 event_filter=all_workflow_event_filter,
                 node_output_mocks=self._context._get_all_node_output_mocks(),
+                event_max_size=self._context.event_max_size,
             )
 
         outputs: Optional[BaseOutputs] = None

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

@@ -1,11 +1,12 @@
-from typing import Optional, Union
+from typing import Generic, Optional, Union
 
 from vellum.workflows.constants import AuthorizationType
 from vellum.workflows.nodes.displayable.bases.api_node import BaseAPINode
 from vellum.workflows.types.core import MergeBehavior, VellumSecret
+from vellum.workflows.types.generics import StateType
 
 
-class APINode(BaseAPINode):
+class APINode(BaseAPINode[StateType], Generic[StateType]):
     """
     Used to execute an API call. This node exists to be backwards compatible with Vellum's API Node, and for most cases,
     you should extend from `BaseAPINode` directly.

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

@@ -269,3 +269,41 @@ def test_api_node_passes_timeout_to_vellum_client(vellum_client):
     # AND the call should include RequestOptions with the correct timeout
     assert request_options is not None
     assert request_options["timeout_in_seconds"] == 25
+
+
+def test_api_node_supports_state_type_parameter(vellum_client):
+    """Test that APINode supports StateType generic parameter."""
+
+    # GIVEN a state class
+    from vellum.workflows.state.base import BaseState
+
+    class TestState(BaseState):
+        counter: int = 0
+
+    # AND an API node that uses the StateType parameter
+    class TestAPINode(APINode[TestState]):
+        method = APIRequestMethod.GET
+        authorization_type = AuthorizationType.BEARER_TOKEN
+        url = "https://example.com"
+        bearer_token_value = VellumSecret(name="secret")
+
+    # AND a mock response from the API
+    vellum_client.execute_api.return_value = ExecuteApiResponse(
+        status_code=200,
+        text='{"result": "success"}',
+        json_={"result": "success"},
+        headers={"content-type": "application/json"},
+    )
+
+    # WHEN we create and run the node with a state instance
+    state = TestState()
+    node = TestAPINode(state=state)
+    outputs = node.run()
+
+    # THEN the node should run successfully
+    assert vellum_client.execute_api.call_count == 1
+    assert outputs.status_code == 200
+    assert outputs.json == {"result": "success"}
+
+    # AND the state should be accessible (though not modified by the API node)
+    assert state.counter == 0

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

@@ -40,7 +40,7 @@ class BaseAPINode(BaseNode, Generic[StateType]):
     timeout: Optional[int] = None
 
     class Outputs(BaseOutputs):
-        json: Optional[Json]
+        json: Json
         headers: Dict[str, str]
         status_code: int
         text: str

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

@@ -85,6 +85,8 @@ class BasePromptNode(BaseNode[StateType], Generic[StateType]):
                 continue
             elif event.state == "STREAMING":
                 yield BaseOutput(name="results", delta=event.output.value)
+                if event.output.type == "STRING":
+                    yield BaseOutput(name="text", delta=event.output.value)
             elif event.state == "FULFILLED":
                 outputs = event.outputs
                 yield BaseOutput(name="results", value=event.outputs)
@@ -123,12 +125,24 @@ class BasePromptNode(BaseNode[StateType], Generic[StateType]):
         event: NodeExecutionStreamingEvent,
         workflow_output_descriptor: OutputReference,
     ) -> bool:
+        # Check if workflow output directly references this node's text output
+        text_output = getattr(event.node_definition.Outputs, "text", None)
+
+        if (
+            text_output is not None
+            and event.output.name == "text"
+            and isinstance(workflow_output_descriptor.instance, BaseDescriptor)
+            and _contains_reference_to_output(workflow_output_descriptor.instance, text_output)
+        ):
+            return True
+
         if event.output.name != "results":
             return False
 
         if not isinstance(event.output.delta, str) and not event.output.is_initiated:
             return False
 
+        # Check if workflow output references this node's text output through a FinalOutputNode
         target_nodes = [e.to_node for port in self.Ports for e in port.edges if e.to_node.__simulates_workflow_output__]
         target_node_output = next(
             (
@@ -145,4 +159,7 @@ class BasePromptNode(BaseNode[StateType], Generic[StateType]):
         if not isinstance(target_node_output.instance, BaseDescriptor):
             return False
 
-        return _contains_reference_to_output(target_node_output.instance, event.node_definition.Outputs.text)
+        if text_output is None:
+            return False
+
+        return _contains_reference_to_output(target_node_output.instance, text_output)

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

@@ -1,9 +1,23 @@
 from itertools import chain
 import json
 from uuid import uuid4
-from typing import Callable, ClassVar, Generator, Generic, Iterator, List, Optional, Set, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    ClassVar,
+    Generator,
+    Generic,
+    Iterator,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
 
 import httpx
+import jsonschema
 
 from vellum import (
     AdHocExecutePromptEvent,
@@ -53,6 +67,7 @@ from vellum.workflows.types.definition import (
     ComposioToolDefinition,
     DeploymentDefinition,
     MCPServer,
+    MCPToolDefinition,
     VellumIntegrationToolDefinition,
 )
 from vellum.workflows.types.generics import StateType, is_workflow_class
@@ -67,6 +82,60 @@ from vellum.workflows.utils.functions import (
 )
 from vellum.workflows.utils.pydantic_schema import normalize_json
 
+if TYPE_CHECKING:
+    from vellum.workflows.workflows.base import BaseWorkflow
+
+
+def _get_json_schema_to_validate(parameters_ref: object) -> Optional[dict]:
+    """
+    Extracts the JSON schema to validate from a parameters reference.
+
+    Also normalizes Pydantic models to JSON schema dicts so they can be validated.
+
+    Args:
+        parameters_ref: The parameters reference (NodeReference wrapping PromptParameters)
+
+    Returns:
+        The JSON schema dict to validate, or None if no schema found.
+    """
+    parameters_instance = getattr(parameters_ref, "instance", None)
+    if not parameters_instance:
+        return None
+
+    custom_params = getattr(parameters_instance, "custom_parameters", None)
+    if not isinstance(custom_params, dict):
+        return None
+
+    json_schema = custom_params.get("json_schema")
+    if json_schema is None:
+        return None
+
+    # Normalize Pydantic models to JSON schema dicts before validation
+    # This handles cases like {"name": "...", "schema": SomePydanticModel}
+    json_schema = normalize_json(json_schema)
+
+    # After normalization, we expect a dict; anything else we ignore
+    if not isinstance(json_schema, dict):
+        return None
+
+    return json_schema
+
+
+def _validate_json_schema_structure(schema: dict) -> None:
+    """
+    Validates the structure of a JSON schema using the jsonschema library.
+
+    This uses the JSON Schema meta-schema to validate that the provided schema
+    is a valid JSON Schema according to the specification.
+
+    Args:
+        schema: The JSON schema dictionary to validate
+
+    Raises:
+        jsonschema.exceptions.SchemaError: If the schema structure is invalid
+    """
+    jsonschema.Draft7Validator.check_schema(schema)
+
 
 class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
     """
@@ -87,7 +156,19 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
     blocks: ClassVar[List[PromptBlock]]
 
     # The functions/tools that a Prompt has access to
-    functions: Optional[List[Union[FunctionDefinition, Callable]]] = None
+    functions: Optional[
+        List[
+            Union[
+                FunctionDefinition,
+                Callable,
+                DeploymentDefinition,
+                Type["BaseWorkflow"],
+                VellumIntegrationToolDefinition,
+                MCPServer,
+                MCPToolDefinition,
+            ]
+        ]
+    ] = None
 
     parameters: PromptParameters = DEFAULT_PROMPT_PARAMETERS
     expand_meta: Optional[AdHocExpandMeta] = AdHocExpandMeta(finish_reason=True)
@@ -159,6 +240,14 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                     normalized_functions.append(
                         compile_vellum_integration_tool_definition(function, self._context.vellum_client)
                     )
+                elif isinstance(function, MCPToolDefinition):
+                    normalized_functions.append(
+                        FunctionDefinition(
+                            name=get_mcp_tool_name(function),
+                            description=function.description,
+                            parameters=function.parameters,
+                        )
+                    )
                 elif isinstance(function, MCPServer):
                     tool_definitions = compile_mcp_tool_definition(function)
                     for tool_def in tool_definitions:
@@ -255,6 +344,8 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
                 continue
             elif event.state == "STREAMING":
                 yield BaseOutput(name="results", delta=event.output.value)
+                if event.output.type == "STRING":
+                    yield BaseOutput(name="text", delta=event.output.value)
             elif event.state == "FULFILLED":
                 if event.meta and event.meta.finish_reason == "LENGTH":
                     text_value, json_value = process_additional_prompt_outputs(event.outputs)
@@ -465,3 +556,19 @@ class BaseInlinePromptNode(BasePromptNode[StateType], Generic[StateType]):
         Override this method to process the blocks before they are executed.
         """
         return blocks
+
+    @classmethod
+    def __validate__(cls) -> None:
+        """
+        Validates the node configuration, including JSON schema structure in parameters.
+
+        Raises:
+            jsonschema.exceptions.SchemaError: If the JSON schema structure is invalid
+        """
+        parameters_ref = getattr(cls, "parameters", None)
+        if parameters_ref is None:
+            return
+
+        schema = _get_json_schema_to_validate(parameters_ref)
+        if schema is not None:
+            _validate_json_schema_structure(schema)

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

@@ -1,3 +1,4 @@
+from itertools import chain
 import json
 from uuid import UUID
 from typing import Any, ClassVar, Dict, Generator, Generic, Iterator, List, Optional, Sequence, Set, Union
@@ -113,12 +114,18 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
         if prompt_event_stream is None:
             try:
                 prompt_event_stream = self._get_prompt_event_stream()
-                next(prompt_event_stream)
+                first_event = next(prompt_event_stream)
             except ApiError as e:
                 if e.status_code and e.status_code < 500 and self.ml_model_fallbacks is not None:
                     prompt_event_stream = self._retry_prompt_stream_with_fallbacks(tried_fallbacks)
                 else:
                     self._handle_api_error(e)
+            else:
+                if first_event.state == "REJECTED":
+                    workflow_error = vellum_error_to_workflow_error(first_event.error)
+                    raise NodeException.of(workflow_error)
+                if first_event.state != "INITIATED":
+                    prompt_event_stream = chain([first_event], prompt_event_stream)
 
         outputs: Optional[List[PromptOutput]] = None
         if prompt_event_stream is not None:
@@ -127,6 +134,8 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                     continue
                 elif event.state == "STREAMING":
                     yield BaseOutput(name="results", delta=event.output.value)
+                    if event.output.type == "STRING":
+                        yield BaseOutput(name="text", delta=event.output.value)
                 elif event.state == "FULFILLED":
                     outputs = event.outputs
                     yield BaseOutput(name="results", value=event.outputs)
@@ -159,7 +168,9 @@ class BasePromptDeploymentNode(BasePromptNode, Generic[StateType]):
                 try:
                     tried_fallbacks.add(ml_model_fallback)
                     prompt_event_stream = self._get_prompt_event_stream(ml_model_fallback=ml_model_fallback)
-                    next(prompt_event_stream)
+                    first_event = next(prompt_event_stream)
+                    if first_event.state != "INITIATED":
+                        prompt_event_stream = chain([first_event], prompt_event_stream)
                     return prompt_event_stream
                 except ApiError:
                     continue

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

@@ -104,7 +104,7 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
         output_type = self.__class__.get_output_type()
         code, filepath = self._resolve_code()
         if not self.packages and self.runtime == "PYTHON_3_11_6" and not self._has_secrets_in_code_inputs():
-            logs, result = run_code_inline(code, self.code_inputs, output_type, filepath)
+            logs, result = run_code_inline(code, self.code_inputs, output_type, filepath, self._context.vellum_client)
             return self.Outputs(result=result, log=logs)
 
         else:
@@ -133,22 +133,16 @@ class CodeExecutionNode(BaseNode[StateType], Generic[StateType, _OutputType], me
             return self.Outputs(result=code_execution_result.output.value, log=code_execution_result.log)
 
     def _handle_api_error(self, e: ApiError) -> None:
-        if e.status_code and e.status_code == 403 and isinstance(e.body, dict):
-            raise NodeException(
-                message=e.body.get("detail", "Provider credentials is missing or unavailable"),
-                code=WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE,
-            ) from e
+        body = e.body if isinstance(e.body, dict) else {}
+        message = body.get("detail") or body.get("message") or "Failed to execute code"
 
-        if e.status_code and e.status_code >= 400 and e.status_code < 500 and isinstance(e.body, dict):
-            raise NodeException(
-                message=e.body.get("message", e.body.get("detail", "Failed to execute code")),
-                code=WorkflowErrorCode.INVALID_INPUTS,
-            ) from e
+        if e.status_code == 400:
+            raise NodeException(message=message, code=WorkflowErrorCode.INVALID_INPUTS) from e
+
+        if e.status_code and e.status_code >= 500:
+            raise NodeException(message=message, code=WorkflowErrorCode.INTERNAL_ERROR) from e
 
-        raise NodeException(
-            message="Failed to execute code",
-            code=WorkflowErrorCode.INTERNAL_ERROR,
-        ) from e
+        raise NodeException(message=message, code=WorkflowErrorCode.NODE_EXECUTION) from e
 
     def _has_secrets_in_code_inputs(self) -> bool:
         """Check if any code_inputs contain VellumSecret instances that require API execution."""

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

@@ -1,12 +1,12 @@
 import pytest
+from datetime import datetime
 import os
 import re
 from typing import Any, List, Union
 
 from pydantic import BaseModel
 
-from vellum import ArrayInput, CodeExecutorResponse, NumberVellumValue, StringInput, StringVellumValue
-from vellum.client.core.api_error import ApiError
+from vellum import ArrayInput, CodeExecutorResponse, MlModelRead, NumberVellumValue, StringInput, StringVellumValue
 from vellum.client.errors.bad_request_error import BadRequestError
 from vellum.client.errors.forbidden_error import ForbiddenError
 from vellum.client.errors.internal_server_error import InternalServerError
@@ -24,6 +24,7 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.displayable.code_execution_node import CodeExecutionNode
 from vellum.workflows.references.vellum_secret import VellumSecretReference
 from vellum.workflows.state.base import BaseState, StateMeta
+from vellum.workflows.state.context import WorkflowContext
 from vellum.workflows.types.core import Json
 
 
@@ -551,6 +552,51 @@ def main(word: str) -> dict:
     }
 
 
+def test_run_node__run_inline__vellum_client(vellum_client):
+    """Confirm that CodeExecutionNodes can convert a dict to a Pydantic model during inline execution."""
+
+    # GIVEN a node that subclasses CodeExecutionNode that returns a dict matching Any
+    vellum_client.ml_models.retrieve.return_value = MlModelRead(
+        id="test-ml-model-id",
+        name="Test ML Model",
+        description="Test ML Model Description",
+        introduced_on=datetime(2025, 1, 2),
+    )
+
+    class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, Any]):
+        code = """\
+def main(model_id: str) -> dict:
+    ml_model = vellum_client.ml_models.retrieve(model_id)
+    return {
+        "model_name": ml_model.name,
+        "model_id": ml_model.id,
+    }
+    """
+        runtime = "PYTHON_3_11_6"
+
+        code_inputs = {
+            "model_id": "test-ml-model-id",
+        }
+
+    # WHEN we run the node
+    node = ExampleCodeExecutionNode(context=WorkflowContext(vellum_client=vellum_client))
+    outputs = node.run()
+
+    # THEN the node should have produced the outputs we expect
+    assert outputs == {
+        "result": {
+            "model_name": "Test ML Model",
+            "model_id": "test-ml-model-id",
+        },
+        "log": "",
+    }
+
+    # AND
+    vellum_client.ml_models.retrieve.assert_called_once_with(
+        "test-ml-model-id",
+    )
+
+
 def test_run_node__array_input_with_vellum_values(vellum_client):
     """Confirm that CodeExecutionNodes can handle arrays containing VellumValue objects."""
 
@@ -796,34 +842,34 @@ Node.js v21.7.3
     assert exc_info.value.message == message
 
 
-def test_run_node__execute_code_api_fails_403__provider_credentials_unavailable(vellum_client):
-    """
-    Tests that a 403 error from the API is handled with PROVIDER_CREDENTIALS_UNAVAILABLE error code.
-    """
+def test_run_node__execute_code_api_fails_403__node_execution(vellum_client):
+    """Tests that a 403 error from the code execution API is handled with NODE_EXECUTION error code."""
 
+    # GIVEN a code execution node
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
         code = "def main(): return 'test'"
         runtime = "PYTHON_3_11_6"
         packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
 
+    # AND the API returns a 403 error
     vellum_client.execute_code.side_effect = ForbiddenError(
         body={
-            "detail": "Provider credentials is missing or unavailable",
+            "detail": "Access denied to this resource",
         }
     )
 
+    # WHEN we run the node
     node = ExampleCodeExecutionNode()
     with pytest.raises(NodeException) as exc_info:
         node.run()
 
-    assert exc_info.value.code == WorkflowErrorCode.PROVIDER_CREDENTIALS_UNAVAILABLE
-    assert "Provider credentials is missing or unavailable" in exc_info.value.message
+    # THEN it should raise NODE_EXECUTION (not PROVIDER_CREDENTIALS_UNAVAILABLE)
+    assert exc_info.value.code == WorkflowErrorCode.NODE_EXECUTION
+    assert exc_info.value.message == "Access denied to this resource"
 
 
-def test_run_node__execute_code_api_fails_404__invalid_inputs(vellum_client):
-    """
-    Tests that a 404 error from the API is handled with INVALID_INPUTS error code.
-    """
+def test_run_node__execute_code_api_fails_404__node_execution(vellum_client):
+    """Tests that a 404 error from the API is handled with NODE_EXECUTION error code."""
 
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
         code = "def main(): return 'test'"
@@ -840,14 +886,12 @@ def test_run_node__execute_code_api_fails_404__invalid_inputs(vellum_client):
     with pytest.raises(NodeException) as exc_info:
         node.run()
 
-    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert exc_info.value.code == WorkflowErrorCode.NODE_EXECUTION
     assert "Resource not found" in exc_info.value.message
 
 
 def test_run_node__execute_code_api_fails_500__internal_error(vellum_client):
-    """
-    Tests that a 500 error from the API is handled with INTERNAL_ERROR error code.
-    """
+    """Tests that a 500 error from the API is handled with INTERNAL_ERROR error code."""
 
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
         code = "def main(): return 'test'"
@@ -865,21 +909,18 @@ def test_run_node__execute_code_api_fails_500__internal_error(vellum_client):
         node.run()
 
     assert exc_info.value.code == WorkflowErrorCode.INTERNAL_ERROR
-    assert exc_info.value.message == "Failed to execute code"
+    assert exc_info.value.message == "Internal server error occurred"
 
 
-def test_run_node__execute_code_api_fails_4xx_no_message__uses_detail(vellum_client):
-    """
-    Tests that a 4xx error without a 'message' field falls back to 'detail' field.
-    """
+def test_run_node__execute_code_api_fails_400__invalid_inputs(vellum_client):
+    """Tests that a 400 error from the API is handled with INVALID_INPUTS error code."""
 
     class ExampleCodeExecutionNode(CodeExecutionNode[BaseState, str]):
         code = "def main(): return 'test'"
         runtime = "PYTHON_3_11_6"
         packages = [CodeExecutionPackage(name="requests", version="2.28.0")]
 
-    vellum_client.execute_code.side_effect = ApiError(
-        status_code=422,
+    vellum_client.execute_code.side_effect = BadRequestError(
         body={
             "detail": "Invalid request parameters",
         },

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

@@ -4,6 +4,7 @@ import sys
 import traceback
 from typing import Any, Optional, Tuple, Union
 
+from vellum import Vellum
 from vellum.workflows.constants import undefined
 from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
@@ -41,6 +42,7 @@ def run_code_inline(
     inputs: EntityInputsInterface,
     output_type: Any,
     filepath: str,
+    vellum_client: Vellum,
 ) -> Tuple[str, Any]:
     log_buffer = io.StringIO()
 
@@ -54,6 +56,7 @@ def run_code_inline(
         "__arg__inputs": wrapped_inputs,
         "__arg__out": None,
         "print": _inline_print,
+        "vellum_client": vellum_client,
     }
     run_args = [f"{name}=__arg__inputs['{name}']" for name, value in inputs.items() if value is not undefined]
     execution_code = f"""\