Added tools

Author: 3coins

jupyter_ai_jupyternaut/jupyternaut/jupyternaut.py

@@ -15,6 +15,11 @@
     JupyternautSystemPromptArgs,
 )
 
+from .toolkits.code_execution import toolkit as exec_toolkit
+from .toolkits.filesystem import toolkit as fs_toolkit
+from .toolkits.notebook import toolkit as nb_toolkit
+
+
 MEMORY_STORE_PATH = os.path.join(jupyter_data_dir(), "jupyter_ai", "memory.sqlite")
 
 
@@ -40,11 +45,23 @@ async def get_memory_store(self):
             conn = await aiosqlite.connect(MEMORY_STORE_PATH, check_same_thread=False)
             self._memory_store = AsyncSqliteSaver(conn)        
         return self._memory_store
+    
+    def get_tools(self):
+        tools = []
+        tools += nb_toolkit
+        tools += fs_toolkit
 
-    async def _create_agent(self, model_id: str, model_args, system_prompt: str):
+        return tools
+
+    async def get_agent(self, model_id: str, model_args, system_prompt: str):
         model = ChatLiteLLM(**model_args, model_id=model_id, streaming=True)
         memory_store = await self.get_memory_store()
-        return create_agent(model, system_prompt=system_prompt, checkpointer=memory_store)
+        return create_agent(
+            model, 
+            system_prompt=system_prompt, 
+            checkpointer=memory_store,
+            tools=self.get_tools()
+        )
 
     async def process_message(self, message: Message) -> None:
         if not hasattr(self, "config_manager"):
@@ -62,7 +79,7 @@ async def process_message(self, message: Message) -> None:
         model_id = self.config_manager.chat_model
         model_args = self.config_manager.chat_model_args
         system_prompt = self.get_system_prompt(model_id=model_id, message=message)
-        agent = await self._create_agent(
+        agent = await self.get_agent(
             model_id=model_id, 
             model_args=model_args, 
             system_prompt=system_prompt

jupyter_ai_jupyternaut/jupyternaut/toolkits/__init__.py

@@ -0,0 +1,41 @@
+"""Tools that provide code execution features"""
+
+import asyncio
+import shlex
+from typing import Optional
+
+
+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)
+    """
+
+    proc = await asyncio.create_subprocess_exec(
+        *shlex.split(command),
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.PIPE,
+    )
+
+    try:
+        stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout)
+        output = stdout.decode("utf-8")
+        error = stderr.decode("utf-8")
+
+        if proc.returncode != 0:
+            if error:
+                return f"Error: {error}"
+            return f"Command failed with exit code {proc.returncode}"
+
+        return output if output else "Command executed successfully with no output."
+    except asyncio.TimeoutError:
+        proc.kill()
+        return f"Command timed out after {timeout} seconds"
+
+
+toolkit = [bash]