WIP: Toolkits for Jupyter

Author: 3coins

jupyter_ai_tools/__init__.py

@@ -1,4 +1,4 @@
-from .extension import get_git_tools, get_notebook_tools
+from .tools import get_git_tools, get_notebook_tools
 from jupyter_server_ai_tools.models import ToolSet, Toolkit
 
 __version__ = "0.1.2"
@@ -7,23 +7,5 @@
 def _jupyter_server_extension_points():
     return [{"module": "jupyter_ai_tools"}]
 
-
 def _load_jupyter_server_extension(serverapp):
-    serverapp.log.info("✅ jupyter_ai_tools extension loaded.")
-
-async def _start_jupyter_server_extension(serverapp):
-    registry = serverapp.web_app.settings["toolkit_registry"]
-    if registry:
-        notebook_tools = ToolSet(get_notebook_tools())
-        registry.register_toolkit(
-            Toolkit(
-                name="notebook_toolkit", tools=notebook_tools
-            )
-        )
-
-        git_tools = ToolSet(get_git_tools())
-        registry.register_toolkit(
-            Toolkit(
-                name="git_toolkit", tools=git_tools
-            )
-        )
+    serverapp.log.info("✅ jupyter_ai_tools extension loaded.")

jupyter_ai_tools/extension.py

@@ -0,0 +1 @@
+"""Toolkits for Jupyter"""

jupyter_ai_tools/toolkits/__init__.py

@@ -0,0 +1,49 @@
+"""Tools that provide code execution features"""
+
+
+import asyncio
+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
+    
+    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 = Toolkit(
+    name="code_execution_toolkit",
+    description="Tools to execute code in different environments.",
+)
+toolkit.add(Tool(callable=bash, execute=True))
+