vellum-ai 1.0.4__py3-none-any.whl → 1.0.6__py3-none-any.whl

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

@@ -1,72 +1,65 @@
 import pytest
 from unittest.mock import Mock, patch
 
-from vellum.workflows.nodes.displayable.tool_calling_node.composio_service import ComposioService, ConnectionInfo
+from vellum.workflows.integrations.composio_service import ComposioService, ConnectionInfo
 
 
 @pytest.fixture
-def mock_composio_client():
-    """Mock the Composio client completely"""
-    with patch("vellum.workflows.nodes.displayable.tool_calling_node.composio_service.ComposioClient") as mock_composio:
-        yield mock_composio.return_value
+def mock_requests():
+    """Mock requests module"""
+    with patch("vellum.workflows.integrations.composio_service.requests") as mock_requests:
+        yield mock_requests
 
 
 @pytest.fixture
 def mock_connected_accounts_response():
-    """Mock response for connected accounts"""
-    mock_item1 = Mock()
-    mock_item1.id = "conn-123"
-    mock_item1.toolkit.slug = "github"
-    mock_item1.status = "ACTIVE"
-    mock_item1.created_at = "2023-01-01T00:00:00Z"
-    mock_item1.updated_at = "2023-01-15T10:30:00Z"
-
-    mock_item2 = Mock()
-    mock_item2.id = "conn-456"
-    mock_item2.toolkit.slug = "slack"
-    mock_item2.status = "ACTIVE"
-    mock_item2.created_at = "2023-01-01T00:00:00Z"
-    mock_item2.updated_at = "2023-01-10T08:00:00Z"
-
-    mock_response = Mock()
-    mock_response.items = [mock_item1, mock_item2]
-
-    return mock_response
+    """Mock response for connected accounts API"""
+    return {
+        "items": [
+            {
+                "id": "conn-123",
+                "toolkit": {"slug": "github"},
+                "status": "ACTIVE",
+                "created_at": "2023-01-01T00:00:00Z",
+                "updated_at": "2023-01-15T10:30:00Z",
+            },
+            {
+                "id": "conn-456",
+                "toolkit": {"slug": "slack"},
+                "status": "ACTIVE",
+                "created_at": "2023-01-01T00:00:00Z",
+                "updated_at": "2023-01-10T08:00:00Z",
+            },
+        ]
+    }
 
 
 @pytest.fixture
-def mock_composio_core_client():
-    """Mock the composio-core Composio client"""
-    with patch("vellum.workflows.nodes.displayable.tool_calling_node.composio_service.Composio") as mock_composio:
-        yield mock_composio.return_value
+def mock_tool_execution_response():
+    """Mock response for tool execution API"""
+    return {
+        "data": {"items": [], "total": 0},
+        "successful": True,
+        "error": None,
+    }
 
 
 @pytest.fixture
-def mock_action():
-    """Mock the Action class and specific actions"""
-    with patch("vellum.workflows.nodes.displayable.tool_calling_node.composio_service.Action") as mock_action_class:
-        # Mock a specific action
-        mock_hackernews_action = Mock()
-        mock_action_class.HACKERNEWS_GET_USER = mock_hackernews_action
-        mock_action_class.GITHUB_GET_USER = Mock()
-        yield mock_action_class
-
-
-@pytest.fixture
-def composio_service(mock_composio_client, mock_composio_core_client):
-    """Create ComposioService with mocked clients"""
+def composio_service():
+    """Create ComposioService with test API key"""
     return ComposioService(api_key="test-key")
 
 
 class TestComposioAccountService:
     """Test suite for ComposioAccountService"""
 
-    def test_get_user_connections_success(
-        self, composio_service, mock_composio_client, mock_connected_accounts_response
-    ):
+    def test_get_user_connections_success(self, composio_service, mock_requests, mock_connected_accounts_response):
         """Test successful retrieval of user connections"""
-        # GIVEN the Composio client returns a valid response with two connections
-        mock_composio_client.connected_accounts.list.return_value = mock_connected_accounts_response
+        # GIVEN the requests mock returns a valid response with two connections
+        mock_response = Mock()
+        mock_response.json.return_value = mock_connected_accounts_response
+        mock_response.raise_for_status.return_value = None
+        mock_requests.get.return_value = mock_response
 
         # WHEN we request user connections
         result = composio_service.get_user_connections()
@@ -86,14 +79,21 @@ class TestComposioAccountService:
         assert result[1].created_at == "2023-01-01T00:00:00Z"
         assert result[1].updated_at == "2023-01-10T08:00:00Z"
 
-        mock_composio_client.connected_accounts.list.assert_called_once()
+        # Verify the correct API endpoint was called
+        mock_requests.get.assert_called_once_with(
+            "https://backend.composio.dev/api/v3/connected_accounts",
+            headers={"x-api-key": "test-key", "Content-Type": "application/json"},
+            params={},
+            timeout=30,
+        )
 
-    def test_get_user_connections_empty_response(self, composio_service, mock_composio_client):
+    def test_get_user_connections_empty_response(self, composio_service, mock_requests):
         """Test handling of empty connections response"""
-        # GIVEN the Composio client returns an empty response
+        # GIVEN the requests mock returns an empty response
         mock_response = Mock()
-        mock_response.items = []
-        mock_composio_client.connected_accounts.list.return_value = mock_response
+        mock_response.json.return_value = {"items": []}
+        mock_response.raise_for_status.return_value = None
+        mock_requests.get.return_value = mock_response
 
         # WHEN we request user connections
         result = composio_service.get_user_connections()
@@ -105,18 +105,23 @@ class TestComposioAccountService:
 class TestComposioCoreService:
     """Test suite for ComposioCoreService"""
 
-    def test_execute_tool_success(self, composio_service, mock_composio_core_client, mock_action):
+    def test_execute_tool_success(self, composio_service, mock_requests, mock_tool_execution_response):
         """Test executing a tool with complex argument structure"""
         # GIVEN complex arguments and a mock response
         complex_args = {"filters": {"status": "active"}, "limit": 10, "sort": "created_at"}
-        expected_result = {"items": [], "total": 0}
-        mock_composio_core_client.actions.execute.return_value = expected_result
+        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 complex arguments
         result = composio_service.execute_tool("HACKERNEWS_GET_USER", complex_args)
 
-        # THEN the arguments are passed through correctly
-        mock_composio_core_client.actions.execute.assert_called_once_with(
-            mock_action.HACKERNEWS_GET_USER, params=complex_args
+        # THEN the arguments are passed through correctly and we get the expected result
+        mock_requests.post.assert_called_once_with(
+            "https://backend.composio.dev/api/v3/tools/execute/HACKERNEWS_GET_USER",
+            headers={"x-api-key": "test-key", "Content-Type": "application/json"},
+            json={"arguments": complex_args},
+            timeout=30,
         )
-        assert result == expected_result
+        assert result == {"items": [], "total": 0}

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

@@ -1,10 +1,12 @@
+import pytest
+
 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 get_function_name
 from vellum.workflows.outputs.base import BaseOutputs
 from vellum.workflows.state.base import BaseState
-from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
 
 
 def test_get_function_name_callable():
@@ -56,3 +58,21 @@ def test_get_function_name_subworkflow_deployment_uuid():
     result = get_function_name(deployment_config)
 
     assert result == "57f09bebb46340e0bf9ec972e664352f"
+
+
+@pytest.mark.parametrize(
+    "toolkit,action,description,expected_result",
+    [
+        ("SLACK", "SLACK_SEND_MESSAGE", "Send message to Slack", "slack_send_message"),
+        ("GMAIL", "GMAIL_CREATE_EMAIL_DRAFT", "Create Gmail draft", "gmail_create_email_draft"),
+    ],
+)
+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)
+
+    result = get_function_name(composio_tool)
+
+    assert result == expected_result

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

@@ -1,11 +1,13 @@
 import json
-from typing import Any, Callable, Iterator, List, Optional, Type, cast
+import logging
+from typing import Any, Callable, Iterator, List, Optional, Type, Union, cast
 
 from pydash import snake_case
 
 from vellum import ChatMessage, PromptBlock
 from vellum.client.types.function_call_chat_message_content import FunctionCallChatMessageContent
 from vellum.client.types.function_call_chat_message_content_value import FunctionCallChatMessageContentValue
+from vellum.client.types.function_definition import FunctionDefinition
 from vellum.client.types.prompt_output import PromptOutput
 from vellum.client.types.prompt_parameters import PromptParameters
 from vellum.client.types.string_chat_message_content import StringChatMessageContent
@@ -14,6 +16,7 @@ from vellum.workflows.errors.types import WorkflowErrorCode
 from vellum.workflows.exceptions import NodeException
 from vellum.workflows.expressions.concat import ConcatExpression
 from vellum.workflows.inputs import BaseInputs
+from vellum.workflows.integrations.composio_service import ComposioService
 from vellum.workflows.nodes.bases import BaseNode
 from vellum.workflows.nodes.core.inline_subworkflow_node.node import InlineSubworkflowNode
 from vellum.workflows.nodes.displayable.inline_prompt_node.node import InlinePromptNode
@@ -31,6 +34,32 @@ from vellum.workflows.types.generics import is_workflow_class
 CHAT_HISTORY_VARIABLE = "chat_history"
 
 
+logger = logging.getLogger(__name__)
+
+
+class FunctionCallNodeMixin:
+    """Mixin providing common functionality for nodes that handle function calls."""
+
+    function_call_output: List[PromptOutput]
+
+    def _extract_function_arguments(self) -> dict:
+        """Extract arguments from function call output."""
+        if self.function_call_output and len(self.function_call_output) > 0:
+            function_call = self.function_call_output[0]
+            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
+                return function_call.value.arguments or {}
+        return {}
+
+    def _add_function_result_to_chat_history(self, result: Any, state: ToolCallingState) -> None:
+        """Add function execution result to chat history."""
+        state.chat_history.append(
+            ChatMessage(
+                role="FUNCTION",
+                content=StringChatMessageContent(value=json.dumps(result, cls=DefaultStateEncoder)),
+            )
+        )
+
+
 class ToolRouterNode(InlinePromptNode[ToolCallingState]):
     max_prompt_iterations: Optional[int] = 5
 
@@ -69,20 +98,11 @@ class ToolRouterNode(InlinePromptNode[ToolCallingState]):
             yield output
 
 
-class DynamicSubworkflowDeploymentNode(SubworkflowDeploymentNode[ToolCallingState]):
+class DynamicSubworkflowDeploymentNode(SubworkflowDeploymentNode[ToolCallingState], FunctionCallNodeMixin):
     """Node that executes a deployment definition with function call output."""
 
-    function_call_output: List[PromptOutput]
-
     def run(self) -> Iterator[BaseOutput]:
-        if self.function_call_output and len(self.function_call_output) > 0:
-            function_call = self.function_call_output[0]
-            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
-                arguments = function_call.value.arguments
-            else:
-                arguments = {}
-        else:
-            arguments = {}
+        arguments = self._extract_function_arguments()
 
         # Mypy doesn't like instance assignments of class attributes. It's safe in our case tho bc it's what
         # we do in the `__init__` method. Long term, instead of the function_call_output attribute above, we
@@ -103,28 +123,16 @@ class DynamicSubworkflowDeploymentNode(SubworkflowDeploymentNode[ToolCallingStat
             yield output
 
         # Add the result to the chat history
-        self.state.chat_history.append(
-            ChatMessage(
-                role="FUNCTION",
-                content=StringChatMessageContent(value=json.dumps(outputs, cls=DefaultStateEncoder)),
-            )
-        )
+        self._add_function_result_to_chat_history(outputs, self.state)
 
 
-class DynamicInlineSubworkflowNode(InlineSubworkflowNode[ToolCallingState, BaseInputs, BaseState]):
+class DynamicInlineSubworkflowNode(
+    InlineSubworkflowNode[ToolCallingState, BaseInputs, BaseState], FunctionCallNodeMixin
+):
     """Node that executes an inline subworkflow with function call output."""
 
-    function_call_output: List[PromptOutput]
-
     def run(self) -> Iterator[BaseOutput]:
-        if self.function_call_output and len(self.function_call_output) > 0:
-            function_call = self.function_call_output[0]
-            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
-                arguments = function_call.value.arguments
-            else:
-                arguments = {}
-        else:
-            arguments = {}
+        arguments = self._extract_function_arguments()
 
         self.subworkflow_inputs = arguments  # type: ignore[misc]
 
@@ -137,29 +145,16 @@ class DynamicInlineSubworkflowNode(InlineSubworkflowNode[ToolCallingState, BaseI
             yield output
 
         # Add the result to the chat history
-        self.state.chat_history.append(
-            ChatMessage(
-                role="FUNCTION",
-                content=StringChatMessageContent(value=json.dumps(outputs, cls=DefaultStateEncoder)),
-            )
-        )
+        self._add_function_result_to_chat_history(outputs, self.state)
 
 
-class FunctionNode(BaseNode[ToolCallingState]):
+class FunctionNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
     """Node that executes a regular Python function with function call output."""
 
-    function_call_output: List[PromptOutput]
     function_definition: Callable[..., Any]
 
     def run(self) -> Iterator[BaseOutput]:
-        if self.function_call_output and len(self.function_call_output) > 0:
-            function_call = self.function_call_output[0]
-            if function_call.type == "FUNCTION_CALL" and function_call.value is not None:
-                arguments = function_call.value.arguments
-            else:
-                arguments = {}
-        else:
-            arguments = {}
+        arguments = self._extract_function_arguments()
 
         try:
             result = self.function_definition(**arguments)
@@ -171,16 +166,73 @@ class FunctionNode(BaseNode[ToolCallingState]):
             )
 
         # Add the result to the chat history
-        self.state.chat_history.append(
-            ChatMessage(
-                role="FUNCTION",
-                content=StringChatMessageContent(value=json.dumps(result, cls=DefaultStateEncoder)),
+        self._add_function_result_to_chat_history(result, self.state)
+
+        yield from []
+
+
+class ComposioNode(BaseNode[ToolCallingState], FunctionCallNodeMixin):
+    """Node that executes a Composio tool with function call output."""
+
+    composio_tool: ComposioToolDefinition
+
+    def run(self) -> Iterator[BaseOutput]:
+        # Extract arguments from function call
+        arguments = self._extract_function_arguments()
+
+        try:
+            # Execute using ComposioService
+            composio_service = ComposioService()
+            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,
             )
-        )
+
+        # Add result to chat history
+        self._add_function_result_to_chat_history(result, self.state)
 
         yield from []
 
 
+def _hydrate_composio_tool_definition(tool_def: ComposioToolDefinition) -> ComposioToolDefinition:
+    """Hydrate a ComposioToolDefinition with detailed information from the Composio API.
+
+    Args:
+        tool_def: The basic ComposioToolDefinition to enhance
+
+    Returns:
+        ComposioToolDefinition with detailed parameters and description
+    """
+    try:
+        composio_service = ComposioService()
+        tool_details = composio_service.get_tool_by_slug(tool_def.action)
+
+        # Extract toolkit information from API response
+        toolkit_info = tool_details.get("toolkit", {})
+        toolkit_slug = (
+            toolkit_info.get("slug", tool_def.toolkit) if isinstance(toolkit_info, dict) else tool_def.toolkit
+        )
+
+        # Create a version of the tool definition with proper field extraction
+        return ComposioToolDefinition(
+            type=tool_def.type,
+            toolkit=toolkit_slug.upper() if toolkit_slug else tool_def.toolkit,
+            action=tool_details.get("slug", tool_def.action),
+            description=tool_details.get("description", tool_def.description),
+            display_name=tool_details.get("name", tool_def.display_name),
+            parameters=tool_details.get("input_parameters", tool_def.parameters),
+            version=tool_details.get("version", tool_def.version),
+            tags=tool_details.get("tags", tool_def.tags),
+        )
+
+    except Exception as e:
+        # If hydration fails (including no API key), log and return original
+        logger.warning(f"Failed to enhance Composio tool '{tool_def.action}': {e}")
+        return tool_def
+
+
 def create_tool_router_node(
     ml_model: str,
     blocks: List[PromptBlock],
@@ -190,9 +242,25 @@ def create_tool_router_node(
     max_prompt_iterations: Optional[int] = None,
 ) -> Type[ToolRouterNode]:
     if functions and len(functions) > 0:
-        # If we have functions, create dynamic ports for each function
+        # Create dynamic ports and convert functions in a single loop
         Ports = type("Ports", (), {})
+        prompt_functions: List[Union[Tool, FunctionDefinition]] = []
+
         for function in functions:
+            if isinstance(function, ComposioToolDefinition):
+                # Get Composio tool details and hydrate the function definition
+                enhanced_function = _hydrate_composio_tool_definition(function)
+                prompt_functions.append(
+                    FunctionDefinition(
+                        name=enhanced_function.name,
+                        description=enhanced_function.description,
+                        parameters=enhanced_function.parameters,
+                    )
+                )
+            else:
+                prompt_functions.append(function)
+
+            # Create port for this function (using original function for get_function_name)
             function_name = get_function_name(function)
 
             # Avoid using lambda to capture function_name
@@ -215,6 +283,7 @@ def create_tool_router_node(
     else:
         # If no functions exist, create a simple Ports class with just a default port
         Ports = type("Ports", (), {"default": Port(default=True)})
+        prompt_functions = []
 
     # Add a chat history block to blocks only if one doesn't already exist
     has_chat_history_block = any(
@@ -247,7 +316,7 @@ def create_tool_router_node(
             {
                 "ml_model": ml_model,
                 "blocks": blocks,
-                "functions": functions,
+                "functions": prompt_functions,  # Use converted functions for prompt layer
                 "prompt_inputs": node_prompt_inputs,
                 "parameters": parameters,
                 "max_prompt_iterations": max_prompt_iterations,
@@ -291,20 +360,16 @@ def create_function_node(
         return node
 
     elif isinstance(function, ComposioToolDefinition):
-        # ComposioToolDefinition execution not yet implemented
-        def composio_not_implemented(**kwargs):
-            raise NotImplementedError("ComposioToolDefinition execution not yet implemented")
-
         node = type(
             f"ComposioNode_{function.name}",
-            (FunctionNode,),
+            (ComposioNode,),
             {
-                "function_definition": composio_not_implemented,
+                "composio_tool": function,
                 "function_call_output": tool_router_node.Outputs.results,
                 "__module__": __name__,
             },
         )
-
+        return node
     elif is_workflow_class(function):
         node = type(
             f"DynamicInlineSubworkflowNode_{function.__name__}",

vellum/workflows/types/definition.py

@@ -2,7 +2,7 @@ import importlib
 import inspect
 from types import FrameType
 from uuid import UUID
-from typing import Annotated, Any, Dict, Literal, Optional, Union
+from typing import Annotated, Any, Dict, List, Literal, Optional, Union
 
 from pydantic import BeforeValidator
 
@@ -109,8 +109,10 @@ class ComposioToolDefinition(UniversalBaseModel):
     action: str  # Specific action like "GITHUB_CREATE_AN_ISSUE"
     description: str
 
-    # Optional cached metadata
     display_name: Optional[str] = None
+    parameters: Optional[Dict[str, Any]] = None
+    version: Optional[str] = None
+    tags: Optional[List[str]] = None
 
     @property
     def name(self) -> str:

vellum/workflows/utils/functions.py

@@ -1,6 +1,6 @@
 import dataclasses
 import inspect
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional, Type, Union, get_args, get_origin
+from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Type, Union, get_args, get_origin
 
 from pydantic import BaseModel
 from pydantic_core import PydanticUndefined
@@ -33,6 +33,18 @@ def compile_annotation(annotation: Optional[Any], defs: dict[str, Any]) -> dict:
     if get_origin(annotation) is Union:
         return {"anyOf": [compile_annotation(a, defs) for a in get_args(annotation)]}
 
+    if get_origin(annotation) is Literal:
+        values = list(get_args(annotation))
+        types = {type(value) for value in values}
+        if len(types) == 1:
+            value_type = types.pop()
+            if value_type in type_map:
+                return {"type": type_map[value_type], "enum": values}
+            else:
+                return {"enum": values}
+        else:
+            return {"enum": values}
+
     if get_origin(annotation) is dict:
         _, value_type = get_args(annotation)
         return {"type": "object", "additionalProperties": compile_annotation(value_type, defs)}

vellum/workflows/utils/tests/test_functions.py

@@ -1,6 +1,8 @@
+import pytest
 from dataclasses import dataclass
+from enum import Enum
 from unittest.mock import Mock
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 from pydantic import BaseModel
 
@@ -581,3 +583,32 @@ def test_compile_workflow_deployment_function_definition__defaults():
             "required": ["no_default"],
         },
     )
+
+
+@pytest.mark.parametrize(
+    "annotation,expected_schema",
+    [
+        (Literal["a", "b"], {"type": "string", "enum": ["a", "b"]}),
+        (Literal["a", 1], {"enum": ["a", 1]}),
+    ],
+)
+def test_compile_function_definition__literal(annotation, expected_schema):
+    def my_function(a: annotation):  # type: ignore
+        pass
+
+    compiled_function = compile_function_definition(my_function)
+    assert isinstance(compiled_function.parameters, dict)
+    assert compiled_function.parameters["properties"]["a"] == expected_schema
+
+
+def test_compile_function_definition__literal_type_not_in_map():
+    class MyEnum(Enum):
+        FOO = "foo"
+        BAR = "bar"
+
+    def my_function(a: Literal[MyEnum.FOO, MyEnum.BAR]):
+        pass
+
+    compiled_function = compile_function_definition(my_function)
+    assert isinstance(compiled_function.parameters, dict)
+    assert compiled_function.parameters["properties"]["a"] == {"enum": [MyEnum.FOO, MyEnum.BAR]}

{vellum_ai-1.0.4.dist-info → vellum_ai-1.0.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.0.4
+Version: 1.0.6
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -22,8 +22,6 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Typing :: Typed
 Requires-Dist: Jinja2 (>=3.1.0,<4.0.0)
 Requires-Dist: click (>=8.1.7,<9.0.0)
-Requires-Dist: composio-client (>=1.5.0,<2.0.0)
-Requires-Dist: composio-core (>=0.7.20,<1.0.0)
 Requires-Dist: docker (>=7.1.0,<8.0.0)
 Requires-Dist: httpx (>=0.21.2)
 Requires-Dist: openai (>=1.0.0,<2.0.0)