From 1eeebed4d77fc52e5bab00fdbbb9b07f2c8337c2 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 14:59:43 -0700 Subject: [PATCH 1/7] Simplified toolkits as list of tools, enabled imports from root --- jupyter_ai_tools/__init__.py | 11 ++++++++++ jupyter_ai_tools/toolkits/code_execution.py | 10 +++------ jupyter_ai_tools/toolkits/file_system.py | 23 ++++++++------------- jupyter_ai_tools/toolkits/git.py | 23 +++++++++------------ jupyter_ai_tools/toolkits/notebook.py | 23 +++++++++------------ pyproject.toml | 2 +- 6 files changed, 44 insertions(+), 48 deletions(-) diff --git a/jupyter_ai_tools/__init__.py b/jupyter_ai_tools/__init__.py index 987cc59..4a7db5d 100644 --- a/jupyter_ai_tools/__init__.py +++ b/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"}] diff --git a/jupyter_ai_tools/toolkits/code_execution.py b/jupyter_ai_tools/toolkits/code_execution.py index 36023f6..cb7bad1 100644 --- a/jupyter_ai_tools/toolkits/code_execution.py +++ b/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, +] diff --git a/jupyter_ai_tools/toolkits/file_system.py b/jupyter_ai_tools/toolkits/file_system.py index a296d18..11a7723 100644 --- a/jupyter_ai_tools/toolkits/file_system.py +++ b/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, +] diff --git a/jupyter_ai_tools/toolkits/git.py b/jupyter_ai_tools/toolkits/git.py index 9f93384..4bf858b 100644 --- a/jupyter_ai_tools/toolkits/git.py +++ b/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, +] diff --git a/jupyter_ai_tools/toolkits/notebook.py b/jupyter_ai_tools/toolkits/notebook.py index fba849b..8080f2c 100644 --- a/jupyter_ai_tools/toolkits/notebook.py +++ b/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, +] diff --git a/pyproject.toml b/pyproject.toml index d4ba5b5..7d356e3 100644 --- a/pyproject.toml +++ b/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" ] From dd8d96ff543ae6c0ceb0beb86c441984ff86190c Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 15:05:34 -0700 Subject: [PATCH 2/7] Fixes typing errors --- jupyter_ai_tools/toolkits/notebook.py | 14 +++++++------- jupyter_ai_tools/utils.py | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/jupyter_ai_tools/toolkits/notebook.py b/jupyter_ai_tools/toolkits/notebook.py index 8080f2c..204c23b 100644 --- a/jupyter_ai_tools/toolkits/notebook.py +++ b/jupyter_ai_tools/toolkits/notebook.py @@ -211,8 +211,8 @@ async def get_cell_id_from_index(file_path: str, cell_index: int) -> str: async def add_cell( file_path: str, - content: str | None = None, - cell_id: str | None = None, + content: Optional[str] = None, + cell_id: Optional[str] = None, add_above: bool = False, cell_type: Literal["code", "markdown", "raw"] = "code", ): @@ -295,8 +295,8 @@ async def add_cell( async def insert_cell( file_path: str, - content: str | None = None, - insert_index: int | None = None, + content: Optional[str] = None, + insert_index: Optional[int] = None, cell_type: Literal["code", "markdown", "raw"] = "code", ): """Inserts a new cell to the Jupyter notebook at the specified cell index. @@ -887,7 +887,7 @@ def read_cell_nbformat(file_path: str, cell_id: str) -> Dict[str, Any]: raise ValueError(f"Cell with {cell_id=} not found in notebook at {file_path=}") -def _get_cell_index_from_id_json(notebook_json, cell_id: str) -> int | None: +def _get_cell_index_from_id_json(notebook_json, cell_id: str) -> Optional[int]: """Get cell index from cell_id by notebook json dict. Args: @@ -905,7 +905,7 @@ def _get_cell_index_from_id_json(notebook_json, cell_id: str) -> int | None: return None -def _get_cell_index_from_id_ydoc(ydoc, cell_id: str) -> int | None: +def _get_cell_index_from_id_ydoc(ydoc, cell_id: str) -> Optional[int]: """Get cell index from cell_id using YDoc interface. Args: @@ -924,7 +924,7 @@ def _get_cell_index_from_id_ydoc(ydoc, cell_id: str) -> int | None: return None -def _get_cell_index_from_id_nbformat(notebook, cell_id: str) -> int | None: +def _get_cell_index_from_id_nbformat(notebook, cell_id: str) -> Optional[int]: """Get cell index from cell_id using nbformat interface. Args: diff --git a/jupyter_ai_tools/utils.py b/jupyter_ai_tools/utils.py index 33f1ce7..1c81952 100644 --- a/jupyter_ai_tools/utils.py +++ b/jupyter_ai_tools/utils.py @@ -2,7 +2,7 @@ import inspect import os from pathlib import Path -from typing import Optional +from typing import Dict, List, Optional from urllib.parse import unquote from jupyter_server.auth.identity import User @@ -269,7 +269,7 @@ def notebook_json_to_md(notebook_json: dict, include_outputs: bool = True) -> st return "\n\n".join(md_parts) -def metadata_to_md(metadata_json: dict) -> str: +def metadata_to_md(metadata_json: Dict) -> str: """Converts notebook or cell metadata to markdown string in YAML format. Args: @@ -339,7 +339,7 @@ def cell_to_md(cell_json: dict, index: int = 0, include_outputs: bool = True) -> return "\n\n".join(md_parts) -def format_outputs(outputs: list) -> str: +def format_outputs(outputs: List) -> str: """Formats cell outputs into markdown. Args: From 677ae92e851ca039cf7b4df466f43210f61fb6e9 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 15:10:05 -0700 Subject: [PATCH 3/7] lint --- jupyter_ai_tools/__init__.py | 1 + jupyter_ai_tools/toolkits/file_system.py | 1 + 2 files changed, 2 insertions(+) diff --git a/jupyter_ai_tools/__init__.py b/jupyter_ai_tools/__init__.py index 4a7db5d..a5bbbbe 100644 --- a/jupyter_ai_tools/__init__.py +++ b/jupyter_ai_tools/__init__.py @@ -12,6 +12,7 @@ "nb_toolkit", ] + def _jupyter_server_extension_points(): return [{"module": "jupyter_ai_tools"}] diff --git a/jupyter_ai_tools/toolkits/file_system.py b/jupyter_ai_tools/toolkits/file_system.py index 11a7723..cdc1d69 100644 --- a/jupyter_ai_tools/toolkits/file_system.py +++ b/jupyter_ai_tools/toolkits/file_system.py @@ -353,6 +353,7 @@ 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 = [ read, edit, From f2bf3e98848cad1efee2cc61cfe8d9a1bb381095 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 16:06:24 -0700 Subject: [PATCH 4/7] Updated to align with tool updates --- README.md | 50 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/README.md b/README.md index 163799f..3442ed3 100644 --- a/README.md +++ b/README.md @@ -10,25 +10,41 @@ ______________________________________________________________________ This extension provides runtime-discoverable tools compatible with OpenAI-style function calling or MCP tool schemas. These tools can be invoked by agents to: -### 🧠 YNotebook Tools - -- `read_cell`: Return the full content of a cell by index -- `read_notebook`: Return all cells as a JSON-formatted list -- `add_cell`: Insert a blank cell at a specific index -- `delete_cell`: Remove a cell and return its contents -- `write_to_cell`: Overwrite the content of a cell with new source -- `get_max_cell_index`: Return the last valid cell index - -### 🌀 Git Tools - -- `git_clone`: Clone a Git repo into a given path -- `git_status`: Get the working tree status +### 📁 File System Tools (`fs_toolkit`) + +- `read`: Read file contents from the filesystem +- `edit`: Edit file contents with search and replace functionality +- `write`: Write content to a file +- `search_and_replace`: Search and replace text patterns in files +- `glob`: Find files matching a glob pattern +- `grep`: Search for text patterns within file contents +- `ls`: List directory contents + +### 🧠 Notebook Tools (`nb_toolkit`) + +- `read_notebook`: Read entire notebook contents as markdown +- `read_cell`: Read a specific notebook cell by index +- `add_cell`: Add a new cell to a notebook +- `insert_cell`: Insert a cell at a specific index in the notebook +- `delete_cell`: Remove a cell from the notebook +- `edit_cell`: Modify a cell's content +- `get_cell_id_from_index`: Get cell ID from its index position +- `create_notebook`: Create a new Jupyter notebook + +### 🌀 Git Tools (`git_toolkit`) + +- `git_clone`: Clone a Git repository to a specified path +- `git_status`: Get the current working tree status - `git_log`: View recent commit history -- `git_add`: Stage files (individually or all) +- `git_pull`: Pull changes from remote repository +- `git_push`: Push local changes to remote branch - `git_commit`: Commit staged changes with a message -- `git_push`: Push local changes to a remote branch -- `git_pull`: Pull remote updates -- `git_get_repo_root_from_notebookpath`: Find the Git root from a notebook path +- `git_add`: Stage files for commit (individually or all) +- `git_get_repo_root`: Get the root directory of the Git repository + +### ⚙️ Code Execution Tools (`exec_toolkit`) + +- `bash`: Execute bash commands in the system shell These tools are ideal for agents that assist users with code editing, version control, or dynamic notebook interaction. From c6863ac6f7970f505b7386efce1c239047400401 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 18:00:40 -0700 Subject: [PATCH 5/7] Added compatibility for collab, remove jsd as dependency --- jupyter_ai_tools/utils.py | 54 ++++++++++++++++++++++++++++----------- pyproject.toml | 3 +-- 2 files changed, 40 insertions(+), 17 deletions(-) diff --git a/jupyter_ai_tools/utils.py b/jupyter_ai_tools/utils.py index 1c81952..f17281b 100644 --- a/jupyter_ai_tools/utils.py +++ b/jupyter_ai_tools/utils.py @@ -2,13 +2,21 @@ import inspect import os from pathlib import Path -from typing import Dict, List, Optional +from typing import Any, Dict, List, Optional from urllib.parse import unquote from jupyter_server.auth.identity import User from jupyter_server.serverapp import ServerApp from pycrdt import Awareness +JSD_PRESENT = False +try: + import jupyter_server_documents + + JSD_PRESENT = True +except ImportError: + pass + def get_serverapp(): """Returns the server app from the request context""" @@ -66,6 +74,18 @@ def normalize_filepath(file_path: str) -> str: return str(resolved_path.resolve()) +async def get_room(room_id: str): + """Returns the yroom.""" + + serverapp = get_serverapp() + if JSD_PRESENT: + manager = serverapp.web_app.settings["yroom_manager"] + return manager.get_room(room_id) if manager.has_room(room_id) else None + else: + manager = serverapp.web_app.settings["jupyter_server_ydoc"].ywebsocket_server + return await manager.get_room(room_id) if manager.room_exists(room_id) else None + + async def get_jupyter_ydoc(file_id: str): """Returns the notebook ydoc @@ -75,27 +95,31 @@ async def get_jupyter_ydoc(file_id: str): Returns: `YNotebook` ydoc for the notebook """ - serverapp = get_serverapp() - yroom_manager = serverapp.web_app.settings["yroom_manager"] room_id = f"json:notebook:{file_id}" + yroom = await get_room(room_id) + if not yroom: + return - if yroom_manager.has_room(room_id): - yroom = yroom_manager.get_room(room_id) - notebook = await yroom.get_jupyter_ydoc() - return notebook + if JSD_PRESENT: + return await yroom.get_jupyter_ydoc() + else: + return yroom._document async def get_global_awareness() -> Optional[Awareness]: - serverapp = get_serverapp() - yroom_manager = serverapp.web_app.settings["yroom_manager"] + """Returns the global awareness object for JupyterLab collaboration. - room_id = "JupyterLab:globalAwareness" - if yroom_manager.has_room(room_id): - yroom = yroom_manager.get_room(room_id) - return yroom.get_awareness() + Returns: + Awareness instance for global collaboration, or None if not available + """ + yroom = await get_room("JupyterLab:globalAwareness") + if not yroom: + return - # Return None if room doesn't exist - return None + if JSD_PRESENT: + return yroom.get_awareness() + else: + return yroom.awareness async def get_file_id(file_path: str) -> str: diff --git a/pyproject.toml b/pyproject.toml index 7d356e3..39fe099 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,8 +22,7 @@ classifiers = [ dependencies = [ "jupyter_server>=1.6,<3", - "jupyterlab_git", - "jupyter_server_documents>=0.1.a8" + "jupyterlab_git" ] From e65a62e4756f3ae5acf5bdb26b54e9088cbd4ca8 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 18:10:34 -0700 Subject: [PATCH 6/7] Added missing deps --- jupyter_ai_tools/utils.py | 2 +- pyproject.toml | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/jupyter_ai_tools/utils.py b/jupyter_ai_tools/utils.py index f17281b..b5de290 100644 --- a/jupyter_ai_tools/utils.py +++ b/jupyter_ai_tools/utils.py @@ -2,7 +2,7 @@ import inspect import os from pathlib import Path -from typing import Any, Dict, List, Optional +from typing import Dict, List, Optional from urllib.parse import unquote from jupyter_server.auth.identity import User diff --git a/pyproject.toml b/pyproject.toml index 39fe099..e4a8e49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,7 +22,9 @@ classifiers = [ dependencies = [ "jupyter_server>=1.6,<3", - "jupyterlab_git" + "jupyterlab_git", + "pycrdt>=0.12.0,<0.13.0", + "jupyter_ydoc>=3.0.0,<4.0.0", ] From dffc69e0fac46148e3b84e700669c91759143e94 Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Wed, 29 Oct 2025 18:15:07 -0700 Subject: [PATCH 7/7] lint --- jupyter_ai_tools/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/jupyter_ai_tools/utils.py b/jupyter_ai_tools/utils.py index b5de290..2335a91 100644 --- a/jupyter_ai_tools/utils.py +++ b/jupyter_ai_tools/utils.py @@ -11,7 +11,7 @@ JSD_PRESENT = False try: - import jupyter_server_documents + import jupyter_server_documents # noqa: F401 JSD_PRESENT = True except ImportError: @@ -114,7 +114,7 @@ async def get_global_awareness() -> Optional[Awareness]: """ yroom = await get_room("JupyterLab:globalAwareness") if not yroom: - return + return None if JSD_PRESENT: return yroom.get_awareness()