feat: get rid of client injection in args and just make it always available (#6459)
This commit is contained in:
committed by
Caren Thomas
parent
f417e53638
commit
86023db9b1
@@ -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"]
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user