Fixed register tool

Author: 3coins

jupyter_ai_tools/__init__.py

@@ -1,5 +1,5 @@
 from .extension import get_git_tools, get_notebook_tools
-from jupyter_server_ai_tools.models import Tool, Toolkit
+from jupyter_server_ai_tools.models import ToolSet, Toolkit
 
 __version__ = "0.1.2"
 
@@ -12,18 +12,18 @@ def _load_jupyter_server_extension(serverapp):
     serverapp.log.info("✅ jupyter_ai_tools extension loaded.")
 
 async def _start_jupyter_server_extension(serverapp):
-    registry = serverapp.extension_manager.extensions.get(
-        "jupyter_server_ai_tools"
-    )
+    registry = serverapp.web_app.settings["toolkit_registry"]
     if registry:
+        notebook_tools = ToolSet(get_notebook_tools())
         registry.register_toolkit(
             Toolkit(
-                name="notebook_toolkit", tools=get_notebook_tools()
+                name="notebook_toolkit", tools=notebook_tools
             )
         )
 
+        git_tools = ToolSet(get_git_tools())
         registry.register_toolkit(
             Toolkit(
-                name="git_toolkit", tools=get_git_tools()
+                name="git_toolkit", tools=git_tools
             )
         )