Improve tools

jupyter_ai_jupyternaut/jupyternaut/jupyternaut.py

@@ -22,6 +22,7 @@
 )
 from .toolkits.notebook import toolkit as nb_toolkit
 from .toolkits.jupyterlab import toolkit as jlab_toolkit
+from .toolkits.code_execution import toolkit as exec_toolkit
 
 MEMORY_STORE_PATH = os.path.join(jupyter_data_dir(), "jupyter_ai", "memory.sqlite")
 
@@ -123,29 +124,18 @@ async def get_memory_store(self):
     def get_tools(self):
         tools = nb_toolkit
         tools += jlab_toolkit
-        return nb_toolkit
+        tools += exec_toolkit
+        return tools
 
     async def get_agent(self, model_id: str, model_args, system_prompt: str):
         model = ChatLiteLLM(**model_args, model=model_id, streaming=True)
         memory_store = await self.get_memory_store()
-
-        if not hasattr(self, "search_tool"):
-            self.search_tool = FilesystemFileSearchMiddleware(
-                root_path=self.parent.root_dir
-            )
-        if not hasattr(self, "shell_tool"):
-            self.shell_tool = ShellToolMiddleware(workspace_root=self.parent.root_dir)
-        if not hasattr(self, "tool_call_handler"):
-            self.tool_call_handler = ToolMonitoringMiddleware(
-                persona=self
-            )
-
+
         return create_agent(
             model,
             system_prompt=system_prompt,
             checkpointer=memory_store,
-            tools=self.get_tools(), # notebook and jlab tools
-            middleware=[self.shell_tool, self.tool_call_handler],
+            tools=self.get_tools(),
         )
 
     async def process_message(self, message: Message) -> None:

jupyter_ai_jupyternaut/jupyternaut/toolkits/jupyterlab.py

@@ -5,12 +5,21 @@ async def open_file(file_path: str):
     """
     Opens a file in JupyterLab main area
     """
-    await execute_command("docmanager:open", {"path": file_path})
+    return await execute_command("docmanager:open", {"path": file_path})
+
 
 async def run_all_cells():
     """
     Runs all cells in the currently active Jupyter notebook
     """
     return await execute_command("notebook:run-all-cells")
 
-toolkit = [open_file, run_all_cells]
+
+async def restart_kernel():
+    """
+    Restarts the notebook kernel, useful when new packages are installed
+    """
+    return await execute_command("notebook:restart-kernel")
+
+
+toolkit = [open_file, run_all_cells, restart_kernel]