Simplified toolkits as list of tools, enabled imports from root

Author: 3coins

jupyter_ai_tools/__init__.py

@@ -1,5 +1,16 @@
+from .toolkits.code_execution import toolkit as exec_toolkit
+from .toolkits.file_system import toolkit as fs_toolkit
+from .toolkits.git import toolkit as git_toolkit
+from .toolkits.notebook import toolkit as nb_toolkit
+
 __version__ = "0.2.1"
 
+__all__ = [
+    "fs_toolkit",
+    "exec_toolkit",
+    "git_toolkit",
+    "nb_toolkit",
+]
 
 def _jupyter_server_extension_points():
     return [{"module": "jupyter_ai_tools"}]

jupyter_ai_tools/toolkits/code_execution.py

@@ -4,8 +4,6 @@
 import shlex
 from typing import Optional
 
-from jupyter_ai.tools.models import Tool, Toolkit
-
 
 async def bash(command: str, timeout: Optional[int] = None) -> str:
     """Executes a bash command and returns the result
@@ -40,8 +38,6 @@ async def bash(command: str, timeout: Optional[int] = None) -> str:
         return f"Command timed out after {timeout} seconds"
 
 
-toolkit = Toolkit(
-    name="code_execution_toolkit",
-    description="Tools to execute code in different environments.",
-)
-toolkit.add_tool(Tool(callable=bash, execute=True))
+toolkit = [
+    bash,
+]

jupyter_ai_tools/toolkits/file_system.py

@@ -6,8 +6,6 @@
 import os
 from typing import List, Optional
 
-from jupyter_ai.tools.models import Tool, Toolkit
-
 from ..utils import normalize_filepath
 
 
@@ -355,15 +353,12 @@ async def ls(path: str, ignore: Optional[List[str]] = None) -> str:
     except Exception as e:
         return f"Error: Failed to list directory: {str(e)}"
 
-
-toolkit = Toolkit(
-    name="file_system_toolkit",
-    description="Tools to do search, list, read, write and edit operations on files.",
-)
-toolkit.add_tool(Tool(callable=read, read=True))
-toolkit.add_tool(Tool(callable=edit, read=True, write=True))
-toolkit.add_tool(Tool(callable=write, write=True))
-toolkit.add_tool(Tool(callable=search_and_replace, read=True, write=True))
-toolkit.add_tool(Tool(callable=glob, read=True))
-toolkit.add_tool(Tool(callable=grep, read=True))
-toolkit.add_tool(Tool(callable=ls, read=True))
+toolkit = [
+    read,
+    edit,
+    write,
+    search_and_replace,
+    glob,
+    grep,
+    ls,
+]

jupyter_ai_tools/toolkits/git.py

@@ -1,7 +1,6 @@
 import json
 import os
 
-from jupyter_ai.tools.models import Tool, Toolkit
 from jupyterlab_git.git import Git
 
 from ..utils import normalize_filepath
@@ -175,15 +174,13 @@ async def git_get_repo_root(path: str) -> str:
     return f"❌ Not inside a Git repo. {res.get('message', '')}"
 
 
-toolkit = Toolkit(
-    name="git_toolkit",
-    description="Tools for working with Git repositories.",
-)
-toolkit.add_tool(Tool(callable=git_clone, execute=True))
-toolkit.add_tool(Tool(callable=git_status, read=True))
-toolkit.add_tool(Tool(callable=git_log, read=True))
-toolkit.add_tool(Tool(callable=git_pull, execute=True))
-toolkit.add_tool(Tool(callable=git_push, execute=True))
-toolkit.add_tool(Tool(callable=git_commit, execute=True))
-toolkit.add_tool(Tool(callable=git_add, execute=True))
-toolkit.add_tool(Tool(callable=git_get_repo_root, read=True))
+toolkit = [
+    git_clone,
+    git_status,
+    git_log,
+    git_pull,
+    git_push,
+    git_commit,
+    git_add,
+    git_get_repo_root,
+]

jupyter_ai_tools/toolkits/notebook.py

@@ -6,7 +6,6 @@
 from typing import Any, Dict, Literal, Optional, Tuple
 
 import nbformat
-from jupyter_ai.tools.models import Tool, Toolkit
 from jupyter_ydoc import YNotebook
 from pycrdt import Assoc, Text
 
@@ -1006,15 +1005,13 @@ async def create_notebook(file_path: str) -> str:
         return f"Error: Failed to create notebook: {str(e)}"
 
 
-toolkit = Toolkit(
-    name="notebook_toolkit",
-    description="Tools for reading and manipulating Jupyter notebooks.",
-)
-toolkit.add_tool(Tool(callable=read_notebook, read=True))
-toolkit.add_tool(Tool(callable=read_cell, read=True))
-toolkit.add_tool(Tool(callable=add_cell, read=True, write=True))
-toolkit.add_tool(Tool(callable=insert_cell, read=True, write=True))
-toolkit.add_tool(Tool(callable=delete_cell, delete=True))
-toolkit.add_tool(Tool(callable=edit_cell, read=True, write=True))
-toolkit.add_tool(Tool(callable=get_cell_id_from_index, read=True))
-toolkit.add_tool(Tool(callable=create_notebook, write=True))
+toolkit = [
+    read_notebook,
+    read_cell,
+    add_cell,
+    insert_cell,
+    delete_cell,
+    edit_cell,
+    get_cell_id_from_index,
+    create_notebook,
+]

pyproject.toml

@@ -23,7 +23,7 @@ classifiers = [
 dependencies = [
   "jupyter_server>=1.6,<3",
   "jupyterlab_git",
-  "jupyter_ai>=3.0.0-beta.1"
+  "jupyter_server_documents>=0.1.a8"
 ]