feat: Pass in entity_id explicitly to auto-gen composio tool (#535)

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.

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']
     """

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:
     """