vellum-ai 1.7.5__tar.gz → 1.7.6__tar.gz

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.7.5
+Version: 1.7.6
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/ee/vellum_cli/push.py

@@ -77,11 +77,7 @@ def push_command(
 
     logger.info(f"Loading workflow from {workflow_config.module}")
     resolved_workspace = workspace or workflow_config.workspace or DEFAULT_WORKSPACE_CONFIG.name
-    workspace_config = (
-        next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
-        if workspace
-        else DEFAULT_WORKSPACE_CONFIG
-    )
+    workspace_config = next((w for w in config.workspaces if w.name == resolved_workspace), DEFAULT_WORKSPACE_CONFIG)
     api_key = os.getenv(workspace_config.api_key)
     if not api_key:
         raise ValueError(f"No API key value found in environment for workspace '{workspace_config.name}'.")

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/ee/vellum_cli/tests/test_push.py

@@ -1134,3 +1134,89 @@ def test_push__deploy_stores_deployment_config_in_lock_file(mock_module, vellum_
     deployment_config_str = call_args["deployment_config"]
     deployment_config = json.loads(deployment_config_str)
     assert deployment_config["name"] == "test-push-deploy-stores-deployment-config-in-lock-file"
+
+
+def test_push__custom_workspace_in_lockfile__uses_custom_workspace_without_flag(mock_module, vellum_client_class):
+    """
+    Tests that push respects the custom workspace from the lockfile when no --workspace flag is provided.
+    """
+
+    # GIVEN a workflow already configured in the lockfile with a custom workspace
+    temp_dir = mock_module.temp_dir
+    module = mock_module.module
+    workflow_sandbox_id = mock_module.workflow_sandbox_id
+    set_pyproject_toml = mock_module.set_pyproject_toml
+
+    # AND the lockfile has the workflow with a custom workspace
+    with open(os.path.join(temp_dir, "vellum.lock.json"), "w") as f:
+        json.dump(
+            {
+                "version": "1.0",
+                "workflows": [
+                    {
+                        "module": module,
+                        "workflow_sandbox_id": workflow_sandbox_id,
+                        "workspace": "my_custom_workspace",
+                        "container_image_name": None,
+                        "container_image_tag": None,
+                        "deployments": [],
+                        "ignore": None,
+                        "target_directory": None,
+                    }
+                ],
+                "workspaces": [],
+            },
+            f,
+            indent=2,
+        )
+
+    # AND the custom workspace is defined in pyproject.toml
+    set_pyproject_toml(
+        {
+            "workflows": [],
+            "workspaces": [
+                {
+                    "name": "my_custom_workspace",
+                    "api_key": "MY_CUSTOM_VELLUM_API_KEY",
+                }
+            ],
+        }
+    )
+
+    # AND the .env file has the custom api key
+    with open(os.path.join(temp_dir, ".env"), "w") as f:
+        f.write(
+            """\
+VELLUM_API_KEY=abcdef123456
+MY_CUSTOM_VELLUM_API_KEY=custom-key-xyz
+"""
+        )
+
+    # AND a workflow exists in the module
+    _ensure_workflow_py(temp_dir, module)
+
+    # AND the push API returns successfully
+    vellum_client_class.return_value.workflows.push.return_value = WorkflowPushResponse(
+        workflow_sandbox_id=workflow_sandbox_id,
+    )
+    vellum_client_class.return_value._client_wrapper._environment.default = "https://api.vellum.ai/v1"
+
+    # WHEN calling `vellum push` WITHOUT the --workspace flag
+    runner = CliRunner()
+    result = runner.invoke(cli_main, ["push", module])
+
+    # THEN it should succeed
+    assert result.exit_code == 0, result.output
+
+    # AND the custom workspace API key should have been used
+    vellum_client_class.assert_called_once_with(
+        api_key="custom-key-xyz",
+        environment=mock.ANY,
+        api_version=None,
+    )
+
+    with open(os.path.join(temp_dir, "vellum.lock.json")) as f:
+        lock_file_content = json.load(f)
+        assert len(lock_file_content["workflows"]) == 1
+        assert lock_file_content["workflows"][0]["workspace"] == "my_custom_workspace"
+        assert lock_file_content["workflows"][0]["workflow_sandbox_id"] == workflow_sandbox_id

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_manual_trigger_serialization.py

@@ -3,8 +3,6 @@
 import pytest
 from typing import cast
 
-from deepdiff import DeepDiff
-
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases.base import BaseNode
@@ -42,27 +40,24 @@ def serialize(workflow_class) -> JsonObject:
 
 
 def test_manual_trigger_serialization():
-    """Workflow with ManualTrigger serializes with trigger field."""
+    """Workflow with ManualTrigger serializes with triggers field."""
     result = serialize(create_workflow(ManualTrigger))
-    workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
-    trigger = cast(JsonObject, workflow_raw_data["trigger"])
+    triggers = cast(JsonArray, result["triggers"])
+
+    assert len(triggers) == 1
+    trigger = cast(JsonObject, triggers[0])
 
     assert trigger["type"] == "MANUAL"
-    assert not DeepDiff(
-        {
-            "type": "MANUAL",
-            "definition": {"name": "ManualTrigger", "module": ["vellum", "workflows", "triggers", "manual"]},
-        },
-        trigger,
-        ignore_order=True,
-    )
+    assert "id" in trigger
+    assert "attributes" in trigger
+    assert trigger["attributes"] == []
+    assert "definition" not in trigger
 
 
 def test_no_trigger_serialization():
-    """Workflow without trigger has no trigger field."""
+    """Workflow without trigger has no triggers field."""
     result = serialize(create_workflow())
-    workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
-    assert "trigger" not in workflow_raw_data
+    assert "triggers" not in result
 
 
 def test_manual_trigger_multiple_entrypoints():
@@ -84,10 +79,12 @@ def test_manual_trigger_multiple_entrypoints():
             output_b = NodeB.Outputs.output
 
     result = serialize(MultiWorkflow)
+    triggers = cast(JsonArray, result["triggers"])
     workflow_data = cast(JsonObject, result["workflow_raw_data"])
-    trigger = cast(JsonObject, workflow_data["trigger"])
     nodes = cast(JsonArray, workflow_data["nodes"])
 
+    assert len(triggers) == 1
+    trigger = cast(JsonObject, triggers[0])
     assert trigger["type"] == "MANUAL"
     assert len([n for n in nodes if cast(JsonObject, n)["type"] == "GENERIC"]) >= 2
 
@@ -98,8 +95,8 @@ def test_serialized_workflow_structure():
     workflow_raw_data = cast(JsonObject, result["workflow_raw_data"])
     definition = cast(JsonObject, workflow_raw_data["definition"])
 
-    assert result.keys() == {"workflow_raw_data", "input_variables", "state_variables", "output_variables"}
-    assert workflow_raw_data.keys() == {"nodes", "edges", "display_data", "definition", "output_values", "trigger"}
+    assert result.keys() == {"workflow_raw_data", "input_variables", "state_variables", "output_variables", "triggers"}
+    assert workflow_raw_data.keys() == {"nodes", "edges", "display_data", "definition", "output_values"}
     assert definition["name"] == "TestWorkflow"
 
 

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/ee/vellum_ee/workflows/display/workflows/base_workflow_display.py

@@ -410,7 +410,7 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 self.display_context.add_error(e)
 
         # Serialize workflow-level trigger if present
-        trigger_data: Optional[JsonObject] = self._serialize_workflow_trigger()
+        triggers: Optional[JsonArray] = self._serialize_workflow_trigger()
 
         workflow_raw_data: JsonObject = {
             "nodes": cast(JsonArray, nodes_dict_list),
@@ -423,22 +423,25 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
             "output_values": output_values,
         }
 
-        if trigger_data is not None:
-            workflow_raw_data["trigger"] = trigger_data
-
-        return {
+        result: JsonObject = {
             "workflow_raw_data": workflow_raw_data,
             "input_variables": input_variables,
             "state_variables": state_variables,
             "output_variables": output_variables,
         }
 
-    def _serialize_workflow_trigger(self) -> Optional[JsonObject]:
+        if triggers is not None:
+            result["triggers"] = triggers
+
+        return result
+
+    def _serialize_workflow_trigger(self) -> Optional[JsonArray]:
         """
         Serialize workflow-level trigger information.
 
         Returns:
-            JsonObject with trigger data if a trigger is present, None otherwise
+            JsonArray with trigger data if a trigger is present, None otherwise.
+            Each trigger in the array has: id (UUID), type (str), attributes (list)
         """
         # Get all trigger edges from the workflow's subgraphs
         trigger_edges = []
@@ -469,13 +472,19 @@ class BaseWorkflowDisplay(Generic[WorkflowType]):
                 f"Please add it to the trigger type mapping in get_trigger_type_mapping()."
             )
 
-        return {
-            "type": trigger_type.value,
-            "definition": {
-                "name": trigger_class.__name__,
-                "module": cast(JsonArray, trigger_class.__module__.split(".")),
-            },
-        }
+        # Return as a list with a single trigger object matching Django schema
+        trigger_id = uuid4_from_hash(f"{trigger_class.__module__} | {trigger_class.__qualname__}")
+
+        return cast(
+            JsonArray,
+            [
+                {
+                    "id": str(trigger_id),
+                    "type": trigger_type.value,
+                    "attributes": [],
+                }
+            ],
+        )
 
     def _serialize_edge_display_data(self, edge_display: EdgeDisplay) -> Optional[JsonObject]:
         """Serialize edge display data, returning None if no display data is present."""

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/ee/vellum_ee/workflows/tests/test_server.py

@@ -536,10 +536,49 @@ class BrokenNode(BaseNode)  # Missing colon
 
     # AND the error message should be user-friendly
     error_message = str(exc_info.value)
-    assert "Failed to load workflow module:" in error_message
+    assert "Syntax Error raised while loading Workflow:" in error_message
     assert "invalid syntax" in error_message or "expected ':'" in error_message
 
 
+def test_load_from_module__name_error_in_node_file():
+    """
+    Tests that a NameError in a node file raises WorkflowInitializationException with user-facing message.
+    """
+    # GIVEN a workflow module with a node file containing a NameError (undefined class reference)
+    files = {
+        "__init__.py": "",
+        "workflow.py": """\
+from vellum.workflows import BaseWorkflow
+from .nodes.broken_node import BrokenNode
+
+class Workflow(BaseWorkflow):
+    graph = BrokenNode
+""",
+        "nodes/__init__.py": "",
+        "nodes/broken_node.py": """\
+from vellum.workflows.nodes import BaseNode
+
+class BrokenNode(BaseNode):
+    some_attribute = UndefinedClass()
+""",
+    }
+
+    namespace = str(uuid4())
+
+    # AND the virtual file loader is registered
+    sys.meta_path.append(VirtualFileFinder(files, namespace))
+
+    # WHEN we attempt to load the workflow
+    # THEN it should raise WorkflowInitializationException
+    with pytest.raises(WorkflowInitializationException) as exc_info:
+        BaseWorkflow.load_from_module(namespace)
+
+    # AND the error message should be user-friendly
+    error_message = str(exc_info.value)
+    assert "Invalid variable reference:" in error_message
+    assert "UndefinedClass" in error_message or "not defined" in error_message
+
+
 def test_serialize_module__tool_calling_node_with_single_tool():
     """Test that serialize_module works with a tool calling node that has a single tool."""
 

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/pyproject.toml

@@ -3,7 +3,7 @@ name = "vellum-ai"
 
 [tool.poetry]
 name = "vellum-ai"
-version = "1.7.5"
+version = "1.7.6"
 description = ""
 readme = "README.md"
 authors = []

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/client/core/client_wrapper.py

@@ -27,10 +27,10 @@ class BaseClientWrapper:
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
-            "User-Agent": "vellum-ai/1.7.5",
+            "User-Agent": "vellum-ai/1.7.6",
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "vellum-ai",
-            "X-Fern-SDK-Version": "1.7.5",
+            "X-Fern-SDK-Version": "1.7.6",
             **(self.get_custom_headers() or {}),
         }
         if self._api_version is not None:

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/bases/base.py

@@ -1,10 +1,26 @@
 from abc import ABC, ABCMeta, abstractmethod
+from collections.abc import Callable as CollectionsCallable
 from dataclasses import field
 from functools import cached_property, reduce
 import inspect
 from types import MappingProxyType
 from uuid import UUID, uuid4
-from typing import Any, Dict, Generic, Iterator, Optional, Set, Tuple, Type, TypeVar, Union, cast, get_args
+from typing import (
+    Any,
+    Callable as TypingCallable,
+    Dict,
+    Generic,
+    Iterator,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    TypeVar,
+    Union,
+    cast,
+    get_args,
+    get_origin,
+)
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.descriptors.base import BaseDescriptor
@@ -43,15 +59,15 @@ def _is_nested_class(nested: Any, parent: Type) -> bool:
     ) or any(_is_nested_class(nested, base) for base in parent.__bases__)
 
 
-def _is_annotated(cls: Type, name: str) -> bool:
+def _is_annotated(cls: Type, name: str) -> Any:
     if name in cls.__annotations__:
-        return True
+        return cls.__annotations__[name]
 
     for base in cls.__bases__:
-        if _is_annotated(base, name):
-            return True
+        if annotation := _is_annotated(base, name):
+            return annotation
 
-    return False
+    return None
 
 
 class BaseNodeMeta(ABCMeta):
@@ -151,8 +167,10 @@ class BaseNodeMeta(ABCMeta):
         try:
             attribute = super().__getattribute__(name)
         except AttributeError as e:
-            if _is_annotated(cls, name):
-                attribute = None
+            annotation = _is_annotated(cls, name)
+            origin_annotation = get_origin(annotation)
+            if origin_annotation is not CollectionsCallable and origin_annotation is not TypingCallable:
+                attribute = undefined
             else:
                 raise e
 
@@ -482,7 +500,8 @@ class BaseNode(Generic[StateType], ABC, BaseExecutable, metaclass=BaseNodeMeta):
                     setattr(base, leaf, input_value)
 
         for descriptor in self.__class__:
-            if not descriptor.instance:
+            if descriptor.instance is undefined:
+                setattr(self, descriptor.name, undefined)
                 continue
 
             if any(isinstance(t, type) and issubclass(t, BaseDescriptor) for t in descriptor.types):

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/displayable/search_node/node.py

@@ -1,6 +1,7 @@
 import json
 from typing import ClassVar
 
+from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.nodes.displayable.bases import BaseSearchNode as BaseSearchNode
@@ -37,7 +38,7 @@ class SearchNode(BaseSearchNode[StateType]):
         text: str
 
     def run(self) -> Outputs:
-        if self.query is None or self.query == "":
+        if self.query is undefined or self.query is None or self.query == "":
             raise NodeException(
                 message="Search query is required but was not provided",
                 code=WorkflowErrorCode.INVALID_INPUTS,

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/displayable/search_node/tests/test_node.py

@@ -234,3 +234,17 @@ def test_run_workflow__invalid_query_raises_validation_error(invalid_query):
     assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
     assert "query" in exc_info.value.message.lower()
     assert "required" in exc_info.value.message.lower() or "missing" in exc_info.value.message.lower()
+
+
+def test_run_workflow__missing_query_attribute_raises_validation_error():
+    """Confirm that a SearchNode without a query attribute defined raises INVALID_INPUTS"""
+
+    class MySearchNode(SearchNode):
+        document_index = "document_index"
+
+    with pytest.raises(NodeException) as exc_info:
+        MySearchNode().run()
+
+    assert exc_info.value.code == WorkflowErrorCode.INVALID_INPUTS
+    assert "query" in exc_info.value.message.lower()
+    assert "required" in exc_info.value.message.lower()

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/displayable/subworkflow_deployment_node/node.py

@@ -15,7 +15,7 @@ from vellum import (
 from vellum.client.core import RequestOptions
 from vellum.client.core.api_error import ApiError
 from vellum.client.types.chat_message_request import ChatMessageRequest
-from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT
+from vellum.workflows.constants import LATEST_RELEASE_TAG, OMIT, undefined
 from vellum.workflows.context import execution_context, get_execution_context, get_parent_context
 from vellum.workflows.errors import WorkflowErrorCode
 from vellum.workflows.errors.types import workflow_event_error_to_workflow_error
@@ -226,6 +226,12 @@ class SubworkflowDeploymentNode(BaseNode[StateType], Generic[StateType]):
             **request_options.get("additional_body_parameters", {}),
         }
 
+        if self.deployment is undefined:
+            raise NodeException(
+                code=WorkflowErrorCode.NODE_EXECUTION,
+                message="Expected subworkflow deployment attribute to be either a UUID or STR, got `undefined` instead",
+            )
+
         try:
             deployment_id = str(self.deployment) if isinstance(self.deployment, UUID) else None
             deployment_name = self.deployment if isinstance(self.deployment, str) else None

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

@@ -265,7 +265,7 @@ def test_run_workflow__no_deployment():
 
     # AND the error message should be correct
     assert exc_info.value.code == WorkflowErrorCode.NODE_EXECUTION
-    assert "Expected subworkflow deployment attribute to be either a UUID or STR, got None instead" in str(
+    assert "Expected subworkflow deployment attribute to be either a UUID or STR, got `undefined` instead" in str(
         exc_info.value
     )
 

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/displayable/tool_calling_node/tests/test_node.py

@@ -1,3 +1,4 @@
+import pytest
 import json
 from uuid import uuid4
 from typing import Any, Iterator, List
@@ -14,6 +15,8 @@ from vellum.client.types.string_vellum_value import StringVellumValue
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
 from vellum.prompts.constants import DEFAULT_PROMPT_PARAMETERS
 from vellum.workflows import BaseWorkflow
+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.tool_calling_node.node import ToolCallingNode
@@ -372,3 +375,54 @@ def test_tool_calling_node_workflow_is_dynamic(vellum_adhoc_prompt_client):
     assert initiated_events[0].body.workflow_definition.is_dynamic is False  # Main workflow
     assert initiated_events[1].body.workflow_definition.is_dynamic is True  # Tool calling internal
     assert initiated_events[2].body.workflow_definition.is_dynamic is True  # Inline workflow
+
+
+def test_tool_node_preserves_node_exception():
+    """Test that tool nodes preserve NodeException error codes and raw_data."""
+
+    def failing_function() -> str:
+        raise NodeException(
+            message="Custom error",
+            code=WorkflowErrorCode.INVALID_INPUTS,
+            raw_data={"key": "value"},
+        )
+
+    tool_prompt_node = create_tool_prompt_node(
+        ml_model="test-model",
+        blocks=[],
+        functions=[failing_function],
+        prompt_inputs=None,
+        parameters=DEFAULT_PROMPT_PARAMETERS,
+    )
+
+    function_node_class = create_function_node(
+        function=failing_function,
+        tool_prompt_node=tool_prompt_node,
+    )
+
+    state = ToolCallingState(
+        meta=StateMeta(
+            node_outputs={
+                tool_prompt_node.Outputs.results: [
+                    FunctionCallVellumValue(
+                        value=FunctionCall(
+                            arguments={},
+                            id="call_123",
+                            name="failing_function",
+                            state="FULFILLED",
+                        ),
+                    )
+                ],
+            },
+        )
+    )
+
+    function_node = function_node_class(state=state)
+
+    with pytest.raises(NodeException) as exc_info:
+        list(function_node.run())
+
+    e = exc_info.value
+    assert e.code == WorkflowErrorCode.INVALID_INPUTS
+    assert e.raw_data == {"key": "value"}
+    assert "Custom error" in e.message

{vellum_ai-1.7.5 → vellum_ai-1.7.6}/src/vellum/workflows/nodes/displayable/tool_calling_node/utils.py

@@ -55,6 +55,28 @@ class FunctionCallNodeMixin:
 
     function_call_output: List[PromptOutput]
 
+    def _handle_tool_exception(self, e: Exception, tool_type: str, tool_name: str) -> None:
+        """
+        Re-raise exceptions with contextual information while preserving NodeException details.
+
+        Args:
+            e: The caught exception
+            tool_type: Type of tool (e.g., "function", "MCP tool", "Vellum Integration tool")
+            tool_name: Name of the tool that failed
+        """
+        if isinstance(e, NodeException):
+            # Preserve original error code and raw_data while adding context
+            raise NodeException(
+                message=f"Error executing {tool_type} '{tool_name}': {e.message}",
+                code=e.code,
+                raw_data=e.raw_data,
+            ) from e
+        else:
+            raise NodeException(
+                message=f"Error executing {tool_type} '{tool_name}': {str(e)}",
+                code=WorkflowErrorCode.NODE_EXECUTION,
+            ) from e
+
     def _extract_function_arguments(self) -> dict:
         """Extract arguments from function call output."""
         current_index = getattr(self, "state").current_prompt_output_index
@@ -201,11 +223,7 @@ class FunctionNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
         try:
             result = self.function_definition(**arguments)
         except Exception as e:
-            function_name = self.function_definition.__name__
-            raise NodeException(
-                message=f"Error executing function '{function_name}': {str(e)}",
-                code=WorkflowErrorCode.NODE_EXECUTION,
-            ) from e
+            self._handle_tool_exception(e, "function", self.function_definition.__name__)
 
         # Add the result to the chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -232,10 +250,7 @@ class ComposioNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
             else:
                 result = composio_service.execute_tool(tool_name=self.composio_tool.action, arguments=arguments)
         except Exception as e:
-            raise NodeException(
-                message=f"Error executing Composio tool '{self.composio_tool.action}': {str(e)}",
-                code=WorkflowErrorCode.NODE_EXECUTION,
-            ) from e
+            self._handle_tool_exception(e, "Composio tool", self.composio_tool.action)
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -255,10 +270,7 @@ class MCPNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
             mcp_service = MCPService()
             result = mcp_service.execute_tool(tool_def=self.mcp_tool, arguments=arguments)
         except Exception as e:
-            raise NodeException(
-                message=f"Error executing MCP tool '{self.mcp_tool.name}': {str(e)}",
-                code=WorkflowErrorCode.NODE_EXECUTION,
-            ) from e
+            self._handle_tool_exception(e, "MCP tool", self.mcp_tool.name)
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)
@@ -283,18 +295,8 @@ class VellumIntegrationNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
                 tool_name=self.vellum_integration_tool.name,
                 arguments=arguments,
             )
-        except NodeException as e:
-            # Preserve original error code and raw_data while adding context
-            raise NodeException(
-                message=f"Error executing Vellum Integration tool '{self.vellum_integration_tool.name}': {e.message}",
-                code=e.code,
-                raw_data=e.raw_data,
-            ) from e
         except Exception as e:
-            raise NodeException(
-                message=f"Error executing Vellum Integration tool '{self.vellum_integration_tool.name}': {str(e)}",
-                code=WorkflowErrorCode.NODE_EXECUTION,
-            ) from e
+            self._handle_tool_exception(e, "Vellum Integration tool", self.vellum_integration_tool.name)
 
         # Add result to chat history
         self._add_function_result_to_chat_history(result, self.state)