From 622de59fd042e20ed9993fecaef41d4e60f40ed7 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Mon, 30 Sep 2024 09:24:10 -0700
Subject: [PATCH] fix: use JSON schema name instead of tool name for loading
 from env (#1798)

---
 letta/agent.py         | 5 ++++-
 letta/server/server.py | 2 +-
 tests/test_tools.py    | 6 +++++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/letta/agent.py b/letta/agent.py
index 7b460388..3fbc7016 100644
--- a/letta/agent.py
+++ b/letta/agent.py
@@ -350,7 +350,10 @@ class Agent(BaseAgent):
                 exec(tool.module, env)
             else:
                 exec(tool.source_code, env)
-            self.functions_python[tool.name] = env[tool.name]
+            from pprint import pprint
+
+            pprint(tool.json_schema)
+            self.functions_python[tool.name] = env[tool.json_schema["name"]]
             self.functions.append(tool.json_schema)
         assert all([callable(f) for k, f in self.functions_python.items()]), self.functions_python
 
diff --git a/letta/server/server.py b/letta/server/server.py
index 2ab79144..6fdcff99 100644
--- a/letta/server/server.py
+++ b/letta/server/server.py
@@ -1752,7 +1752,7 @@ class SyncServer(Server):
 
             # TODO: not sure if this always works
             func = env[functions[-1]]
-            json_schema = generate_schema(func, request.name)
+            json_schema = generate_schema(func)
         else:
             # provided by client
             json_schema = request.json_schema
diff --git a/tests/test_tools.py b/tests/test_tools.py
index f7750b73..3fa6eec7 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -106,7 +106,7 @@ def test_create_tool(client: Union[LocalClient, RESTClient]):
     tools = client.list_tools()
     print(f"Original tools {[t.name for t in tools]}")
 
-    tool = client.create_tool(print_tool, tags=["extras"])
+    tool = client.create_tool(print_tool, name="my_name", tags=["extras"])
 
     tools = client.list_tools()
     assert tool in tools, f"Expected {tool.name} in {[t.name for t in tools]}"
@@ -117,6 +117,10 @@ def test_create_tool(client: Union[LocalClient, RESTClient]):
     assert tool is not None, "Expected tool to be created"
     assert tool.id == tool.id, f"Expected {tool.id} to be {tool.id}"
 
+    # create agent with tool
+    agent_state = client.create_agent(tools=[tool.name])
+    response = client.user_message(agent_id=agent_state.id, message="hi")
+
 
 # TODO: add back once we fix admin client tool creation
 # def test_create_agent_tool_admin(admin_client):