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

jupyter_ext/resource_usage.py

@@ -0,0 +1,180 @@
+"""
+Resource Usage Utilities for the Jupyter server host.
+
+Collects CPU, memory, disk, and GPU usage to guide client-side execution.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from typing import Optional
+
+try:
+    import psutil
+except ImportError:  # pragma: no cover - optional dependency fallback
+    psutil = None
+
+
+def _read_cgroup_value(path: str) -> Optional[str]:
+    try:
+        with open(path, "r", encoding="utf-8") as handle:
+            return handle.read().strip()
+    except Exception:
+        return None
+
+
+def _format_gb(value_bytes: float) -> float:
+    return round(value_bytes / (1024**3), 2)
+
+
+def _safe_float(value: Optional[str]) -> Optional[float]:
+    if value is None:
+        return None
+    try:
+        return float(value)
+    except (TypeError, ValueError):
+        return None
+
+
+def _safe_cpu_count() -> Optional[int]:
+    count = os.cpu_count()
+    return int(count) if count is not None else None
+
+
+def get_integrated_resources(
+    env_type: str = "auto", workspace_root: str = "."
+) -> dict[str, object]:
+    """
+    Collect resource usage snapshot for the Jupyter server host.
+
+    Args:
+        env_type: "auto", "pod", or "host" (auto detects Kubernetes)
+        workspace_root: Path used to compute disk usage
+
+    Returns:
+        JSON-serializable resource snapshot dict.
+    """
+    is_pod = False
+    if env_type == "auto":
+        is_pod = os.path.exists("/var/run/secrets/kubernetes.io") or bool(
+            os.environ.get("KUBERNETES_SERVICE_HOST")
+        )
+    else:
+        is_pod = env_type.lower() == "pod"
+
+    environment = "Kubernetes Pod" if is_pod else "Host/VM"
+    cpu: dict[str, Optional[float]] = {"cores": None, "usage_percent": None}
+    memory: dict[str, Optional[float]] = {
+        "total_gb": None,
+        "available_gb": None,
+        "used_gb": None,
+    }
+    disk: dict[str, Optional[object]] = {
+        "path": None,
+        "total_gb": None,
+        "free_gb": None,
+        "used_gb": None,
+    }
+    gpus: list[dict[str, Optional[object]]] = []
+    gpu_status = "not_detected"
+
+    if is_pod:
+        try:
+            cpu_max = _read_cgroup_value("/sys/fs/cgroup/cpu.max")
+            if cpu_max:
+                quota, period = cpu_max.split()
+                if quota != "max":
+                    cpu_limit = float(quota) / float(period)
+                else:
+                    cpu_limit = psutil.cpu_count() if psutil else _safe_cpu_count() or 0
+            else:
+                cpu_limit = psutil.cpu_count() if psutil else _safe_cpu_count() or 0
+
+            mem_limit_raw = _read_cgroup_value("/sys/fs/cgroup/memory.max")
+            mem_current_raw = _read_cgroup_value("/sys/fs/cgroup/memory.current")
+
+            if mem_limit_raw and mem_limit_raw != "max":
+                mem_limit_gb = _format_gb(float(mem_limit_raw))
+            else:
+                mem_limit_gb = (
+                    _format_gb(psutil.virtual_memory().total) if psutil else 0.0
+                )
+
+            mem_used_gb = _format_gb(float(mem_current_raw)) if mem_current_raw else 0.0
+
+            cpu["cores"] = round(cpu_limit, 2)
+            memory["total_gb"] = mem_limit_gb
+            memory["used_gb"] = mem_used_gb
+            if mem_limit_gb is not None and mem_used_gb is not None:
+                memory["available_gb"] = round(mem_limit_gb - mem_used_gb, 2)
+        except Exception:
+            cpu_count = psutil.cpu_count() if psutil else _safe_cpu_count()
+            cpu["cores"] = float(cpu_count) if cpu_count is not None else None
+            if psutil:
+                vm = psutil.virtual_memory()
+                memory["total_gb"] = _format_gb(vm.total)
+                memory["available_gb"] = _format_gb(vm.available)
+                memory["used_gb"] = _format_gb(vm.total - vm.available)
+    else:
+        cpu_count = psutil.cpu_count() if psutil else _safe_cpu_count()
+        cpu_percent = psutil.cpu_percent() if psutil else None
+        cpu["cores"] = float(cpu_count) if cpu_count is not None else None
+        cpu["usage_percent"] = cpu_percent
+        if psutil:
+            vm = psutil.virtual_memory()
+            memory["total_gb"] = _format_gb(vm.total)
+            memory["available_gb"] = _format_gb(vm.available)
+            memory["used_gb"] = _format_gb(vm.total - vm.available)
+
+    disk_path = workspace_root if workspace_root else "."
+    if not os.path.exists(disk_path):
+        disk_path = "."
+    disk["path"] = os.path.abspath(disk_path)
+    try:
+        total, used, free = shutil.disk_usage(disk_path)
+        disk["total_gb"] = _format_gb(total)
+        disk["free_gb"] = _format_gb(free)
+        disk["used_gb"] = _format_gb(used)
+    except Exception:
+        pass
+
+    if shutil.which("nvidia-smi"):
+        gpu_status = "available"
+        try:
+            result = subprocess.check_output(
+                [
+                    "nvidia-smi",
+                    "--query-gpu=name,utilization.gpu,memory.used,memory.total",
+                    "--format=csv,noheader,nounits",
+                ],
+                encoding="utf-8",
+                timeout=2,
+            ).strip()
+            if result:
+                for line in result.split("\n"):
+                    name, gpu_util, mem_used, mem_total = [
+                        value.strip() for value in line.split(",")
+                    ]
+                    gpus.append(
+                        {
+                            "name": name,
+                            "utilization_percent": _safe_float(gpu_util),
+                            "memory_used_mb": _safe_float(mem_used),
+                            "memory_total_mb": _safe_float(mem_total),
+                        }
+                    )
+        except Exception:
+            gpu_status = "unavailable"
+    else:
+        gpu_status = "not_detected"
+
+    return {
+        "environment": environment,
+        "cpu": cpu,
+        "memory": memory,
+        "disk": disk,
+        "gpus": gpus,
+        "gpu_status": gpu_status,
+    }

jupyter_ext/tests/test_handlers.py

@@ -0,0 +1,58 @@
+import os
+
+from jupyter_ext.handlers import (
+    DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS,
+    _append_stream_output,
+    _resolve_path_in_workspace,
+    _resolve_stream_timeout_ms,
+    _resolve_timeout_ms,
+)
+
+
+def test_resolve_path_strips_duplicate_cwd_prefix() -> None:
+    workspace_root = "/workspace"
+    requested_cwd = os.path.join(workspace_root, "extensions", "jupyter")
+    path = os.path.join("extensions", "jupyter", "cal.py")
+    resolved = _resolve_path_in_workspace(path, workspace_root, requested_cwd)
+    assert resolved == os.path.join(requested_cwd, "cal.py")
+
+
+def test_resolve_path_relative_to_cwd() -> None:
+    workspace_root = "/workspace"
+    requested_cwd = os.path.join(workspace_root, "extensions", "jupyter")
+    resolved = _resolve_path_in_workspace("cal.py", workspace_root, requested_cwd)
+    assert resolved == os.path.join(requested_cwd, "cal.py")
+
+
+def test_resolve_timeout_ms_uses_default_on_invalid() -> None:
+    assert _resolve_timeout_ms(None) == DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS
+    assert _resolve_timeout_ms("not-a-number") == DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS
+    assert _resolve_timeout_ms(-1) == DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS
+
+
+def test_resolve_timeout_ms_accepts_value() -> None:
+    assert _resolve_timeout_ms(120000) == 120000
+
+
+def test_resolve_stream_timeout_ms_uses_default_on_invalid() -> None:
+    assert _resolve_stream_timeout_ms(None) == DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS
+    assert (
+        _resolve_stream_timeout_ms("not-a-number") == DEFAULT_EXECUTE_COMMAND_TIMEOUT_MS
+    )
+
+
+def test_resolve_stream_timeout_ms_disables_on_non_positive() -> None:
+    assert _resolve_stream_timeout_ms(0) is None
+    assert _resolve_stream_timeout_ms(-1) is None
+
+
+def test_append_stream_output_appends() -> None:
+    output, truncated = _append_stream_output("abc", "def", max_bytes=10)
+    assert output == "abcdef"
+    assert truncated is False
+
+
+def test_append_stream_output_truncates() -> None:
+    output, truncated = _append_stream_output("abc", "def", max_bytes=5)
+    assert output == "abcde"
+    assert truncated is True

agent_server/langchain/executors/jupyter_executor.py

@@ -1,429 +0,0 @@
-"""
-Jupyter Executor (Embedded Mode)
-
-Provides direct access to Jupyter kernel for code execution
-when running inside JupyterLab server.
-
-This executor uses the Jupyter server's kernel manager to:
-- Execute code in the current notebook's kernel
-- Add cells to the notebook
-- Retrieve execution results
-"""
-
-import asyncio
-import logging
-from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional
-
-logger = logging.getLogger(__name__)
-
-
-@dataclass
-class ExecutionResult:
-    """Result of code execution in Jupyter kernel"""
-    success: bool
-    output: str = ""
-    error_type: Optional[str] = None
-    error_message: Optional[str] = None
-    traceback: Optional[List[str]] = None
-    execution_count: int = 0
-    cell_index: int = -1
-    display_data: List[Dict[str, Any]] = field(default_factory=list)
-
-    def to_dict(self) -> Dict[str, Any]:
-        return {
-            "success": self.success,
-            "output": self.output,
-            "error_type": self.error_type,
-            "error_message": self.error_message,
-            "traceback": self.traceback,
-            "execution_count": self.execution_count,
-            "cell_index": self.cell_index,
-            "display_data": self.display_data,
-        }
-
-
-class JupyterExecutor:
-    """
-    Executes code in Jupyter kernel (Embedded Mode).
-    
-    In Embedded Mode, this class directly accesses the Jupyter server's
-    kernel manager and contents manager to execute code and modify notebooks.
-    
-    Usage:
-        executor = JupyterExecutor()
-        await executor.initialize(kernel_id, notebook_path)
-        result = await executor.execute_code("print('hello')")
-    """
-
-    def __init__(self):
-        self._kernel_manager = None
-        self._contents_manager = None
-        self._kernel_id: Optional[str] = None
-        self._notebook_path: Optional[str] = None
-        self._kernel_client = None
-        self._initialized = False
-
-    @property
-    def is_initialized(self) -> bool:
-        return self._initialized
-
-    async def initialize(
-        self,
-        kernel_id: str,
-        notebook_path: str,
-        kernel_manager: Any = None,
-        contents_manager: Any = None,
-    ) -> bool:
-        """
-        Initialize the executor with kernel and notebook information.
-        
-        Args:
-            kernel_id: ID of the kernel to use
-            notebook_path: Path to the notebook file
-            kernel_manager: Jupyter's MappingKernelManager (optional, auto-detect)
-            contents_manager: Jupyter's ContentsManager (optional, auto-detect)
-            
-        Returns:
-            True if initialization successful
-        """
-        self._kernel_id = kernel_id
-        self._notebook_path = notebook_path
-
-        # Try to get kernel manager from Jupyter server if not provided
-        if kernel_manager is None:
-            kernel_manager = self._get_kernel_manager()
-
-        if contents_manager is None:
-            contents_manager = self._get_contents_manager()
-
-        self._kernel_manager = kernel_manager
-        self._contents_manager = contents_manager
-
-        if self._kernel_manager is None:
-            logger.warning("Kernel manager not available. Running in mock mode.")
-            self._initialized = True
-            return True
-
-        # Get kernel client
-        try:
-            self._kernel_client = self._kernel_manager.get_kernel(kernel_id)
-            self._initialized = True
-            logger.info(f"JupyterExecutor initialized with kernel {kernel_id}")
-            return True
-        except Exception as e:
-            logger.error(f"Failed to get kernel client: {e}")
-            self._initialized = False
-            return False
-
-    def _get_kernel_manager(self) -> Optional[Any]:
-        """Try to get kernel manager from Jupyter server app"""
-        try:
-            from jupyter_server.serverapp import ServerApp
-            app = ServerApp.instance()
-            return app.kernel_manager
-        except Exception:
-            try:
-                # Fallback for older versions
-                from notebook.notebookapp import NotebookApp
-                app = NotebookApp.instance()
-                return app.kernel_manager
-            except Exception:
-                return None
-
-    def _get_contents_manager(self) -> Optional[Any]:
-        """Try to get contents manager from Jupyter server app"""
-        try:
-            from jupyter_server.serverapp import ServerApp
-            app = ServerApp.instance()
-            return app.contents_manager
-        except Exception:
-            try:
-                from notebook.notebookapp import NotebookApp
-                app = NotebookApp.instance()
-                return app.contents_manager
-            except Exception:
-                return None
-
-    async def execute_code(
-        self,
-        code: str,
-        timeout: float = 60.0,
-        add_to_notebook: bool = True,
-    ) -> ExecutionResult:
-        """
-        Execute Python code in the Jupyter kernel.
-        
-        Args:
-            code: Python code to execute
-            timeout: Execution timeout in seconds
-            add_to_notebook: Whether to add the code as a new cell
-            
-        Returns:
-            ExecutionResult with output or error
-        """
-        if not self._initialized:
-            return ExecutionResult(
-                success=False,
-                error_type="NotInitialized",
-                error_message="Executor not initialized. Call initialize() first."
-            )
-
-        # If no kernel manager, use mock execution
-        if self._kernel_manager is None:
-            return await self._mock_execute(code)
-
-        try:
-            # Add cell to notebook if requested
-            cell_index = -1
-            if add_to_notebook and self._contents_manager:
-                cell_index = await self._add_cell_to_notebook(code)
-
-            # Execute code in kernel
-            result = await self._execute_in_kernel(code, timeout)
-            result.cell_index = cell_index
-
-            return result
-
-        except asyncio.TimeoutError:
-            return ExecutionResult(
-                success=False,
-                error_type="TimeoutError",
-                error_message=f"Execution timed out after {timeout} seconds"
-            )
-        except Exception as e:
-            logger.error(f"Execution failed: {e}")
-            return ExecutionResult(
-                success=False,
-                error_type=type(e).__name__,
-                error_message=str(e)
-            )
-
-    async def _execute_in_kernel(
-        self,
-        code: str,
-        timeout: float
-    ) -> ExecutionResult:
-        """Execute code using kernel client"""
-        # This is a simplified implementation
-        # In production, you would use jupyter_client's async API
-
-        try:
-            from jupyter_client import KernelClient
-
-            # Get connection info from kernel manager
-            km = self._kernel_manager
-            kernel = km.get_kernel(self._kernel_id)
-
-            # Create a client and execute
-            client = kernel.client()
-            client.start_channels()
-
-            try:
-                # Send execute request
-                msg_id = client.execute(code)
-
-                # Wait for results
-                output_parts = []
-                error_info = None
-                execution_count = 0
-                display_data = []
-
-                deadline = asyncio.get_event_loop().time() + timeout
-
-                while True:
-                    remaining = deadline - asyncio.get_event_loop().time()
-                    if remaining <= 0:
-                        raise asyncio.TimeoutError()
-
-                    try:
-                        msg = client.get_iopub_msg(timeout=min(remaining, 1.0))
-                    except Exception:
-                        continue
-
-                    if msg["parent_header"].get("msg_id") != msg_id:
-                        continue
-
-                    msg_type = msg["msg_type"]
-                    content = msg["content"]
-
-                    if msg_type == "stream":
-                        output_parts.append(content.get("text", ""))
-                    elif msg_type == "execute_result":
-                        output_parts.append(str(content.get("data", {}).get("text/plain", "")))
-                        execution_count = content.get("execution_count", 0)
-                    elif msg_type == "display_data":
-                        display_data.append(content.get("data", {}))
-                    elif msg_type == "error":
-                        error_info = {
-                            "ename": content.get("ename", "Error"),
-                            "evalue": content.get("evalue", ""),
-                            "traceback": content.get("traceback", []),
-                        }
-                    elif msg_type == "status" and content.get("execution_state") == "idle":
-                        break
-
-                if error_info:
-                    return ExecutionResult(
-                        success=False,
-                        output="".join(output_parts),
-                        error_type=error_info["ename"],
-                        error_message=error_info["evalue"],
-                        traceback=error_info["traceback"],
-                        execution_count=execution_count,
-                        display_data=display_data,
-                    )
-
-                return ExecutionResult(
-                    success=True,
-                    output="".join(output_parts),
-                    execution_count=execution_count,
-                    display_data=display_data,
-                )
-
-            finally:
-                client.stop_channels()
-
-        except ImportError:
-            logger.warning("jupyter_client not available, using mock execution")
-            return await self._mock_execute(code)
-        except Exception as e:
-            return ExecutionResult(
-                success=False,
-                error_type=type(e).__name__,
-                error_message=str(e),
-            )
-
-    async def _add_cell_to_notebook(self, code: str) -> int:
-        """Add a new code cell to the notebook"""
-        if not self._contents_manager or not self._notebook_path:
-            return -1
-
-        try:
-            # Read current notebook
-            model = self._contents_manager.get(self._notebook_path, content=True)
-            notebook = model["content"]
-
-            # Create new cell
-            new_cell = {
-                "cell_type": "code",
-                "execution_count": None,
-                "metadata": {},
-                "outputs": [],
-                "source": code,
-            }
-
-            # Add cell
-            notebook["cells"].append(new_cell)
-            cell_index = len(notebook["cells"]) - 1
-
-            # Save notebook
-            self._contents_manager.save(model, self._notebook_path)
-
-            return cell_index
-
-        except Exception as e:
-            logger.error(f"Failed to add cell to notebook: {e}")
-            return -1
-
-    async def add_markdown_cell(self, content: str) -> int:
-        """Add a markdown cell to the notebook"""
-        if not self._contents_manager or not self._notebook_path:
-            return -1
-
-        try:
-            model = self._contents_manager.get(self._notebook_path, content=True)
-            notebook = model["content"]
-
-            new_cell = {
-                "cell_type": "markdown",
-                "metadata": {},
-                "source": content,
-            }
-
-            notebook["cells"].append(new_cell)
-            cell_index = len(notebook["cells"]) - 1
-
-            self._contents_manager.save(model, self._notebook_path)
-
-            return cell_index
-
-        except Exception as e:
-            logger.error(f"Failed to add markdown cell: {e}")
-            return -1
-
-    async def _mock_execute(self, code: str) -> ExecutionResult:
-        """Mock execution for testing or when kernel is not available"""
-        logger.info(f"Mock executing code: {code[:100]}...")
-
-        # Simple mock that just returns success
-        return ExecutionResult(
-            success=True,
-            output=f"[Mock] Code executed successfully:\n{code[:200]}",
-            execution_count=1,
-        )
-
-    async def get_notebook_state(self) -> Dict[str, Any]:
-        """Get current notebook state"""
-        if not self._contents_manager or not self._notebook_path:
-            return {
-                "cell_count": 0,
-                "imported_libraries": [],
-                "defined_variables": [],
-            }
-
-        try:
-            model = self._contents_manager.get(self._notebook_path, content=True)
-            notebook = model["content"]
-            cells = notebook.get("cells", [])
-
-            # Extract imports and variables from code cells
-            imported_libraries = set()
-            defined_variables = set()
-
-            import re
-            import_pattern = re.compile(r'^(?:import|from)\s+([\w.]+)', re.MULTILINE)
-            var_pattern = re.compile(r'^(\w+)\s*=', re.MULTILINE)
-
-            for cell in cells:
-                if cell.get("cell_type") != "code":
-                    continue
-
-                source = cell.get("source", "")
-                if isinstance(source, list):
-                    source = "".join(source)
-
-                # Find imports
-                for match in import_pattern.finditer(source):
-                    lib = match.group(1).split(".")[0]
-                    imported_libraries.add(lib)
-
-                # Find variable definitions
-                for match in var_pattern.finditer(source):
-                    defined_variables.add(match.group(1))
-
-            return {
-                "cell_count": len(cells),
-                "imported_libraries": list(imported_libraries),
-                "defined_variables": list(defined_variables),
-            }
-
-        except Exception as e:
-            logger.error(f"Failed to get notebook state: {e}")
-            return {
-                "cell_count": 0,
-                "imported_libraries": [],
-                "defined_variables": [],
-            }
-
-
-# Singleton instance
-_executor_instance: Optional[JupyterExecutor] = None
-
-
-def get_jupyter_executor() -> JupyterExecutor:
-    """Get or create the JupyterExecutor singleton"""
-    global _executor_instance
-    if _executor_instance is None:
-        _executor_instance = JupyterExecutor()
-    return _executor_instance

agent_server/langchain/middleware/__init__.py

@@ -1,36 +0,0 @@
-"""
-LangChain Middleware for Jupyter Agent
-
-Middleware stack (execution order):
-1. RAGMiddleware: Inject RAG context before model calls
-2. CodeSearchMiddleware: Search workspace/notebook for relevant code
-3. ValidationMiddleware: Validate code before execution
-4. JupyterExecutionMiddleware: Execute code in Jupyter kernel
-5. ErrorHandlingMiddleware: Classify errors and decide recovery strategy
-
-Built-in middleware used:
-- SummarizationMiddleware: Compress long conversations
-- ModelRetryMiddleware: Retry on rate limits
-- ToolRetryMiddleware: Retry failed tool calls
-- ModelCallLimitMiddleware: Prevent infinite loops
-"""
-
-from agent_server.langchain.middleware.code_search_middleware import (
-    CodeSearchMiddleware,
-)
-from agent_server.langchain.middleware.error_handling_middleware import (
-    ErrorHandlingMiddleware,
-)
-from agent_server.langchain.middleware.jupyter_execution_middleware import (
-    JupyterExecutionMiddleware,
-)
-from agent_server.langchain.middleware.rag_middleware import RAGMiddleware
-from agent_server.langchain.middleware.validation_middleware import ValidationMiddleware
-
-__all__ = [
-    "RAGMiddleware",
-    "CodeSearchMiddleware",
-    "ValidationMiddleware",
-    "JupyterExecutionMiddleware",
-    "ErrorHandlingMiddleware",
-]