Fimeg/letta-server
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

feat: Make the tool runner take a schema (#1328)

Browse Source
This commit is contained in:
Matthew Zhou
2025-03-18 12:06:02 -07:00
committed by GitHub
parent e4456acec1
commit c22d1f4af2
4 changed files with 66 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
letta/schemas/tool.py
View File

@@ -250,3 +250,6 @@ class ToolRunFromSource(LettaBase):
name: Optional[str] = Field(None, description="The name of the tool to run.")
source_type: Optional[str] = Field(None, description="The type of the source code.")
args_json_schema: Optional[Dict] = Field(None, description="The args JSON schema of the function.")
json_schema: Optional[Dict] = Field(
None, description="The JSON schema of the function (auto-generated from source_code if not provided)"
)

1
letta/server/rest_api/routers/v1/tools.py
View File

@@ -192,6 +192,7 @@ def run_tool_from_source(
tool_env_vars=request.env_vars,
tool_name=request.name,
tool_args_json_schema=request.args_json_schema,
tool_json_schema=request.json_schema,
actor=actor,
)
except LettaToolCreateError as e:

6
letta/server/server.py
View File

@@ -1202,6 +1202,7 @@ class SyncServer(Server):
tool_source_type: Optional[str] = None,
tool_name: Optional[str] = None,
tool_args_json_schema: Optional[Dict[str, Any]] = None,
tool_json_schema: Optional[Dict[str, Any]] = None,
) -> ToolReturnMessage:
"""Run a tool from source code"""
if tool_source_type is not None and tool_source_type != "python":
@@ -1213,6 +1214,11 @@ class SyncServer(Server):
source_code=tool_source,
args_json_schema=tool_args_json_schema,
)
# If tools_json_schema is explicitly passed in, override it on the created Tool object
if tool_json_schema:
tool.json_schema = tool_json_schema
assert tool.name is not None, "Failed to create tool object"
# TODO eventually allow using agent state in tools

89
tests/test_server.py
View File

@@ -798,22 +798,25 @@ def ingest(message: str):
'''
def test_tool_run(server, mock_e2b_api_key_none, user, agent_id):
"""Test that the server can run tools"""
import pytest
def test_tool_run_basic(server, mock_e2b_api_key_none, user):
"""Test running a simple tool from source"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args={"message": "Hello, world!"},
# tool_name="ingest",
)
print(result)
assert result.status == "success"
assert result.tool_return == "Ingested message Hello, world!", result.tool_return
assert result.tool_return == "Ingested message Hello, world!"
assert not result.stdout
assert not result.stderr
def test_tool_run_with_env_var(server, mock_e2b_api_key_none, user):
"""Test running a tool that uses an environment variable"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_ENV_VAR,
@@ -821,56 +824,45 @@ def test_tool_run(server, mock_e2b_api_key_none, user, agent_id):
tool_args={},
tool_env_vars={"secret": "banana"},
)
print(result)
assert result.status == "success"
assert result.tool_return == "banana", result.tool_return
assert result.tool_return == "banana"
assert not result.stdout
assert not result.stderr
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args={"message": "Well well well"},
# tool_name="ingest",
)
print(result)
assert result.status == "success"
assert result.tool_return == "Ingested message Well well well", result.tool_return
assert not result.stdout
assert not result.stderr
def test_tool_run_invalid_args(server, mock_e2b_api_key_none, user):
"""Test running a tool with incorrect arguments"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args={"bad_arg": "oh no"},
# tool_name="ingest",
)
print(result)
assert result.status == "error"
assert "Error" in result.tool_return, result.tool_return
assert "missing 1 required positional argument" in result.tool_return, result.tool_return
assert "Error" in result.tool_return
assert "missing 1 required positional argument" in result.tool_return
assert not result.stdout
assert result.stderr
assert "missing 1 required positional argument" in result.stderr[0]
# Test that we can still pull the tool out by default (pulls that last tool in the source)
def test_tool_run_with_distractor(server, mock_e2b_api_key_none, user):
"""Test running a tool with a distractor function in the source"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
tool_source_type="python",
tool_args={"message": "Well well well"},
# tool_name="ingest",
)
print(result)
assert result.status == "success"
assert result.tool_return == "Ingested message Well well well", result.tool_return
assert result.tool_return == "Ingested message Well well well"
assert result.stdout
assert "I'm a distractor" in result.stdout[0]
assert not result.stderr
# Test that we can pull the tool out by name
def test_tool_run_explicit_tool_name(server, mock_e2b_api_key_none, user):
"""Test selecting a tool by name when multiple tools exist in the source"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
@@ -878,14 +870,15 @@ def test_tool_run(server, mock_e2b_api_key_none, user, agent_id):
tool_args={"message": "Well well well"},
tool_name="ingest",
)
print(result)
assert result.status == "success"
assert result.tool_return == "Ingested message Well well well", result.tool_return
assert result.tool_return == "Ingested message Well well well"
assert result.stdout
assert "I'm a distractor" in result.stdout[0]
assert not result.stderr
# Test that we can pull a different tool out by name
def test_tool_run_util_function(server, mock_e2b_api_key_none, user):
"""Test selecting a utility function that does not return anything meaningful"""
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE_WITH_DISTRACTOR,
@@ -893,14 +886,44 @@ def test_tool_run(server, mock_e2b_api_key_none, user, agent_id):
tool_args={},
tool_name="util_do_nothing",
)
print(result)
assert result.status == "success"
assert result.tool_return == str(None), result.tool_return
assert result.tool_return == str(None)
assert result.stdout
assert "I'm a distractor" in result.stdout[0]
assert not result.stderr
def test_tool_run_with_explicit_json_schema(server, mock_e2b_api_key_none, user):
"""Test overriding the autogenerated JSON schema with an explicit one"""
explicit_json_schema = {
"name": "ingest",
"description": "Blah blah blah.",
"parameters": {
"type": "object",
"properties": {
"message": {"type": "string", "description": "The message to ingest into the system."},
"request_heartbeat": {
"type": "boolean",
"description": "Request an immediate heartbeat after function execution. Set to `True` if you want to send a follow-up message or run a follow-up function.",
},
},
"required": ["message", "request_heartbeat"],
},
}
result = server.run_tool_from_source(
actor=user,
tool_source=EXAMPLE_TOOL_SOURCE,
tool_source_type="python",
tool_args={"message": "Custom schema test"},
tool_json_schema=explicit_json_schema,
)
assert result.status == "success"
assert result.tool_return == "Ingested message Custom schema test"
assert not result.stdout
assert not result.stderr
def test_composio_client_simple(server):
apps = server.get_composio_apps()
# Assert there's some amount of apps returned