vellum-ai 1.0.3__py3-none-any.whl → 1.0.5__py3-none-any.whl

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

@@ -1,4 +1,5 @@
 import json
+import os
 from typing import Any, Callable, Iterator, List, Optional, Type, cast
 
 from pydash import snake_case
@@ -12,11 +13,13 @@ from vellum.client.types.string_chat_message_content import StringChatMessageCon
 from vellum.client.types.variable_prompt_block import VariablePromptBlock
 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.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
 from vellum.workflows.nodes.displayable.subworkflow_deployment_node.node import SubworkflowDeploymentNode
+from vellum.workflows.nodes.displayable.tool_calling_node.composio_service import ComposioService
 from vellum.workflows.nodes.displayable.tool_calling_node.state import ToolCallingState
 from vellum.workflows.outputs.base import BaseOutput
 from vellum.workflows.ports.port import Port
@@ -24,12 +27,35 @@ from vellum.workflows.references.lazy import LazyReference
 from vellum.workflows.state import BaseState
 from vellum.workflows.state.encoder import DefaultStateEncoder
 from vellum.workflows.types.core import EntityInputsInterface, MergeBehavior, Tool
-from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
 from vellum.workflows.types.generics import is_workflow_class
 
 CHAT_HISTORY_VARIABLE = "chat_history"
 
 
+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
 
@@ -41,10 +67,6 @@ class ToolRouterNode(InlinePromptNode[ToolCallingState]):
             max_iterations_message = f"Maximum number of prompt iterations `{self.max_prompt_iterations}` reached."
             raise NodeException(message=max_iterations_message, code=WorkflowErrorCode.NODE_EXECUTION)
 
-        # Merge user-provided chat history with node's chat history
-        user_chat_history = self.prompt_inputs.get(CHAT_HISTORY_VARIABLE, []) if self.prompt_inputs else []
-        merged_chat_history = user_chat_history + self.state.chat_history
-        self.prompt_inputs = {**self.prompt_inputs, CHAT_HISTORY_VARIABLE: merged_chat_history}  # type: ignore
         generator = super().run()
         for output in generator:
             if output.name == "results" and output.value:
@@ -72,20 +94,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
@@ -106,28 +119,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]
 
@@ -140,29 +141,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)
@@ -174,16 +162,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()
+
+        # HACK: Use first Composio API key found in environment variables
+        composio_api_key = None
+        common_env_var_names = ["COMPOSIO_API_KEY", "COMPOSIO_KEY"]
+
+        for env_var_name in common_env_var_names:
+            value = os.environ.get(env_var_name)
+            if value:
+                composio_api_key = value
+                break
+
+        if not composio_api_key:
+            raise NodeException(
+                message=(
+                    "No Composio API key found in environment variables. "
+                    "Please ensure one of these environment variables is set: "
+                )
+                + ", ".join(common_env_var_names),
+                code=WorkflowErrorCode.NODE_EXECUTION,
             )
-        )
+
+        try:
+            # Execute using ComposioService
+            composio_service = ComposioService(api_key=composio_api_key)
+            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 create_composio_wrapper_function(tool_def: ComposioToolDefinition):
+    """Create a real Python function that wraps the Composio tool for prompt layer compatibility."""
+
+    def wrapper_function(**kwargs):
+        # This should never be called due to routing, but satisfies introspection
+        raise RuntimeError(
+            f"ComposioToolDefinition wrapper for '{tool_def.action}' should not be called directly. "
+            f"Execution should go through ComposioNode. This suggests a routing issue."
+        )
+
+    # Set proper function attributes for prompt layer introspection
+    wrapper_function.__name__ = tool_def.name
+    wrapper_function.__doc__ = tool_def.description
+
+    return wrapper_function
+
+
 def create_tool_router_node(
     ml_model: str,
     blocks: List[PromptBlock],
@@ -193,9 +238,18 @@ 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 = []
+
         for function in functions:
+            # Convert ComposioToolDefinition to wrapper function for prompt layer
+            if isinstance(function, ComposioToolDefinition):
+                prompt_functions.append(create_composio_wrapper_function(function))
+            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
@@ -218,6 +272,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(
@@ -234,6 +289,14 @@ def create_tool_router_node(
             )
         )
 
+    node_prompt_inputs = {
+        **(prompt_inputs or {}),
+        CHAT_HISTORY_VARIABLE: ConcatExpression[List[ChatMessage], List[ChatMessage]](
+            lhs=(prompt_inputs or {}).get(CHAT_HISTORY_VARIABLE, []),
+            rhs=ToolCallingState.chat_history,
+        ),
+    }
+
     node = cast(
         Type[ToolRouterNode],
         type(
@@ -242,8 +305,8 @@ def create_tool_router_node(
             {
                 "ml_model": ml_model,
                 "blocks": blocks,
-                "functions": functions,
-                "prompt_inputs": prompt_inputs,
+                "functions": prompt_functions,  # Use converted functions for prompt layer
+                "prompt_inputs": node_prompt_inputs,
                 "parameters": parameters,
                 "max_prompt_iterations": max_prompt_iterations,
                 "Ports": Ports,
@@ -285,6 +348,17 @@ def create_function_node(
 
         return node
 
+    elif isinstance(function, ComposioToolDefinition):
+        node = type(
+            f"ComposioNode_{function.name}",
+            (ComposioNode,),
+            {
+                "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__}",
@@ -314,5 +388,7 @@ def get_function_name(function: Tool) -> str:
     if isinstance(function, DeploymentDefinition):
         name = str(function.deployment_id or function.deployment_name)
         return name.replace("-", "")
+    elif isinstance(function, ComposioToolDefinition):
+        return function.name
     else:
         return snake_case(function.__name__)

vellum/workflows/types/core.py

@@ -13,7 +13,7 @@ from typing import (  # type: ignore[attr-defined]
 )
 
 from vellum.client.core.pydantic_utilities import UniversalBaseModel
-from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
 
 if TYPE_CHECKING:
     from vellum.workflows.workflows.base import BaseWorkflow
@@ -50,4 +50,4 @@ class ConditionType(Enum):
 
 
 # Type alias for functions that can be called in tool calling nodes
-Tool = Union[Callable[..., Any], DeploymentDefinition, Type["BaseWorkflow"]]
+Tool = Union[Callable[..., Any], DeploymentDefinition, Type["BaseWorkflow"], ComposioToolDefinition]

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, Optional, Union
+from typing import Annotated, Any, Dict, Literal, Optional, Union
 
 from pydantic import BeforeValidator
 
@@ -97,3 +97,22 @@ class DeploymentDefinition(UniversalBaseModel):
         if not self._is_uuid():
             return self.deployment
         return None
+
+
+class ComposioToolDefinition(UniversalBaseModel):
+    """Represents a specific Composio action that can be used in Tool Calling Node"""
+
+    type: Literal["COMPOSIO"] = "COMPOSIO"
+
+    # Core identification
+    toolkit: str  # "GITHUB", "SLACK", etc.
+    action: str  # Specific action like "GITHUB_CREATE_AN_ISSUE"
+    description: str
+
+    # Optional cached metadata
+    display_name: Optional[str] = None
+
+    @property
+    def name(self) -> str:
+        """Generate a function name for this tool"""
+        return self.action.lower()

vellum/workflows/types/tests/test_definition.py

@@ -1,7 +1,7 @@
 import pytest
 from uuid import UUID
 
-from vellum.workflows.types.definition import DeploymentDefinition
+from vellum.workflows.types.definition import ComposioToolDefinition, DeploymentDefinition
 
 
 @pytest.mark.parametrize(
@@ -31,3 +31,16 @@ def test_deployment_definition(deployment_value, expected_deployment_id, expecte
 
     assert deployment.deployment_id == expected_deployment_id
     assert deployment.deployment_name == expected_deployment_name
+
+
+def test_composio_tool_definition_creation():
+    """Test that ComposioToolDefinition can be created with required fields."""
+    composio_tool = ComposioToolDefinition(
+        toolkit="GITHUB", action="GITHUB_CREATE_AN_ISSUE", description="Create a new issue in a GitHub repository"
+    )
+
+    assert composio_tool.toolkit == "GITHUB"
+    assert composio_tool.action == "GITHUB_CREATE_AN_ISSUE"
+    assert composio_tool.description == "Create a new issue in a GitHub repository"
+    assert composio_tool.display_name is None
+    assert composio_tool.name == "github_create_an_issue"

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.3.dist-info → vellum_ai-1.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vellum-ai
-Version: 1.0.3
+Version: 1.0.5
 Summary: 
 License: MIT
 Requires-Python: >=3.9,<4.0
@@ -22,6 +22,8 @@ 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)