Add workspace directory context to file tools

Added cwd parameter to read, edit, and write tools to support relative
file paths. The tools now resolve relative paths against the workspace
directory, allowing the LLM to reference files without absolute paths.

Changes:
- Added cwd parameter to read(), edit(), and write() in default_toolkit.py
- Updated _build_toolkit() to use functools.partial with a helper function
- Replaced verbose wrapper functions with cleaner partial binding approach
- All file tools now receive workspace directory context automatically

Author: joadoumie

packages/jupyter-ai/jupyter_ai/personas/jupyternaut/jupyternaut.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Optional
 from jupyterlab_chat.models import Message
 
@@ -71,40 +72,19 @@ def _build_toolkit(self) -> Toolkit:
         # Get workspace directory for this chat
         workspace_dir = self.get_workspace_dir()
 
-        # Create wrapper functions that bind workspace_dir
-        # We can't use functools.partial because litellm.function_to_dict expects __name__
-        async def bash(command: str, timeout: Optional[int] = None) -> str:
-            """Executes a bash command and returns the result
-
-            Args:
-                command: The bash command to execute
-                timeout: Optional timeout in seconds
-
-            Returns:
-                The command output (stdout and stderr combined)
-            """
-            from ...tools.default_toolkit import bash as bash_orig
-            return await bash_orig(command, timeout=timeout, cwd=workspace_dir)
-
-        async def search_grep(pattern: str, include: str = "*") -> str:
-            """Search for text patterns in files using ripgrep.
-
-            Args:
-                pattern: A regular expression pattern to search for
-                include: A glob pattern to filter which files to search
-
-            Returns:
-                The raw output from ripgrep, including file paths, line numbers, and matching lines
-            """
-            from ...tools.default_toolkit import search_grep as search_grep_orig
-            return await search_grep_orig(pattern, include=include, cwd=workspace_dir)
+        def bind_cwd(func, **kwargs):
+            """Create a partial function with custom __name__ and __doc__ preserved"""
+            bound_func = partial(func, **kwargs)
+            bound_func.__name__ = func.__name__
+            bound_func.__doc__ = func.__doc__
+            return bound_func
 
         # Create toolkit with workspace-aware tools
         toolkit = Toolkit(name="jupyter-ai-contextual-toolkit")
-        toolkit.add_tool(Tool(callable=bash))
-        toolkit.add_tool(Tool(callable=search_grep))
-        toolkit.add_tool(Tool(callable=read))
-        toolkit.add_tool(Tool(callable=edit))
-        toolkit.add_tool(Tool(callable=write))
+        toolkit.add_tool(Tool(callable=bind_cwd(bash, cwd=workspace_dir)))
+        toolkit.add_tool(Tool(callable=bind_cwd(search_grep, cwd=workspace_dir)))
+        toolkit.add_tool(Tool(callable=bind_cwd(read, cwd=workspace_dir)))
+        toolkit.add_tool(Tool(callable=bind_cwd(edit, cwd=workspace_dir)))
+        toolkit.add_tool(Tool(callable=bind_cwd(write, cwd=workspace_dir)))
 
         return toolkit

packages/jupyter-ai/jupyter_ai/tools/default_toolkit.py

@@ -6,20 +6,23 @@
 from .models import Tool, Toolkit
 
 
-def read(file_path: str, offset: int, limit: int) -> str:
+def read(file_path: str, offset: int, limit: int, cwd: Optional[str] = None) -> str:
     """
     Read a subset of lines from a text file.
 
     Parameters
     ----------
     file_path : str
-        Absolute path to the file that should be read.
+        Path to the file that should be read. Can be absolute or relative to cwd.
     offset : int
         The line number at which to start reading (1-based indexing).
     limit : int
-        Number of lines to read starting from *offset*.  
+        Number of lines to read starting from *offset.
         If *offset + limit* exceeds the number of lines in the file,
         all available lines after *offset* are returned.
+    cwd : str, optional
+        The directory to use as the base for relative paths. If not provided,
+        file_path must be absolute.
 
     Returns
     -------
@@ -33,6 +36,8 @@ def read(file_path: str, offset: int, limit: int) -> str:
     ['third line\n', 'fourth line\n', 'fifth line\n', 'sixth line\n']
     """
     path = pathlib.Path(file_path)
+    if cwd and not path.is_absolute():
+        path = pathlib.Path(cwd) / path
     if not path.is_file():
         raise FileNotFoundError(f"File not found: {file_path}")
 
@@ -69,21 +74,25 @@ def edit(
     old_string: str,
     new_string: str,
     replace_all: bool = False,
+    cwd: Optional[str] = None,
 ) -> None:
     """
     Replace occurrences of a substring in a file.
 
     Parameters
     ----------
     file_path : str
-        Absolute path to the file that should be edited.
+        Path to the file that should be edited. Can be absolute or relative to cwd.
     old_string : str
         Text that should be replaced.
     new_string : str
         Text that will replace *old_string*.
     replace_all : bool, optional
         If ``True`` all occurrences of *old_string* are replaced.
         If ``False`` (default), only the first occurrence in the file is replaced.
+    cwd : str, optional
+        The directory to use as the base for relative paths. If not provided,
+        file_path must be absolute.
 
     Returns
     -------
@@ -110,6 +119,8 @@ def edit(
     >>> edit('/tmp/test.txt', 'foo', 'bar', replace_all=True)
     """
     path = pathlib.Path(file_path)
+    if cwd and not path.is_absolute():
+        path = pathlib.Path(cwd) / path
     if not path.is_file():
         raise FileNotFoundError(f"File not found: {file_path}")
 
@@ -129,16 +140,19 @@ def edit(
     path.write_text(new_content, encoding="utf-8")
 
 
-def write(file_path: str, content: str) -> None:
+def write(file_path: str, content: str, cwd: Optional[str] = None) -> None:
     """
     Write content to a file, creating it if it doesn't exist.
 
     Parameters
     ----------
     file_path : str
-        Absolute path to the file that should be written.
+        Path to the file that should be written. Can be absolute or relative to cwd.
     content : str
         Content to write to the file.
+    cwd : str, optional
+        The directory to use as the base for relative paths. If not provided,
+        file_path must be absolute.
 
     Returns
     -------
@@ -160,7 +174,9 @@ def write(file_path: str, content: str) -> None:
     >>> write('/tmp/data.json', '{"key": "value"}')
     """
     path = pathlib.Path(file_path)
-    
+    if cwd and not path.is_absolute():
+        path = pathlib.Path(cwd) / path
+
     # Write the content to the file
     path.write_text(content, encoding="utf-8")