hdsp-jupyter-extension 2.0.6__py3-none-any.whl → 2.0.7__py3-none-any.whl

agent_server/langchain/tools/search_tools.py

@@ -1,131 +1,142 @@
 """
 Search Tools for LangChain Agent
 
-Provides tools for searching code in workspace and notebooks:
-- search_workspace: Search files in the workspace
-- search_notebook_cells: Search cells in Jupyter notebooks
+Provides tools for searching code in workspace and notebooks.
+These tools return pending_execution status and are executed on the client (Jupyter) side
+using subprocess (find/grep/ripgrep).
+
+Key features:
+- Returns command info for client-side execution via subprocess
+- Supports ripgrep (rg) if available, falls back to grep
+- Executes immediately without user approval
+- Shows the command being executed in status messages
 """
 
-import json
-import os
-import re
+import logging
+import shutil
 from typing import Any, Dict, List, Optional
 
 from langchain_core.tools import tool
 from pydantic import BaseModel, Field
 
+logger = logging.getLogger(__name__)
+
 
 class SearchWorkspaceInput(BaseModel):
     """Input schema for search_workspace tool"""
+
     pattern: str = Field(description="Search pattern (regex or text)")
     file_types: List[str] = Field(
         default=["*.py", "*.ipynb"],
-        description="File patterns to search (e.g., ['*.py', '*.ipynb'])"
+        description="File patterns to search (e.g., ['*.py', '*.ipynb'])",
     )
     path: str = Field(default=".", description="Directory to search in")
     max_results: int = Field(default=50, description="Maximum number of results")
     case_sensitive: bool = Field(default=False, description="Case-sensitive search")
+    execution_result: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Execution result payload from the client",
+    )
 
 
 class SearchNotebookCellsInput(BaseModel):
     """Input schema for search_notebook_cells tool"""
+
     pattern: str = Field(description="Search pattern (regex or text)")
     notebook_path: Optional[str] = Field(
-        default=None,
-        description="Specific notebook to search (None = all notebooks)"
+        default=None, description="Specific notebook to search (None = all notebooks)"
     )
     cell_type: Optional[str] = Field(
         default=None,
-        description="Cell type filter: 'code', 'markdown', or None for all"
+        description="Cell type filter: 'code', 'markdown', or None for all",
     )
     max_results: int = Field(default=30, description="Maximum number of results")
     case_sensitive: bool = Field(default=False, description="Case-sensitive search")
+    execution_result: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="Execution result payload from the client",
+    )
+
+
+def _is_ripgrep_available() -> bool:
+    """Check if ripgrep (rg) is installed and available."""
+    return shutil.which("rg") is not None
 
 
-def _search_in_file(
-    file_path: str,
+def _build_grep_command(
     pattern: str,
-    case_sensitive: bool = False
-) -> List[Dict[str, Any]]:
-    """Search for pattern in a regular file"""
-    results = []
-    flags = 0 if case_sensitive else re.IGNORECASE
-
-    try:
-        compiled = re.compile(pattern, flags)
-    except re.error:
-        # If pattern is not valid regex, use literal search
-        compiled = re.compile(re.escape(pattern), flags)
-
-    try:
-        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
-            for line_num, line in enumerate(f, 1):
-                if compiled.search(line):
-                    results.append({
-                        "file_path": file_path,
-                        "line_number": line_num,
-                        "content": line.strip()[:200],
-                        "match_type": "line",
-                    })
-    except Exception:
-        pass
-
-    return results
-
-
-def _search_in_notebook(
-    notebook_path: str,
+    file_types: List[str],
+    path: str,
+    case_sensitive: bool,
+    max_results: int,
+) -> tuple[str, str]:
+    """
+    Build a grep/ripgrep command for searching files.
+
+    Returns:
+        Tuple of (command_string, tool_name) where tool_name is 'rg' or 'grep'
+    """
+    # Check ripgrep availability (this check will also be done on client)
+    use_ripgrep = _is_ripgrep_available()
+
+    if use_ripgrep:
+        # Build ripgrep command
+        cmd_parts = ["rg", "--line-number", "--with-filename"]
+
+        if not case_sensitive:
+            cmd_parts.append("--ignore-case")
+
+        # Add file type filters using glob patterns
+        for ft in file_types:
+            cmd_parts.extend(["--glob", ft])
+
+        # Limit results
+        cmd_parts.extend(["--max-count", str(max_results)])
+
+        # Escape pattern for shell
+        escaped_pattern = pattern.replace("'", "'\\''")
+        cmd_parts.append(f"'{escaped_pattern}'")
+        cmd_parts.append(path)
+
+        return " ".join(cmd_parts), "rg"
+    else:
+        # Build find + grep command for cross-platform compatibility
+        find_parts = ["find", path, "-type", "f", "("]
+
+        for i, ft in enumerate(file_types):
+            if i > 0:
+                find_parts.append("-o")
+            find_parts.extend(["-name", f"'{ft}'"])
+
+        find_parts.append(")")
+
+        # Add grep with proper flags
+        grep_flags = "-n"  # Line numbers
+        if not case_sensitive:
+            grep_flags += "i"
+
+        # Escape pattern for shell
+        escaped_pattern = pattern.replace("'", "'\\''")
+
+        # Combine with xargs for efficiency
+        cmd = f"{' '.join(find_parts)} 2>/dev/null | xargs grep -{grep_flags} '{escaped_pattern}' 2>/dev/null | head -n {max_results}"
+
+        return cmd, "grep"
+
+
+def _build_notebook_search_command(
     pattern: str,
-    cell_type: Optional[str] = None,
-    case_sensitive: bool = False
-) -> List[Dict[str, Any]]:
-    """Search for pattern in a Jupyter notebook"""
-    results = []
-    flags = 0 if case_sensitive else re.IGNORECASE
-
-    try:
-        compiled = re.compile(pattern, flags)
-    except re.error:
-        compiled = re.compile(re.escape(pattern), flags)
-
-    try:
-        with open(notebook_path, "r", encoding="utf-8") as f:
-            notebook = json.load(f)
-
-        cells = notebook.get("cells", [])
-        for idx, cell in enumerate(cells):
-            current_type = cell.get("cell_type", "code")
-
-            # Filter by cell type if specified
-            if cell_type and current_type != cell_type:
-                continue
-
-            source = cell.get("source", [])
-            if isinstance(source, list):
-                source = "".join(source)
-
-            if compiled.search(source):
-                # Find matching lines
-                matching_lines = []
-                for line_num, line in enumerate(source.split("\n"), 1):
-                    if compiled.search(line):
-                        matching_lines.append({
-                            "line": line_num,
-                            "content": line.strip()[:150]
-                        })
-
-                results.append({
-                    "file_path": notebook_path,
-                    "cell_index": idx,
-                    "cell_type": current_type,
-                    "content": source[:300] + "..." if len(source) > 300 else source,
-                    "matching_lines": matching_lines[:5],
-                    "match_type": "cell",
-                })
-    except Exception:
-        pass
-
-    return results
+    notebook_path: Optional[str],
+    path: str,
+    max_results: int,
+) -> str:
+    """Build a command to find notebooks for searching."""
+    if notebook_path:
+        return f"echo '{notebook_path}'"
+    else:
+        return (
+            f"find {path} -name '*.ipynb' -type f 2>/dev/null | head -n {max_results}"
+        )
 
 
 @tool(args_schema=SearchWorkspaceInput)
@@ -135,84 +146,65 @@ def search_workspace_tool(
     path: str = ".",
     max_results: int = 50,
     case_sensitive: bool = False,
-    workspace_root: str = "."
+    execution_result: Optional[Dict[str, Any]] = None,
+    workspace_root: str = ".",
 ) -> Dict[str, Any]:
     """
     Search for a pattern across files in the workspace.
-    
+
+    This tool is executed on the client side using subprocess (grep/ripgrep).
     Searches both regular files and Jupyter notebooks.
-    For notebooks, searches within cell contents.
-    
+
     Args:
         pattern: Search pattern (regex or text)
         file_types: File patterns to search (default: ['*.py', '*.ipynb'])
-        path: Directory to search in
+        path: Directory to search in (relative to workspace)
         max_results: Maximum number of results to return
         case_sensitive: Whether search is case-sensitive
-        
+
     Returns:
-        Dict with search results
+        Dict with search results or pending_execution status
     """
-    import fnmatch
-
     if file_types is None:
         file_types = ["*.py", "*.ipynb"]
 
-    results = []
-    files_searched = 0
-
-    try:
-        search_path = os.path.normpath(os.path.join(workspace_root, path))
-
-        for root, _, filenames in os.walk(search_path):
-            for filename in filenames:
-                # Check if file matches any pattern
-                if not any(fnmatch.fnmatch(filename, ft) for ft in file_types):
-                    continue
-
-                file_path = os.path.join(root, filename)
-                rel_path = os.path.relpath(file_path, workspace_root)
-                files_searched += 1
-
-                if filename.endswith(".ipynb"):
-                    # Search in notebook
-                    matches = _search_in_notebook(
-                        file_path, pattern, None, case_sensitive
-                    )
-                    for m in matches:
-                        m["file_path"] = rel_path
-                    results.extend(matches)
-                else:
-                    # Search in regular file
-                    matches = _search_in_file(file_path, pattern, case_sensitive)
-                    for m in matches:
-                        m["file_path"] = rel_path
-                    results.extend(matches)
-
-                if len(results) >= max_results:
-                    break
-
-            if len(results) >= max_results:
-                break
-
-        return {
-            "tool": "search_workspace",
-            "success": True,
+    # Build the search command
+    command, tool_used = _build_grep_command(
+        pattern=pattern,
+        file_types=file_types,
+        path=path,
+        case_sensitive=case_sensitive,
+        max_results=max_results,
+    )
+
+    response: Dict[str, Any] = {
+        "tool": "search_workspace_tool",
+        "parameters": {
             "pattern": pattern,
+            "file_types": file_types,
             "path": path,
-            "files_searched": files_searched,
-            "total_results": len(results),
-            "results": results[:max_results],
-            "truncated": len(results) > max_results,
-        }
-
-    except Exception as e:
-        return {
-            "tool": "search_workspace",
-            "success": False,
-            "error": f"Search failed: {str(e)}",
-            "pattern": pattern,
-        }
+            "max_results": max_results,
+            "case_sensitive": case_sensitive,
+        },
+        "command": command,
+        "tool_used": tool_used,
+        "status": "pending_execution",
+        "message": "Search queued for execution by client",
+    }
+
+    if execution_result is not None:
+        response["execution_result"] = execution_result
+        response["status"] = "complete"
+        response["message"] = "Search executed with client-reported results"
+        # Parse the execution result to extract search results
+        if isinstance(execution_result, dict):
+            response["success"] = execution_result.get("success", False)
+            response["results"] = execution_result.get("results", [])
+            response["total_results"] = execution_result.get("total_results", 0)
+            if "error" in execution_result:
+                response["error"] = execution_result["error"]
+
+    return response
 
 
 @tool(args_schema=SearchNotebookCellsInput)
@@ -222,84 +214,74 @@ def search_notebook_cells_tool(
     cell_type: Optional[str] = None,
     max_results: int = 30,
     case_sensitive: bool = False,
-    workspace_root: str = "."
+    execution_result: Optional[Dict[str, Any]] = None,
+    workspace_root: str = ".",
 ) -> Dict[str, Any]:
     """
     Search for a pattern in Jupyter notebook cells.
-    
+
+    This tool is executed on the client side.
     Can search a specific notebook or all notebooks in workspace.
     Optionally filter by cell type (code/markdown).
-    
+
     Args:
         pattern: Search pattern (regex or text)
         notebook_path: Specific notebook to search (None = all)
         cell_type: Filter by cell type ('code', 'markdown', or None)
         max_results: Maximum number of results
         case_sensitive: Whether search is case-sensitive
-        
+
     Returns:
-        Dict with matching cells
+        Dict with matching cells or pending_execution status
     """
-    results = []
-    notebooks_searched = 0
-
-    try:
-        if notebook_path:
-            # Search specific notebook
-            full_path = os.path.normpath(
-                os.path.join(workspace_root, notebook_path)
-            )
-            if os.path.exists(full_path) and full_path.endswith(".ipynb"):
-                matches = _search_in_notebook(
-                    full_path, pattern, cell_type, case_sensitive
-                )
-                for m in matches:
-                    m["file_path"] = notebook_path
-                results.extend(matches)
-                notebooks_searched = 1
-        else:
-            # Search all notebooks
-            for root, _, filenames in os.walk(workspace_root):
-                for filename in filenames:
-                    if not filename.endswith(".ipynb"):
-                        continue
-
-                    file_path = os.path.join(root, filename)
-                    rel_path = os.path.relpath(file_path, workspace_root)
-                    notebooks_searched += 1
-
-                    matches = _search_in_notebook(
-                        file_path, pattern, cell_type, case_sensitive
-                    )
-                    for m in matches:
-                        m["file_path"] = rel_path
-                    results.extend(matches)
-
-                    if len(results) >= max_results:
-                        break
-
-                if len(results) >= max_results:
-                    break
-
-        return {
-            "tool": "search_notebook_cells",
-            "success": True,
+    # Build find command for notebooks
+    find_command = _build_notebook_search_command(
+        pattern=pattern,
+        notebook_path=notebook_path,
+        path=".",
+        max_results=max_results,
+    )
+
+    response: Dict[str, Any] = {
+        "tool": "search_notebook_cells_tool",
+        "parameters": {
             "pattern": pattern,
             "notebook_path": notebook_path,
             "cell_type": cell_type,
-            "notebooks_searched": notebooks_searched,
-            "total_results": len(results),
-            "results": results[:max_results],
-            "truncated": len(results) > max_results,
-        }
-
-    except Exception as e:
-        return {
-            "tool": "search_notebook_cells",
-            "success": False,
-            "error": f"Search failed: {str(e)}",
-            "pattern": pattern,
-        }
+            "max_results": max_results,
+            "case_sensitive": case_sensitive,
+        },
+        "find_command": find_command,
+        "status": "pending_execution",
+        "message": "Notebook search queued for execution by client",
+    }
+
+    if execution_result is not None:
+        response["execution_result"] = execution_result
+        response["status"] = "complete"
+        response["message"] = "Notebook search executed with client-reported results"
+        # Parse the execution result
+        if isinstance(execution_result, dict):
+            response["success"] = execution_result.get("success", False)
+            response["results"] = execution_result.get("results", [])
+            response["total_results"] = execution_result.get("total_results", 0)
+            response["notebooks_searched"] = execution_result.get(
+                "notebooks_searched", 0
+            )
+            if "error" in execution_result:
+                response["error"] = execution_result["error"]
+
+    return response
+
+
+def create_search_tools(workspace_root: str = ".") -> List:
+    """
+    Create search tools (for backward compatibility).
+
+    Note: workspace_root is not used since tools return pending_execution
+    and actual execution happens on the client side.
+    """
+    return [search_workspace_tool, search_notebook_cells_tool]
 
 
 # Export all tools

agent_server/langchain/tools/shell_tools.py

@@ -0,0 +1,54 @@
+"""
+Shell command tool for LangChain agent.
+
+Execution is handled by the Jupyter extension (client) after HITL approval.
+"""
+
+from typing import Any, Dict, Optional
+
+from langchain_core.tools import tool
+from pydantic import BaseModel, Field
+
+
+class ExecuteCommandInput(BaseModel):
+    """Input schema for execute_command_tool."""
+
+    command: str = Field(description="Shell command to execute")
+    stdin: Optional[str] = Field(
+        default="y\n",
+        description="Input to provide to the command for interactive prompts (default: 'y\\n' for yes/no prompts)",
+    )
+    timeout: Optional[int] = Field(
+        default=600000, description="Timeout in milliseconds"
+    )
+    execution_result: Optional[Dict[str, Any]] = Field(
+        default=None, description="Optional execution result payload from the client"
+    )
+
+
+@tool(args_schema=ExecuteCommandInput)
+def execute_command_tool(
+    command: str,
+    stdin: Optional[str] = "y\n",
+    timeout: Optional[int] = 600000,
+    execution_result: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """
+    Execute a shell command via the client (Jupyter extension).
+    The agent server only coordinates the request. The client performs the
+    actual execution after user approval and returns execution_result.
+
+    - Interactive prompts are auto-answered with "y" by default
+    - NEVER run long-running commands (servers, watch, dev) or endless processes
+    """
+    response: Dict[str, Any] = {
+        "tool": "execute_command_tool",
+        "parameters": {"command": command, "stdin": stdin, "timeout": timeout},
+        "status": "pending_execution",
+        "message": "Shell command queued for execution by client",
+    }
+    if execution_result is not None:
+        response["execution_result"] = execution_result
+        response["status"] = "complete"
+        response["message"] = "Shell command executed with client-reported results"
+    return response

agent_server/main.py

@@ -19,6 +19,7 @@ from agent_server.routers import agent, chat, config, file_resolver, health, rag
 # Optional LangChain router (requires langchain dependencies)
 try:
     from agent_server.routers import langchain_agent
+
     LANGCHAIN_AVAILABLE = True
 except ImportError:
     LANGCHAIN_AVAILABLE = False
@@ -153,7 +154,9 @@ app.include_router(file_resolver.router, prefix="/file", tags=["File Resolution"
 
 # Register LangChain agent router if available
 if LANGCHAIN_AVAILABLE:
-    app.include_router(langchain_agent.router, prefix="/agent", tags=["LangChain Agent"])
+    app.include_router(
+        langchain_agent.router, prefix="/agent", tags=["LangChain Agent"]
+    )
     logger.info("LangChain agent router registered")
 
 

agent_server/routers/health.py

@@ -36,7 +36,7 @@ async def health_check() -> HealthResponse:
 async def status_check() -> HealthResponse:
     """
     Status check endpoint (alias for health check).
-    
+
     Returns server status, timestamp, and version.
     """
     return HealthResponse(