vellum-ai 1.0.4__tar.gz → 1.0.5__tar.gz

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/PKG-INFO

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

vellum_ai-1.0.5/ee/vellum_ee/workflows/display/tests/workflow_serialization/test_basic_tool_calling_node_composio_serialization.py

@@ -0,0 +1,86 @@
+from vellum_ee.workflows.display.workflows.get_vellum_workflow_display_class import get_workflow_display
+
+from tests.workflows.basic_tool_calling_node_with_composio_tool.workflow import (
+    BasicToolCallingNodeWithComposioToolWorkflow,
+)
+
+
+def test_serialize_workflow():
+    # GIVEN a Workflow that uses a tool calling node with a composio tool
+    # WHEN we serialize it
+    workflow_display = get_workflow_display(workflow_class=BasicToolCallingNodeWithComposioToolWorkflow)
+
+    serialized_workflow: dict = workflow_display.serialize()
+
+    # THEN we should get a serialized representation of the Workflow
+    assert serialized_workflow.keys() == {
+        "workflow_raw_data",
+        "input_variables",
+        "state_variables",
+        "output_variables",
+    }
+
+    # AND its input variables should be what we expect
+    input_variables = serialized_workflow["input_variables"]
+    assert len(input_variables) == 1
+
+    # AND its output variables should be what we expect
+    output_variables = serialized_workflow["output_variables"]
+    assert len(output_variables) == 2
+
+    # Find the text and chat_history outputs
+    text_output = next(var for var in output_variables if var["key"] == "text")
+    chat_history_output = next(var for var in output_variables if var["key"] == "chat_history")
+
+    assert text_output["type"] == "STRING"
+    assert chat_history_output["type"] == "CHAT_HISTORY"
+
+    # AND its raw data should be what we expect
+    workflow_raw_data = serialized_workflow["workflow_raw_data"]
+    tool_calling_node = workflow_raw_data["nodes"][1]
+
+    # AND the tool calling node should have the composio tool properly serialized
+    functions_attribute = next(attr for attr in tool_calling_node["attributes"] if attr["name"] == "functions")
+    assert functions_attribute["value"]["type"] == "CONSTANT_VALUE"
+    assert functions_attribute["value"]["value"]["type"] == "JSON"
+
+    functions_list = functions_attribute["value"]["value"]["value"]
+    assert len(functions_list) == 1
+
+    composio_function = functions_list[0]
+    assert composio_function == {
+        "type": "COMPOSIO",
+        "toolkit": "GITHUB",
+        "action": "GITHUB_CREATE_AN_ISSUE",
+        "description": "Create a new issue in a GitHub repository",
+        "display_name": "Create GitHub Issue",
+    }
+
+    # AND the rest of the node structure should be correct
+    assert tool_calling_node["type"] == "GENERIC"
+    assert tool_calling_node["base"]["name"] == "ToolCallingNode"
+    assert tool_calling_node["definition"]["name"] == "ComposioToolCallingNode"
+
+    # AND the blocks should be properly serialized
+    blocks_attribute = next(attr for attr in tool_calling_node["attributes"] if attr["name"] == "blocks")
+    assert blocks_attribute["value"]["type"] == "CONSTANT_VALUE"
+    blocks_list = blocks_attribute["value"]["value"]["value"]
+    assert len(blocks_list) == 2
+    assert blocks_list[0]["chat_role"] == "SYSTEM"
+    assert blocks_list[1]["chat_role"] == "USER"
+
+    # AND the prompt inputs should be properly serialized
+    prompt_inputs_attribute = next(attr for attr in tool_calling_node["attributes"] if attr["name"] == "prompt_inputs")
+    assert prompt_inputs_attribute["value"]["type"] == "DICTIONARY_REFERENCE"
+    entries = prompt_inputs_attribute["value"]["entries"]
+    assert len(entries) == 1
+    assert entries[0]["key"] == "question"
+    assert entries[0]["value"]["type"] == "WORKFLOW_INPUT"
+
+    # AND the outputs should be correct
+    outputs = tool_calling_node["outputs"]
+    assert len(outputs) == 2
+    assert outputs[0]["name"] == "text"
+    assert outputs[0]["type"] == "STRING"
+    assert outputs[1]["name"] == "chat_history"
+    assert outputs[1]["type"] == "CHAT_HISTORY"

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/pyproject.toml

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

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/src/vellum/__init__.py

@@ -217,6 +217,7 @@ from .client.types import (
     MlModelRead,
     MlModelUsage,
     MlModelUsageWrapper,
+    NameEnum,
     NamedScenarioInputChatHistoryVariableValueRequest,
     NamedScenarioInputJsonVariableValueRequest,
     NamedScenarioInputRequest,
@@ -867,6 +868,7 @@ __all__ = [
     "MlModelRead",
     "MlModelUsage",
     "MlModelUsageWrapper",
+    "NameEnum",
     "NamedScenarioInputChatHistoryVariableValueRequest",
     "NamedScenarioInputJsonVariableValueRequest",
     "NamedScenarioInputRequest",

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/src/vellum/client/core/client_wrapper.py

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

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/src/vellum/client/types/__init__.py

@@ -225,6 +225,7 @@ from .metric_node_result import MetricNodeResult
 from .ml_model_read import MlModelRead
 from .ml_model_usage import MlModelUsage
 from .ml_model_usage_wrapper import MlModelUsageWrapper
+from .name_enum import NameEnum
 from .named_scenario_input_chat_history_variable_value_request import NamedScenarioInputChatHistoryVariableValueRequest
 from .named_scenario_input_json_variable_value_request import NamedScenarioInputJsonVariableValueRequest
 from .named_scenario_input_request import NamedScenarioInputRequest
@@ -847,6 +848,7 @@ __all__ = [
     "MlModelRead",
     "MlModelUsage",
     "MlModelUsageWrapper",
+    "NameEnum",
     "NamedScenarioInputChatHistoryVariableValueRequest",
     "NamedScenarioInputJsonVariableValueRequest",
     "NamedScenarioInputRequest",

vellum_ai-1.0.5/src/vellum/client/types/name_enum.py

@@ -0,0 +1,7 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import typing
+
+NameEnum = typing.Union[
+    typing.Literal["workflow_executions", "prompt_executions", "workflow_runtime_seconds"], typing.Any
+]

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/src/vellum/client/types/quota.py

@@ -1,13 +1,14 @@
 # This file was auto-generated by Fern from our API Definition.
 
 from ..core.pydantic_utilities import UniversalBaseModel
+from .name_enum import NameEnum
 import typing
 from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import pydantic
 
 
 class Quota(UniversalBaseModel):
-    name: str
+    name: NameEnum
     value: typing.Optional[int] = None
     period_seconds: typing.Optional[int] = None
 

vellum_ai-1.0.5/src/vellum/types/name_enum.py

@@ -0,0 +1,3 @@
+# WARNING: This file will be removed in a future release. Please import from "vellum.client" instead.
+
+from vellum.client.types.name_enum import *

{vellum_ai-1.0.4 → vellum_ai-1.0.5}/src/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_ai-1.0.4 → vellum_ai-1.0.5}/src/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
@@ -18,6 +19,7 @@ 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
@@ -31,6 +33,29 @@ 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
 
@@ -69,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
@@ -103,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]
 
@@ -137,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)
@@ -171,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],
@@ -190,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
@@ -215,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(
@@ -247,7 +305,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 +349,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_ai-1.0.4 → vellum_ai-1.0.5}/src/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_ai-1.0.4 → vellum_ai-1.0.5}/src/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]}