From 86023db9b1c9b2e938bb1c854edbee6c413b54d9 Mon Sep 17 00:00:00 2001 From: Sarah Wooders Date: Sun, 30 Nov 2025 00:32:43 -0800 Subject: [PATCH] feat: get rid of client injection in args and just make it always available (#6459) --- letta/constants.py | 2 +- letta/services/tool_sandbox/base.py | 26 +++++++------ tests/integration_test_async_tool_sandbox.py | 39 +++++++++---------- ...integration_test_tool_execution_sandbox.py | 12 +++--- tests/test_tool_schema_parsing.py | 39 +------------------ 5 files changed, 42 insertions(+), 76 deletions(-) diff --git a/letta/constants.py b/letta/constants.py index 72ec844e..4fa9aa1f 100644 --- a/letta/constants.py +++ b/letta/constants.py @@ -469,4 +469,4 @@ DEFAULT_GENERATE_TOOL_MODEL_HANDLE = "openai/gpt-4.1" # Reserved keyword arguments that are injected by the system into tool functions, not provided by the LLM # These parameters are excluded from tool schema generation -TOOL_RESERVED_KWARGS = ["self", "agent_state", "client"] +TOOL_RESERVED_KWARGS = ["self", "agent_state"] diff --git a/letta/services/tool_sandbox/base.py b/letta/services/tool_sandbox/base.py index ab1da864..12f14df0 100644 --- a/letta/services/tool_sandbox/base.py +++ b/letta/services/tool_sandbox/base.py @@ -72,8 +72,8 @@ class AsyncToolSandboxBase(ABC): else: self.inject_agent_state = False - # Check for Letta client and agent_id injection - self.inject_letta_client = "client" in tool_arguments + # Always inject Letta client (available as `client` variable in sandbox) + self.inject_letta_client = True self.inject_agent_id = "agent_id" in tool_arguments self.is_async_function = self._detect_async_function() @@ -183,9 +183,16 @@ class AsyncToolSandboxBase(ABC): if inject_agent_state: lines.extend(["import letta", "from letta import *"]) # noqa: F401 - # Import Letta client if needed + # Import Letta client if available (wrapped in try/except for sandboxes without letta_client installed) if inject_letta_client: - lines.append("from letta_client import Letta") + lines.extend( + [ + "try:", + " from letta_client import Letta", + "except ImportError:", + " Letta = None", + ] + ) if schema_imports: lines.append(schema_imports.rstrip()) @@ -195,14 +202,14 @@ class AsyncToolSandboxBase(ABC): else: lines.append("agent_state = None") - # Initialize Letta client if needed + # Initialize Letta client if needed (client is always available as a variable, may be None) if inject_letta_client: lines.extend( [ "# Initialize Letta client for tool execution", "import os", "client = None", - "if os.getenv('LETTA_API_KEY'):", + "if Letta is not None and os.getenv('LETTA_API_KEY'):", " # Check letta_client version to use correct parameter name", " from packaging import version as pkg_version", " import letta_client as lc_module", @@ -346,11 +353,8 @@ class AsyncToolSandboxBase(ABC): if self.inject_agent_state: param_list.append("agent_state=agent_state") - if self.inject_letta_client: - # Check if the function expects 'client' or 'letta_client' - tool_arguments = parse_function_arguments(self.tool.source_code, self.tool.name) - if "client" in tool_arguments: - param_list.append("client=client") + # Note: client is always available as a variable in the sandbox scope + # Tools should access it directly rather than receiving it as a parameter if self.inject_agent_id: param_list.append("agent_id=agent_id") diff --git a/tests/integration_test_async_tool_sandbox.py b/tests/integration_test_async_tool_sandbox.py index 53a42308..2d84bb39 100644 --- a/tests/integration_test_async_tool_sandbox.py +++ b/tests/integration_test_async_tool_sandbox.py @@ -1217,28 +1217,25 @@ async def test_e2b_sandbox_async_per_agent_env(check_e2b_key_is_set, async_get_e @pytest.fixture async def list_tools_with_client_tool(test_user): - """Tool that uses injected client to list tools. + """Tool that uses the client (available in sandbox scope) to list tools. - Note: This fixture uses ToolManager directly instead of the client API - because it needs to create a tool with a custom schema that excludes - the 'client' parameter (which is injected by the sandbox, not passed by the LLM). + Note: The `client` variable is always available in the sandbox scope, + so tools can access it directly without declaring it as a parameter. """ from letta.schemas.enums import ToolType from letta.schemas.tool import Tool as PydanticTool source_code = ''' -def list_tools_via_client(client: "Letta") -> str: +def list_tools_via_client() -> str: """ - List available tools using the injected Letta client. - - Args: - client: Letta client instance (injected by system) + List available tools using the Letta client available in sandbox scope. Returns: str: Comma-separated list of tool names """ + # `client` is always available in the sandbox scope if not client: - return "ERROR: client not injected" + return "ERROR: client not available in scope" try: tools = client.tools.list() @@ -1248,19 +1245,19 @@ def list_tools_via_client(client: "Letta") -> str: return f"ERROR: {str(e)}" ''' - # Create the tool with proper schema (client is injected, not in schema) + # Create the tool with proper schema tool = PydanticTool( name="list_tools_via_client", - description="List tools using injected client", + description="List tools using client available in sandbox scope", source_code=source_code, source_type="python", tool_type=ToolType.CUSTOM, ) - # Manually set schema without 'client' parameter since it's injected + # Schema has no parameters since client is accessed from scope, not passed as arg tool.json_schema = { "name": "list_tools_via_client", - "description": "List tools using injected client", + "description": "List tools using client available in sandbox scope", "parameters": {"type": "object", "properties": {}, "required": []}, } @@ -1296,8 +1293,8 @@ async def test_local_sandbox_with_client_injection(disable_e2b_api_key, list_too await sandbox._init_async() - # Verify that client injection was detected - assert sandbox.inject_letta_client is True, "Tool should be detected as needing client injection" + # Verify that client injection is enabled (always True now) + assert sandbox.inject_letta_client is True, "Client injection should always be enabled" # Generate the execution script to verify client initialization code is present script = await sandbox.generate_execution_script(agent_state=None) @@ -1324,8 +1321,10 @@ async def test_local_sandbox_with_client_injection(disable_e2b_api_key, list_too print(result) assert "Found" in str(result.func_return), f"Tool should list tools when client is available: {result.func_return}" - # Verify client was injected successfully (connection may fail if no server is running) - assert "ERROR: client not injected" not in str(result.func_return), "Client should be injected when LETTA_API_KEY is set" + # Verify client was available in scope (connection may fail if no server is running) + assert "ERROR: client not available in scope" not in str(result.func_return), ( + "Client should be available in scope when LETTA_API_KEY is set" + ) @pytest.mark.asyncio @@ -1355,8 +1354,8 @@ async def test_e2b_sandbox_with_client_injection(check_e2b_key_is_set, list_tool await sandbox._init_async() - # Verify that client injection was detected - assert sandbox.inject_letta_client is True, "Tool should be detected as needing client injection" + # Verify that client injection is enabled (always True now) + assert sandbox.inject_letta_client is True, "Client injection should always be enabled" # Generate the execution script to verify client initialization code is present script = await sandbox.generate_execution_script(agent_state=None) diff --git a/tests/integration_test_tool_execution_sandbox.py b/tests/integration_test_tool_execution_sandbox.py index 14ff486d..4e8fa07b 100644 --- a/tests/integration_test_tool_execution_sandbox.py +++ b/tests/integration_test_tool_execution_sandbox.py @@ -406,22 +406,22 @@ async def test_tool_with_client_injection(disable_e2b_api_key, server: SyncServe """Test that tools can access injected letta_client and agent_id to modify agent blocks.""" # Create a tool that uses the injected client and agent_id to actually clear a memory block + # Note: `client` is always available as a variable in the sandbox scope memory_clear_source = ''' -def memory_clear(label: str, agent_id: str, client: "Letta"): +def memory_clear(label: str, agent_id: str): """Test tool that clears a memory block using the injected client. Args: label: The label of the memory block to clear agent_id: The agent's ID (injected by Letta system) - client: The Letta client instance (injected by Letta system) """ # Verify that agent_id was injected if not agent_id or not isinstance(agent_id, str): return f"ERROR: agent_id not properly injected: {agent_id}" - # Verify that client was injected + # Verify that client is available in scope (always injected) if not client or not hasattr(client, 'agents'): - return f"ERROR: client not properly injected: {client}" + return f"ERROR: client not available in scope: {client}" # Use the injected client to actually clear the memory block try: @@ -504,8 +504,8 @@ def memory_clear(label: str, agent_id: str, client: "Letta"): # Initialize the sandbox to detect reserved keywords await sandbox._init_async() - # Verify that the tool correctly detects the need for injection - assert sandbox.inject_letta_client == True # Should detect 'client' parameter + # Verify that injection is configured correctly + assert sandbox.inject_letta_client == True # Client is always injected assert sandbox.inject_agent_id == True # Should detect 'agent_id' parameter # Generate the execution script to verify injection code is present diff --git a/tests/test_tool_schema_parsing.py b/tests/test_tool_schema_parsing.py index bdaa4817..1083982f 100644 --- a/tests/test_tool_schema_parsing.py +++ b/tests/test_tool_schema_parsing.py @@ -401,30 +401,6 @@ def agent_state_ok(agent_state, value: int) -> str: return "ok" -def client_ok(client, value: int) -> str: - """Ignores client param (injected Letta client). - - Args: - value (int): Some value. - - Returns: - str: Status. - """ - return "ok" - - -def all_reserved_params_ok(agent_state, client, value: int) -> str: - """Ignores all reserved params. - - Args: - value (int): Some value. - - Returns: - str: Status. - """ - return "ok" - - class Dummy: def method(self, bar: int) -> str: # keeps an explicit self """Bound-method example. @@ -470,8 +446,6 @@ def missing_param_doc(x: int, y: int) -> str: [ (good_function, None), (agent_state_ok, None), - (client_ok, None), # client is a reserved param (injected Letta client) - (all_reserved_params_ok, None), # all reserved params together (Dummy.method, None), # unbound method keeps `self` (good_function_no_return, None), (no_doc, "has no docstring"), @@ -484,24 +458,13 @@ def test_google_style_docstring_validation(fn, regex): def test_reserved_params_excluded_from_schema(): - """Test that reserved params (agent_state, client) are excluded from generated schema.""" + """Test that reserved params (agent_state) are excluded from generated schema.""" from letta.functions.schema_generator import generate_schema - # Test with client param - schema = generate_schema(client_ok) - assert "client" not in schema["parameters"]["properties"], "client should be excluded from schema" - assert "value" in schema["parameters"]["properties"], "value should be in schema" - # Test with agent_state param schema = generate_schema(agent_state_ok) assert "agent_state" not in schema["parameters"]["properties"], "agent_state should be excluded from schema" assert "value" in schema["parameters"]["properties"], "value should be in schema" - - # Test with all reserved params - schema = generate_schema(all_reserved_params_ok) - assert "agent_state" not in schema["parameters"]["properties"], "agent_state should be excluded from schema" - assert "client" not in schema["parameters"]["properties"], "client should be excluded from schema" - assert "value" in schema["parameters"]["properties"], "value should be in schema" assert schema["parameters"]["required"] == ["value"], "only value should be required"