vellum-ai 1.0.9__py3-none-any.whl → 1.0.11__py3-none-any.whl

vellum/workflows/integrations/composio_service.py

@@ -135,16 +135,24 @@ class ComposioService:
             else:
                 raise NodeException(f"Failed to retrieve tool details for '{tool_slug}': {error_message}")
 
-    def execute_tool(self, tool_name: str, arguments: Dict[str, Any]) -> Any:
+    def execute_tool(self, tool_name: str, arguments: Dict[str, Any], user_id: Optional[str] = None) -> Any:
         """Execute a tool using direct API request
 
         Args:
             tool_name: The name of the tool to execute (e.g., "HACKERNEWS_GET_USER")
             arguments: Dictionary of arguments to pass to the tool
+            user_id: Optional user ID to identify which user's Composio connection to use
 
         Returns:
             The result of the tool execution
         """
         endpoint = f"/tools/execute/{tool_name}"
-        response = self._make_request(endpoint, method="POST", json_data={"arguments": arguments})
+        json_data: Dict[str, Any] = {"arguments": arguments}
+        if user_id is not None:
+            json_data["user_id"] = user_id
+        response = self._make_request(endpoint, method="POST", json_data=json_data)
+
+        if not response.get("successful", True):
+            return response.get("error", "Tool execution failed")
+
         return response.get("data", response)

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

@@ -14,7 +14,7 @@ from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class BasePromptNode(BaseNode, Generic[StateType]):
+class BasePromptNode(BaseNode[StateType], Generic[StateType]):
     # Inputs that are passed to the Prompt
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
 

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

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Dict, Iterator, Type, Union
+from typing import Any, Dict, Generic, Iterator, Type, Union
 
 from vellum.workflows.constants import undefined
 from vellum.workflows.errors import WorkflowErrorCode
@@ -10,7 +10,7 @@ from vellum.workflows.types import MergeBehavior
 from vellum.workflows.types.generics import StateType
 
 
-class InlinePromptNode(BaseInlinePromptNode[StateType]):
+class InlinePromptNode(BaseInlinePromptNode[StateType], Generic[StateType]):
     """
     Used to execute a Prompt defined inline.
 

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

@@ -1,4 +1,4 @@
-from typing import Any, ClassVar, Dict, Iterator, List, Optional, Set, Union
+from typing import Any, ClassVar, Dict, Generic, Iterator, List, Optional, Set, Union
 
 from vellum import ChatMessage, PromptBlock
 from vellum.client.types.prompt_parameters import PromptParameters
@@ -12,19 +12,23 @@ from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
 from vellum.workflows.nodes.displayable.tool_calling_node.utils import (
+    create_else_node,
     create_function_node,
+    create_mcp_tool_node,
     create_tool_router_node,
     get_function_name,
+    get_mcp_tool_name,
     hydrate_mcp_tool_definitions,
 )
 from vellum.workflows.outputs.base import BaseOutput, BaseOutputs
 from vellum.workflows.state.context import WorkflowContext
-from vellum.workflows.types.core import EntityInputsInterface, Tool, ToolSource
+from vellum.workflows.types.core import EntityInputsInterface, Tool
 from vellum.workflows.types.definition import MCPServer
+from vellum.workflows.types.generics import StateType
 from vellum.workflows.workflows.event_filters import all_workflow_event_filter
 
 
-class ToolCallingNode(BaseNode):
+class ToolCallingNode(BaseNode[StateType], Generic[StateType]):
     """
     A Node that dynamically invokes the provided functions to the underlying Prompt
 
@@ -32,7 +36,6 @@ class ToolCallingNode(BaseNode):
         ml_model: str - The model to use for tool calling (e.g., "gpt-4o-mini")
         blocks: List[PromptBlock] - The prompt blocks to use (same format as InlinePromptNode)
         functions: List[Tool] - The functions that can be called
-        tool_sources: List[ToolSource] - The tool sources that can be called
         prompt_inputs: Optional[EntityInputsInterface] - Mapping of input variable names to values
         parameters: PromptParameters - The parameters for the Prompt
         max_prompt_iterations: Optional[int] - Maximum number of prompt iterations before stopping
@@ -41,7 +44,6 @@ class ToolCallingNode(BaseNode):
     ml_model: ClassVar[str] = "gpt-4o-mini"
     blocks: ClassVar[List[Union[PromptBlock, Dict[str, Any]]]] = []
     functions: ClassVar[List[Tool]] = []
-    tool_sources: ClassVar[List[ToolSource]] = []
     prompt_inputs: ClassVar[Optional[EntityInputsInterface]] = None
     parameters: PromptParameters = DEFAULT_PROMPT_PARAMETERS
     max_prompt_iterations: ClassVar[Optional[int]] = 5
@@ -138,7 +140,6 @@ class ToolCallingNode(BaseNode):
             ml_model=self.ml_model,
             blocks=self.blocks,
             functions=self.functions,
-            tool_sources=self.tool_sources,
             prompt_inputs=self.prompt_inputs,
             parameters=self.parameters,
             max_prompt_iterations=self.max_prompt_iterations,
@@ -146,23 +147,22 @@ class ToolCallingNode(BaseNode):
 
         self._function_nodes = {}
         for function in self.functions:
-            function_name = get_function_name(function)
-
-            self._function_nodes[function_name] = create_function_node(
-                function=function,
-                tool_router_node=self.tool_router_node,
-            )
-
-        for tool_source in self.tool_sources:
-            if isinstance(tool_source, MCPServer):
-                tool_definitions = hydrate_mcp_tool_definitions(tool_source)
+            if isinstance(function, MCPServer):
+                tool_definitions = hydrate_mcp_tool_definitions(function)
                 for tool_definition in tool_definitions:
-                    function_name = get_function_name(tool_definition)
+                    function_name = get_mcp_tool_name(tool_definition)
 
-                    self._function_nodes[function_name] = create_function_node(
-                        function=tool_definition,
+                    self._function_nodes[function_name] = create_mcp_tool_node(
+                        tool_def=tool_definition,
                         tool_router_node=self.tool_router_node,
                     )
+            else:
+                function_name = get_function_name(function)
+
+                self._function_nodes[function_name] = create_function_node(
+                    function=function,
+                    tool_router_node=self.tool_router_node,
+                )
 
         graph_set = set()
 
@@ -172,7 +172,11 @@ class ToolCallingNode(BaseNode):
             edge_graph = router_port >> FunctionNodeClass >> self.tool_router_node
             graph_set.add(edge_graph)
 
-        default_port = getattr(self.tool_router_node.Ports, "default")
+        else_node = create_else_node(self.tool_router_node)
+        default_port = self.tool_router_node.Ports.default >> {
+            else_node.Ports.loop >> self.tool_router_node,
+            else_node.Ports.end,
+        }
         graph_set.add(default_port)
 
         self._graph = Graph.from_set(graph_set)

vellum/workflows/nodes/displayable/tool_calling_node/state.py

@@ -7,3 +7,5 @@ from vellum.workflows.state.base import BaseState
 class ToolCallingState(BaseState):
     chat_history: List[ChatMessage] = []
     prompt_iterations: int = 0
+    current_prompt_output_index: int = 0
+    current_function_calls_processed: int = 0

vellum/workflows/nodes/displayable/tool_calling_node/tests/test_composio_service.py

@@ -44,6 +44,21 @@ def mock_tool_execution_response():
     }
 
 
+@pytest.fixture
+def mock_tool_execution_error_response():
+    """Mock response for failed tool execution API"""
+    return {
+        "data": {},
+        "successful": False,
+        "error": (
+            'Request failed error: `{"message":"Not Found",'
+            '"documentation_url":"https://docs.github.com/rest/pulls/pulls#get-a-pull-request",'
+            '"status":"404"}`'
+        ),
+        "log_id": "log_raE_fIWNcDPo",
+    }
+
+
 @pytest.fixture
 def composio_service():
     """Create ComposioService with test API key"""
@@ -125,3 +140,80 @@ class TestComposioCoreService:
             timeout=30,
         )
         assert result == {"items": [], "total": 0}
+
+    def test_execute_tool_with_user_id(self, composio_service, mock_requests, mock_tool_execution_response):
+        """Test executing a tool with user_id parameter"""
+        # GIVEN a user_id and tool arguments
+        user_id = "test_user_123"
+        tool_args = {"param1": "value1"}
+        mock_response = Mock()
+        mock_response.json.return_value = mock_tool_execution_response
+        mock_response.raise_for_status.return_value = None
+        mock_requests.post.return_value = mock_response
+
+        # WHEN we execute a tool with user_id
+        result = composio_service.execute_tool("TEST_TOOL", tool_args, user_id=user_id)
+
+        # THEN the user_id should be included in the request payload
+        mock_requests.post.assert_called_once_with(
+            "https://backend.composio.dev/api/v3/tools/execute/TEST_TOOL",
+            headers={"x-api-key": "test-key", "Content-Type": "application/json"},
+            json={"arguments": tool_args, "user_id": user_id},
+            timeout=30,
+        )
+        assert result == {"items": [], "total": 0}
+
+    def test_execute_tool_without_user_id(self, composio_service, mock_requests, mock_tool_execution_response):
+        """Test executing a tool without user_id parameter maintains backward compatibility"""
+        # GIVEN tool arguments without user_id
+        tool_args = {"param1": "value1"}
+        mock_response = Mock()
+        mock_response.json.return_value = mock_tool_execution_response
+        mock_response.raise_for_status.return_value = None
+        mock_requests.post.return_value = mock_response
+
+        # WHEN we execute a tool without user_id
+        result = composio_service.execute_tool("TEST_TOOL", tool_args)
+
+        # THEN the user_id should NOT be included in the request payload
+        mock_requests.post.assert_called_once_with(
+            "https://backend.composio.dev/api/v3/tools/execute/TEST_TOOL",
+            headers={"x-api-key": "test-key", "Content-Type": "application/json"},
+            json={"arguments": tool_args},
+            timeout=30,
+        )
+        assert result == {"items": [], "total": 0}
+
+    def test_execute_tool_failure_surfaces_error(
+        self, composio_service, mock_requests, mock_tool_execution_error_response
+    ):
+        """Test that tool execution failures surface detailed error information"""
+        # GIVEN a mock response indicating tool execution failure
+        mock_response = Mock()
+        mock_response.json.return_value = mock_tool_execution_error_response
+        mock_response.raise_for_status.return_value = None
+        mock_requests.post.return_value = mock_response
+
+        # WHEN we execute a tool that fails
+        result = composio_service.execute_tool("GITHUB_GET_PR", {"repo": "test", "pr_number": 999})
+
+        # THEN the result should contain the detailed error message from the API
+        assert "Request failed error" in result
+        assert "Not Found" in result
+
+    def test_execute_tool_failure_with_generic_error_message(self, composio_service, mock_requests):
+        """Test that tool execution failures with missing error field use generic message"""
+        # GIVEN a mock response indicating tool execution failure without error field
+        mock_response = Mock()
+        mock_response.json.return_value = {
+            "data": {},
+            "successful": False,
+        }
+        mock_response.raise_for_status.return_value = None
+        mock_requests.post.return_value = mock_response
+
+        # WHEN we execute a tool that fails
+        result = composio_service.execute_tool("TEST_TOOL", {"param": "value"})
+
+        # THEN the result should contain the generic error message
+        assert result == "Tool execution failed"

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

@@ -1,6 +1,6 @@
 import json
 from uuid import uuid4
-from typing import Any, Iterator, List
+from typing import Any, Iterator
 
 from vellum import ChatMessage
 from vellum.client.types.fulfilled_execute_prompt_event import FulfilledExecutePromptEvent
@@ -15,6 +15,7 @@ from vellum.workflows import BaseWorkflow
 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
+from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
 from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_function_node, create_tool_router_node
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState, StateMeta
@@ -39,13 +40,12 @@ def test_port_condition_match_function_name():
         ml_model="test-model",
         blocks=[],
         functions=[first_function, second_function],
-        tool_sources=[],
         prompt_inputs=None,
         parameters=DEFAULT_PROMPT_PARAMETERS,
     )
 
     # AND a state with a function call to the first function
-    state = BaseState(
+    state = ToolCallingState(
         meta=StateMeta(
             node_outputs={
                 router_node.Outputs.results: [
@@ -98,7 +98,6 @@ def test_tool_calling_node_inline_workflow_context():
         ml_model="test-model",
         blocks=[],
         functions=[MyWorkflow],
-        tool_sources=[],
         prompt_inputs=None,
         parameters=DEFAULT_PROMPT_PARAMETERS,
     )
@@ -118,12 +117,6 @@ def test_tool_calling_node_inline_workflow_context():
     )
     function_node._context = parent_context
 
-    # Create a state with chat_history for the function node
-    class TestState(BaseState):
-        chat_history: List[ChatMessage] = []
-
-    function_node.state = TestState(meta=StateMeta(node_outputs={tool_router_node.Outputs.text: '{"arguments": {}}'}))
-
     # WHEN the function node runs
     outputs = list(function_node.run())
 
@@ -213,3 +206,25 @@ def test_tool_calling_node_with_user_provided_chat_history_block(vellum_adhoc_pr
     ]
     assert len(chat_history_inputs) == 1
     assert chat_history_inputs[0].value == [ChatMessage(role="USER", text="Hello from user")]
+
+
+def test_tool_calling_node_with_generic_type_parameter():
+    # GIVEN a custom state class
+    class State(BaseState):
+        pass
+
+    # AND a ToolCallingNode that uses the generic type parameter
+    class TestToolCallingNode(ToolCallingNode[State]):
+        ml_model = "gpt-4o-mini"
+        blocks = []
+        functions = [first_function]
+        max_prompt_iterations = 1
+
+    # WHEN we create an instance of the node
+    state = State()
+    node = TestToolCallingNode(state=state)
+
+    # THEN the node should be created successfully
+    assert node is not None
+    assert isinstance(node, TestToolCallingNode)
+    assert node.state == state

vellum/workflows/nodes/displayable/tool_calling_node/tests/test_utils.py

@@ -12,7 +12,11 @@ from vellum.prompts.constants import DEFAULT_PROMPT_PARAMETERS
 from vellum.workflows import BaseWorkflow
 from vellum.workflows.inputs.base import BaseInputs
 from vellum.workflows.nodes.bases import BaseNode
-from vellum.workflows.nodes.displayable.tool_calling_node.utils import create_tool_router_node, get_function_name
+from vellum.workflows.nodes.displayable.tool_calling_node.utils import (
+    create_tool_router_node,
+    get_function_name,
+    get_mcp_tool_name,
+)
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
 from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition, MCPServer, MCPToolDefinition
@@ -77,7 +81,7 @@ def test_get_function_name_mcp_tool_definition():
         parameters={"repository_name": "string", "description": "string"},
     )
 
-    result = get_function_name(mcp_tool)
+    result = get_mcp_tool_name(mcp_tool)
 
     assert result == "github__create_repository"
 
@@ -93,7 +97,7 @@ def test_get_function_name_composio_tool_definition_various_toolkits(
     toolkit: str, action: str, description: str, expected_result: str
 ):
     """Test ComposioToolDefinition function name generation with various toolkits."""
-    composio_tool = ComposioToolDefinition(toolkit=toolkit, action=action, description=description)
+    composio_tool = ComposioToolDefinition(toolkit=toolkit, action=action, description=description, user_id=None)
 
     result = get_function_name(composio_tool)
 
@@ -106,7 +110,6 @@ def test_create_tool_router_node_max_prompt_iterations(vellum_adhoc_prompt_clien
         ml_model="gpt-4o-mini",
         blocks=[],
         functions=[],
-        tool_sources=[],
         prompt_inputs=None,
         parameters=DEFAULT_PROMPT_PARAMETERS,
         max_prompt_iterations=None,
@@ -166,7 +169,6 @@ def test_create_tool_router_node_chat_history_block_dict(vellum_adhoc_prompt_cli
         ml_model="gpt-4o-mini",
         blocks=blocks,  # type: ignore
         functions=[],
-        tool_sources=[],
         prompt_inputs=None,
         parameters=DEFAULT_PROMPT_PARAMETERS,
     )