diff --git a/examples/composio_tool_usage.py b/examples/composio_tool_usage.py index 26508fcd..877d3754 100644 --- a/examples/composio_tool_usage.py +++ b/examples/composio_tool_usage.py @@ -1,10 +1,14 @@ import json +import os import uuid from letta import create_client from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.llm_config import LLMConfig from letta.schemas.memory import ChatMemory +from letta.schemas.sandbox_config import SandboxEnvironmentVariableCreate, SandboxType +from letta.services.sandbox_config_manager import SandboxConfigManager +from letta.settings import tool_settings """ Setup here. @@ -25,6 +29,17 @@ for agent_state in client.list_agents(): print(f"Deleted agent: {agent_state.name} with ID {str(agent_state.id)}") +# Add sandbox env +manager = SandboxConfigManager(tool_settings) +# Ensure you have e2b key set +sandbox_config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=client.user) +manager.create_sandbox_env_var( + SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=os.environ.get("COMPOSIO_API_KEY")), + sandbox_config_id=sandbox_config.id, + actor=client.user, +) + + """ This example show how you can add Composio tools . diff --git a/letta/client/client.py b/letta/client/client.py index 485ce300..b56bc19f 100644 --- a/letta/client/client.py +++ b/letta/client/client.py @@ -2683,7 +2683,7 @@ class LocalClient(AbstractClient): return self.server.tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=self.user) def load_composio_tool(self, action: "ActionType") -> Tool: - tool_create = ToolCreate.from_composio(action=action) + tool_create = ToolCreate.from_composio(action_name=action.name) return self.server.tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=self.user) def create_tool( diff --git a/letta/functions/helpers.py b/letta/functions/helpers.py index d94eb7da..d58efc46 100644 --- a/letta/functions/helpers.py +++ b/letta/functions/helpers.py @@ -4,12 +4,12 @@ import humps from pydantic import BaseModel -def generate_composio_tool_wrapper(action: "ActionType") -> tuple[str, str]: +def generate_composio_tool_wrapper(action_name: str) -> tuple[str, str]: # Instantiate the object - tool_instantiation_str = f"composio_toolset.get_tools(actions=[Action.{str(action)}])[0]" + tool_instantiation_str = f"composio_toolset.get_tools(actions=['{action_name}'])[0]" # Generate func name - func_name = action.name.lower() + func_name = action_name.lower() wrapper_function_str = f""" def {func_name}(**kwargs): diff --git a/letta/schemas/tool.py b/letta/schemas/tool.py index ba844551..f02cf1fd 100644 --- a/letta/schemas/tool.py +++ b/letta/schemas/tool.py @@ -93,7 +93,7 @@ class ToolCreate(LettaBase): ) @classmethod - def from_composio(cls, action: "ActionType") -> "ToolCreate": + def from_composio(cls, action_name: str, api_key: Optional[str] = None) -> "ToolCreate": """ Class method to create an instance of Letta-compatible Composio Tool. Check https://docs.composio.dev/introduction/intro/overview to look at options for from_composio @@ -101,15 +101,20 @@ class ToolCreate(LettaBase): This function will error if we find more than one tool, or 0 tools. Args: - action ActionType: A action name to filter tools by. + action_name str: A action name to filter tools by. Returns: Tool: A Letta Tool initialized with attributes derived from the Composio tool. """ from composio import LogLevel from composio_langchain import ComposioToolSet - composio_toolset = ComposioToolSet(logging_level=LogLevel.ERROR) - composio_tools = composio_toolset.get_tools(actions=[action]) + if api_key: + # Pass in an external API key + composio_toolset = ComposioToolSet(logging_level=LogLevel.ERROR, api_key=api_key) + else: + # Use environmental variable + composio_toolset = ComposioToolSet(logging_level=LogLevel.ERROR) + composio_tools = composio_toolset.get_tools(actions=[action_name]) assert len(composio_tools) > 0, "User supplied parameters do not match any Composio tools" assert len(composio_tools) == 1, f"User supplied parameters match too many Composio tools; {len(composio_tools)} > 1" @@ -119,7 +124,7 @@ class ToolCreate(LettaBase): description = composio_tool.description source_type = "python" tags = ["composio"] - wrapper_func_name, wrapper_function_str = generate_composio_tool_wrapper(action) + wrapper_func_name, wrapper_function_str = generate_composio_tool_wrapper(action_name) json_schema = generate_schema_from_args_schema_v2(composio_tool.args_schema, name=wrapper_func_name, description=description) return cls( @@ -179,10 +184,10 @@ class ToolCreate(LettaBase): def load_default_composio_tools(cls) -> List["ToolCreate"]: from composio_langchain import Action - calculator = ToolCreate.from_composio(action=Action.MATHEMATICAL_CALCULATOR) - serp_news = ToolCreate.from_composio(action=Action.SERPAPI_NEWS_SEARCH) - serp_google_search = ToolCreate.from_composio(action=Action.SERPAPI_SEARCH) - serp_google_maps = ToolCreate.from_composio(action=Action.SERPAPI_GOOGLE_MAPS_SEARCH) + calculator = ToolCreate.from_composio(action_name=Action.MATHEMATICAL_CALCULATOR.name) + serp_news = ToolCreate.from_composio(action_name=Action.SERPAPI_NEWS_SEARCH.name) + serp_google_search = ToolCreate.from_composio(action_name=Action.SERPAPI_SEARCH.name) + serp_google_maps = ToolCreate.from_composio(action_name=Action.SERPAPI_GOOGLE_MAPS_SEARCH.name) return [calculator, serp_news, serp_google_search, serp_google_maps] diff --git a/letta/server/rest_api/routers/v1/tools.py b/letta/server/rest_api/routers/v1/tools.py index c9b60f78..1b5e2eac 100644 --- a/letta/server/rest_api/routers/v1/tools.py +++ b/letta/server/rest_api/routers/v1/tools.py @@ -7,6 +7,7 @@ from letta.errors import LettaToolCreateError from letta.orm.errors import UniqueConstraintViolationError from letta.schemas.letta_message import FunctionReturn from letta.schemas.tool import Tool, ToolCreate, ToolRunFromSource, ToolUpdate +from letta.schemas.user import User from letta.server.rest_api.utils import get_letta_server from letta.server.server import SyncServer @@ -213,22 +214,27 @@ def run_tool_from_source( @router.get("/composio/apps", response_model=List[AppModel], operation_id="list_composio_apps") -def list_composio_apps(server: SyncServer = Depends(get_letta_server)): +def list_composio_apps(server: SyncServer = Depends(get_letta_server), user_id: Optional[str] = Header(None, alias="user_id")): """ Get a list of all Composio apps """ - return server.get_composio_apps() + actor = server.get_user_or_default(user_id=user_id) + composio_api_key = get_composio_key(server, actor=actor) + return server.get_composio_apps(api_key=composio_api_key) @router.get("/composio/apps/{composio_app_name}/actions", response_model=List[ActionModel], operation_id="list_composio_actions_by_app") def list_composio_actions_by_app( composio_app_name: str, server: SyncServer = Depends(get_letta_server), + user_id: Optional[str] = Header(None, alias="user_id"), ): """ Get a list of all Composio actions for a specific app """ - return server.get_composio_actions_from_app_name(composio_app_name=composio_app_name) + actor = server.get_user_or_default(user_id=user_id) + composio_api_key = get_composio_key(server, actor=actor) + return server.get_composio_actions_from_app_name(composio_app_name=composio_app_name, api_key=composio_api_key) @router.post("/composio/{composio_action_name}", response_model=Tool, operation_id="add_composio_tool") @@ -241,5 +247,21 @@ def add_composio_tool( Add a new Composio tool by action name (Composio refers to each tool as an `Action`) """ actor = server.get_user_or_default(user_id=user_id) - tool_create = ToolCreate.from_composio(action=composio_action_name) + composio_api_key = get_composio_key(server, actor=actor) + tool_create = ToolCreate.from_composio(action_name=composio_action_name, api_key=composio_api_key) return server.tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=actor) + + +# TODO: Factor this out to somewhere else +def get_composio_key(server: SyncServer, actor: User): + api_keys = server.sandbox_config_manager.list_sandbox_env_vars_by_key(key="COMPOSIO_API_KEY", actor=actor) + if not api_keys: + raise HTTPException( + status_code=400, # Bad Request + detail=f"No API keys found for Composio. Please add your Composio API Key as an environment variable for your sandbox configuration.", + ) + + # TODO: Add more protections around this + # Ideally, not tied to a specific sandbox, but for now we just get the first one + # Theoretically possible for someone to have different composio api keys per sandbox + return api_keys[0].value diff --git a/letta/server/server.py b/letta/server/server.py index 7a430862..fa3d29fa 100644 --- a/letta/server/server.py +++ b/letta/server/server.py @@ -236,11 +236,6 @@ class SyncServer(Server): # Locks self.send_message_lock = Lock() - # Composio - self.composio_client = None - if tool_settings.composio_api_key: - self.composio_client = Composio(api_key=tool_settings.composio_api_key) - # Initialize the metadata store config = LettaConfig.load() if settings.letta_pg_uri_no_default: @@ -1899,9 +1894,17 @@ class SyncServer(Server): ) # Composio wrappers - def get_composio_apps(self) -> List["AppModel"]: + def get_composio_client(self, api_key: Optional[str] = None): + if api_key: + return Composio(api_key=api_key) + elif tool_settings.composio_api_key: + return Composio(api_key=tool_settings.composio_api_key) + else: + return Composio() + + def get_composio_apps(self, api_key: Optional[str] = None) -> List["AppModel"]: """Get a list of all Composio apps with actions""" - apps = self.composio_client.apps.get() + apps = self.get_composio_client(api_key=api_key).apps.get() apps_with_actions = [] for app in apps: # A bit of hacky logic until composio patches this @@ -1910,6 +1913,6 @@ class SyncServer(Server): return apps_with_actions - def get_composio_actions_from_app_name(self, composio_app_name: str) -> List["ActionModel"]: - actions = self.composio_client.actions.get(apps=[composio_app_name]) + def get_composio_actions_from_app_name(self, composio_app_name: str, api_key: Optional[str] = None) -> List["ActionModel"]: + actions = self.get_composio_client(api_key=api_key).actions.get(apps=[composio_app_name]) return actions diff --git a/letta/services/sandbox_config_manager.py b/letta/services/sandbox_config_manager.py index 48e53f9f..05b669d8 100644 --- a/letta/services/sandbox_config_manager.py +++ b/letta/services/sandbox_config_manager.py @@ -225,6 +225,21 @@ class SandboxConfigManager: ) return [env_var.to_pydantic() for env_var in env_vars] + @enforce_types + def list_sandbox_env_vars_by_key( + self, key: str, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50 + ) -> List[PydanticEnvVar]: + """List all sandbox environment variables with optional pagination.""" + with self.session_maker() as session: + env_vars = SandboxEnvVarModel.list( + db_session=session, + cursor=cursor, + limit=limit, + organization_id=actor.organization_id, + key=key, + ) + return [env_var.to_pydantic() for env_var in env_vars] + @enforce_types def get_sandbox_env_vars_as_dict( self, sandbox_config_id: str, actor: PydanticUser, cursor: Optional[str] = None, limit: Optional[int] = 50 diff --git a/letta/services/tool_manager.py b/letta/services/tool_manager.py index 6e1818e3..63240930 100644 --- a/letta/services/tool_manager.py +++ b/letta/services/tool_manager.py @@ -40,8 +40,6 @@ class ToolManager: if tool: # Put to dict and remove fields that should not be reset update_data = pydantic_tool.model_dump(exclude={"module"}, exclude_unset=True, exclude_none=True) - # Remove redundant update fields - update_data = {key: value for key, value in update_data.items() if getattr(tool, key) != value} # If there's anything to update if update_data: @@ -108,7 +106,7 @@ class ToolManager: tool = ToolModel.read(db_session=session, identifier=tool_id, actor=actor) # Update tool attributes with only the fields that were explicitly set - update_data = tool_update.model_dump(exclude_unset=True, exclude_none=True) + update_data = tool_update.model_dump(exclude_none=True) for key, value in update_data.items(): setattr(tool, key, value) diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 0574e43c..e13275b2 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -5,7 +5,6 @@ from pathlib import Path from unittest.mock import patch import pytest -from composio import Action from sqlalchemy import delete from letta import create_client @@ -200,7 +199,7 @@ def list_tool(test_user): @pytest.fixture def composio_github_star_tool(test_user): tool_manager = ToolManager() - tool_create = ToolCreate.from_composio(action=Action.GITHUB_STAR_A_REPOSITORY_FOR_THE_AUTHENTICATED_USER) + tool_create = ToolCreate.from_composio(action_name="GITHUB_STAR_A_REPOSITORY_FOR_THE_AUTHENTICATED_USER") tool = tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=test_user) yield tool diff --git a/tests/test_v1_routes.py b/tests/test_v1_routes.py index 883395e1..8f9d9972 100644 --- a/tests/test_v1_routes.py +++ b/tests/test_v1_routes.py @@ -162,6 +162,16 @@ def composio_actions(): ] +def configure_mock_sync_server(mock_sync_server): + # Mock sandbox config manager to return a valid API key + mock_api_key = Mock() + mock_api_key.value = "mock_composio_api_key" + mock_sync_server.sandbox_config_manager.list_sandbox_env_vars_by_key.return_value = [mock_api_key] + + # Mock user retrieval + mock_sync_server.get_user_or_default.return_value = Mock() # Provide additional attributes if needed + + # ====================================================================================================================== # Tools Routes Tests # ====================================================================================================================== @@ -274,6 +284,8 @@ def test_add_base_tools(client, mock_sync_server, add_integers_tool): def test_list_composio_apps(client, mock_sync_server, composio_apps): + configure_mock_sync_server(mock_sync_server) + mock_sync_server.get_composio_apps.return_value = composio_apps response = client.get("/v1/tools/composio/apps") @@ -284,16 +296,20 @@ def test_list_composio_apps(client, mock_sync_server, composio_apps): def test_list_composio_actions_by_app(client, mock_sync_server, composio_actions): + configure_mock_sync_server(mock_sync_server) + mock_sync_server.get_composio_actions_from_app_name.return_value = composio_actions response = client.get("/v1/tools/composio/apps/App1/actions") assert response.status_code == 200 assert len(response.json()) == 1 - mock_sync_server.get_composio_actions_from_app_name.assert_called_once_with(composio_app_name="App1") + mock_sync_server.get_composio_actions_from_app_name.assert_called_once_with(composio_app_name="App1", api_key="mock_composio_api_key") def test_add_composio_tool(client, mock_sync_server, add_integers_tool): + configure_mock_sync_server(mock_sync_server) + # Mock ToolCreate.from_composio to return the expected ToolCreate object with patch("letta.schemas.tool.ToolCreate.from_composio") as mock_from_composio: mock_from_composio.return_value = ToolCreate( @@ -314,4 +330,4 @@ def test_add_composio_tool(client, mock_sync_server, add_integers_tool): mock_sync_server.tool_manager.create_or_update_tool.assert_called_once() # Verify the mocked from_composio method was called - mock_from_composio.assert_called_once_with(action=add_integers_tool.name) + mock_from_composio.assert_called_once_with(action_name=add_integers_tool.name, api_key="mock_composio_api_key")