diff --git a/letta/constants.py b/letta/constants.py index 942dd0c9..5a51ef9b 100644 --- a/letta/constants.py +++ b/letta/constants.py @@ -7,6 +7,8 @@ ADMIN_PREFIX = "/v1/admin" API_PREFIX = "/v1" OPENAI_API_PREFIX = "/openai" +COMPOSIO_ENTITY_ENV_VAR_KEY = "COMPOSIO_ENTITY" + # String in the error message for when the context window is too large # Example full message: # This model's maximum context length is 8192 tokens. However, your messages resulted in 8198 tokens (7450 in the messages, 748 in the functions). Please reduce the length of the messages or functions. diff --git a/letta/functions/helpers.py b/letta/functions/helpers.py index d58efc46..fe86ad93 100644 --- a/letta/functions/helpers.py +++ b/letta/functions/helpers.py @@ -1,8 +1,11 @@ from typing import Any, Optional, Union import humps +from composio.constants import DEFAULT_ENTITY_ID from pydantic import BaseModel +from letta.constants import COMPOSIO_ENTITY_ENV_VAR_KEY + def generate_composio_tool_wrapper(action_name: str) -> tuple[str, str]: # Instantiate the object @@ -15,8 +18,10 @@ def generate_composio_tool_wrapper(action_name: str) -> tuple[str, str]: def {func_name}(**kwargs): from composio import Action, App, Tag from composio_langchain import ComposioToolSet + import os - composio_toolset = ComposioToolSet() + entity_id = os.getenv('{COMPOSIO_ENTITY_ENV_VAR_KEY}', '{DEFAULT_ENTITY_ID}') + composio_toolset = ComposioToolSet(entity_id=entity_id) tool = {tool_instantiation_str} return tool.func(**kwargs)['data'] """ diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 04785c64..0f34110e 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -8,6 +8,7 @@ import pytest from sqlalchemy import delete from letta import create_client +from letta.constants import COMPOSIO_ENTITY_ENV_VAR_KEY from letta.functions.function_sets.base import core_memory_append, core_memory_replace from letta.orm.sandbox_config import SandboxConfig, SandboxEnvironmentVariable from letta.schemas.agent import AgentState @@ -186,6 +187,14 @@ def composio_github_star_tool(test_user): yield tool +@pytest.fixture +def composio_gmail_get_profile_tool(test_user): + tool_manager = ToolManager() + tool_create = ToolCreate.from_composio(action_name="GMAIL_GET_PROFILE") + tool = tool_manager.create_or_update_tool(pydantic_tool=Tool(**tool_create.model_dump()), actor=test_user) + yield tool + + @pytest.fixture def clear_core_memory_tool(test_user): def clear_memory(agent_state: AgentState): @@ -374,6 +383,29 @@ def test_local_sandbox_e2e_composio_star_github(mock_e2b_api_key_none, check_com assert result.func_return["details"] == "Action executed successfully" +@pytest.mark.local_sandbox +def test_local_sandbox_multiple_composio_entities( + mock_e2b_api_key_none, check_composio_key_set, composio_gmail_get_profile_tool, agent_state, test_user +): + # Agent state with no composio entity ID + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com" + + # Agent state with the composio entity set to 'matt' + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="matt", agent_id=agent_state.id) + ] + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "matt@letta.com" + + # Agent state with composio entity ID set to default + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="default", agent_id=agent_state.id) + ] + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com" + + @pytest.mark.local_sandbox def test_local_sandbox_e2e_composio_star_github_without_setting_db_env_vars( mock_e2b_api_key_none, check_composio_key_set, composio_github_star_tool, test_user @@ -593,6 +625,38 @@ def test_e2b_e2e_composio_star_github(check_e2b_key_is_set, check_composio_key_s assert result.func_return["details"] == "Action executed successfully" +@pytest.mark.e2b_sandbox +def test_e2b_multiple_composio_entities( + check_e2b_key_is_set, check_composio_key_set, composio_gmail_get_profile_tool, agent_state, test_user +): + manager = SandboxConfigManager(tool_settings) + config = manager.get_or_create_default_sandbox_config(sandbox_type=SandboxType.E2B, actor=test_user) + + manager.create_sandbox_env_var( + SandboxEnvironmentVariableCreate(key="COMPOSIO_API_KEY", value=tool_settings.composio_api_key), + sandbox_config_id=config.id, + actor=test_user, + ) + + # Agent state with no composio entity ID + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com" + + # Agent state with the composio entity set to 'matt' + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="matt", agent_id=agent_state.id) + ] + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "matt@letta.com" + + # Agent state with composio entity ID set to default + agent_state.tool_exec_environment_variables = [ + AgentEnvironmentVariable(key=COMPOSIO_ENTITY_ENV_VAR_KEY, value="default", agent_id=agent_state.id) + ] + result = ToolExecutionSandbox(composio_gmail_get_profile_tool.name, {}, user=test_user).run(agent_state=agent_state) + assert result.func_return["response_data"]["emailAddress"] == "sarah@letta.com" + + # Core memory integration tests class TestCoreMemoryTools: """