feat: get rid of client injection in args and just make it always available (#6459)

This commit is contained in:
Sarah Wooders
2025-11-30 00:32:43 -08:00
committed by Caren Thomas
parent f417e53638
commit 86023db9b1
5 changed files with 42 additions and 76 deletions

View File

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

View File

@@ -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")

View File

@@ -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)

View File

@@ -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

View File

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