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.
 
diff --git a/jupyter_ai_tools/__init__.py b/jupyter_ai_tools/__init__.py
index 987cc59..a5bbbbe 100644
--- a/jupyter_ai_tools/__init__.py
+++ b/jupyter_ai_tools/__init__.py
@@ -1,5 +1,17 @@
+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..cdc1d69 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
 
 
@@ -356,14 +354,12 @@ async def ls(path: str, ignore: Optional[List[str]] = None) -> str:
         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..204c23b 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
 
@@ -212,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",
 ):
@@ -296,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.
@@ -888,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:
@@ -906,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:
@@ -925,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:
@@ -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/jupyter_ai_tools/utils.py b/jupyter_ai_tools/utils.py
index 33f1ce7..2335a91 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 Optional
+from typing import 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  # noqa: F401
+
+    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 None
 
-    # 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:
@@ -269,7 +293,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 +363,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:
diff --git a/pyproject.toml b/pyproject.toml
index d4ba5b5..e4a8e49 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,8 @@ classifiers = [
 dependencies = [
   "jupyter_server>=1.6,<3",
   "jupyterlab_git",
-  "jupyter_ai>=3.0.0-beta.1"
+  "pycrdt>=0.12.0,<0.13.0",
+  "jupyter_ydoc>=3.0.0,<4.0.0",
 ]