From c22d1f4af2bfb5ff073a73a50302665a668a0688 Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Tue, 18 Mar 2025 12:06:02 -0700 Subject: [PATCH] feat: Make the tool runner take a schema (#1328) --- letta/schemas/tool.py | 3 + letta/server/rest_api/routers/v1/tools.py | 1 + letta/server/server.py | 6 ++ tests/test_server.py | 89 ++++++++++++++--------- 4 files changed, 66 insertions(+), 33 deletions(-) diff --git a/letta/schemas/tool.py b/letta/schemas/tool.py index b23d9ac2..a0f7e4d3 100644 --- a/letta/schemas/tool.py +++ b/letta/schemas/tool.py @@ -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)" + ) diff --git a/letta/server/rest_api/routers/v1/tools.py b/letta/server/rest_api/routers/v1/tools.py index b302a569..f197e449 100644 --- a/letta/server/rest_api/routers/v1/tools.py +++ b/letta/server/rest_api/routers/v1/tools.py @@ -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: diff --git a/letta/server/server.py b/letta/server/server.py index fd417f19..ddde20cc 100644 --- a/letta/server/server.py +++ b/letta/server/server.py @@ -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 diff --git a/tests/test_server.py b/tests/test_server.py index bc9562e3..c6db4da6 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -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