From f31855e83cd8d225eecb739d7a1211b5864e9682 Mon Sep 17 00:00:00 2001
From: Sarah Wooders <sarahwooders@gmail.com>
Date: Fri, 16 Aug 2024 18:03:54 -0700
Subject: [PATCH] fix: fix tool creation to accept dev portal POST request 
 (#1656)

---
 memgpt/server/rest_api/tools/index.py | 40 +++++++++++++++++++++++++--
 tests/test_tools.py                   |  3 +-
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/memgpt/server/rest_api/tools/index.py b/memgpt/server/rest_api/tools/index.py
index 05525eab..6c676c39 100644
--- a/memgpt/server/rest_api/tools/index.py
+++ b/memgpt/server/rest_api/tools/index.py
@@ -18,7 +18,7 @@ class ListToolsResponse(BaseModel):
 
 
 class CreateToolRequest(BaseModel):
-    json_schema: dict = Field(..., description="JSON schema of the tool.")
+    json_schema: dict = Field(..., description="JSON schema of the tool.")  # NOT OpenAI - just has `name`
     source_code: str = Field(..., description="The source code of the function.")
     source_type: Optional[Literal["python"]] = Field(None, description="The type of the source code.")
     tags: Optional[List[str]] = Field(None, description="Metadata tags.")
@@ -81,9 +81,45 @@ def setup_user_tools_index_router(server: SyncServer, interface: QueuingInterfac
         """
         Create a new tool
         """
+        # NOTE: horrifying code, should be replaced when we migrate dev portal
+        from memgpt.agent import Agent  # nasty: need agent to be defined
+        from memgpt.functions.schema_generator import generate_schema
+
+        name = request.json_schema["name"]
+
+        import ast
+
+        parsed_code = ast.parse(request.source_code)
+        function_names = []
+
+        # Function to find and print function names
+        def find_function_names(node):
+            for child in ast.iter_child_nodes(node):
+                if isinstance(child, ast.FunctionDef):
+                    # Print the name of the function
+                    function_names.append(child.name)
+                # Recurse into child nodes
+                find_function_names(child)
+
+        # Find and print function names
+        find_function_names(parsed_code)
+        assert len(function_names) == 1, f"Expected 1 function, found {len(function_names)}: {function_names}"
+
+        # generate JSON schema
+        env = {}
+        env.update(globals())
+        exec(request.source_code, env)
+        func = env.get(function_names[0])
+        json_schema = generate_schema(func, name=name)
+        from pprint import pprint
+
+        pprint(json_schema)
+
         try:
+
             return server.create_tool(
-                json_schema=request.json_schema,
+                # json_schema=request.json_schema, # TODO: add back
+                json_schema=json_schema,
                 source_code=request.source_code,
                 source_type=request.source_type,
                 tags=request.tags,
diff --git a/tests/test_tools.py b/tests/test_tools.py
index b0ac1b62..34e2088b 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -7,7 +7,6 @@ import pytest
 from dotenv import load_dotenv
 
 from memgpt import Admin, create_client
-from memgpt.agent import Agent
 from memgpt.config import MemGPTConfig
 from memgpt.constants import DEFAULT_PRESET
 from memgpt.credentials import MemGPTCredentials
@@ -171,7 +170,7 @@ def test_create_agent_tool_admin(admin_client):
 def test_create_agent_tool(client):
     """Test creation of a agent tool"""
 
-    def core_memory_clear(self: Agent):
+    def core_memory_clear(self):
         """
         Args:
             agent (Agent): The agent to delete from memory.