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

agent_server/langchain/middleware/jupyter_execution_middleware.py

@@ -1,301 +0,0 @@
-"""
-Jupyter Execution Middleware
-
-Handles actual execution of code in Jupyter kernel.
-Wraps jupyter_cell tool calls to execute code and return results.
-"""
-
-import asyncio
-import logging
-from typing import Any, Callable, Dict, Optional
-
-from agent_server.langchain.executors.jupyter_executor import (
-    JupyterExecutor,
-    get_jupyter_executor,
-)
-from agent_server.langchain.state import AgentRuntime, AgentState
-
-logger = logging.getLogger(__name__)
-
-
-class JupyterExecutionMiddleware:
-    """
-    Middleware that executes code in Jupyter kernel.
-    
-    This middleware:
-    1. Wraps jupyter_cell and markdown tool calls
-    2. Executes code using JupyterExecutor
-    3. Returns execution results to the agent
-    4. Handles execution errors and timeouts
-    
-    Uses @wrap_tool_call hook pattern from LangChain middleware.
-    """
-
-    def __init__(
-        self,
-        executor: Optional[JupyterExecutor] = None,
-        timeout: float = 60.0,
-        add_to_notebook: bool = True,
-        enabled: bool = True,
-    ):
-        """
-        Initialize Jupyter execution middleware.
-        
-        Args:
-            executor: JupyterExecutor instance
-            timeout: Execution timeout in seconds
-            add_to_notebook: Whether to add cells to notebook
-            enabled: Whether execution is enabled
-        """
-        self._executor = executor
-        self._timeout = timeout
-        self._add_to_notebook = add_to_notebook
-        self._enabled = enabled
-
-    @property
-    def name(self) -> str:
-        return "JupyterExecutionMiddleware"
-
-    def _get_executor(self, runtime: AgentRuntime) -> JupyterExecutor:
-        """Get executor from runtime or create default"""
-        if runtime and runtime.jupyter_executor:
-            return runtime.jupyter_executor
-        if self._executor:
-            return self._executor
-        return get_jupyter_executor()
-
-    async def wrap_tool_call(
-        self,
-        tool_name: str,
-        tool_input: Dict[str, Any],
-        next_call: Callable,
-        state: AgentState,
-        runtime: AgentRuntime,
-    ) -> Dict[str, Any]:
-        """
-        Wrap tool calls to handle Jupyter execution.
-        
-        For jupyter_cell tools, executes code in kernel.
-        For other tools, passes through to next handler.
-        
-        Args:
-            tool_name: Name of the tool
-            tool_input: Tool input parameters
-            next_call: Next handler in chain
-            state: Current agent state
-            runtime: Agent runtime context
-            
-        Returns:
-            Tool execution result
-        """
-        if not self._enabled:
-            return await next_call(tool_name, tool_input)
-
-        # Handle jupyter_cell execution
-        if tool_name == "jupyter_cell_tool":
-            return await self._execute_jupyter_cell(
-                tool_input, state, runtime
-            )
-
-        # Handle markdown cell
-        if tool_name == "markdown_tool":
-            return await self._add_markdown_cell(
-                tool_input, state, runtime
-            )
-
-        # Handle final_answer
-        if tool_name == "final_answer_tool":
-            return await self._handle_final_answer(
-                tool_input, state, runtime
-            )
-
-        # Pass through for other tools
-        return await next_call(tool_name, tool_input)
-
-    async def _execute_jupyter_cell(
-        self,
-        tool_input: Dict[str, Any],
-        state: AgentState,
-        runtime: AgentRuntime,
-    ) -> Dict[str, Any]:
-        """Execute code in Jupyter cell"""
-        code = tool_input.get("code", "")
-        if not code:
-            return {
-                "success": False,
-                "error": "No code provided",
-            }
-
-        # Check for validation failure
-        if tool_input.get("validation_failed"):
-            errors = tool_input.get("validation_errors", [])
-            return {
-                "success": False,
-                "error_type": "ValidationError",
-                "error": "Code validation failed",
-                "validation_errors": errors,
-                "message": "Please fix validation errors and try again",
-            }
-
-        executor = self._get_executor(runtime)
-
-        # Initialize executor if needed
-        if not executor.is_initialized:
-            notebook_context = state.get("notebook_context", {})
-            kernel_id = notebook_context.get("kernel_id")
-            notebook_path = notebook_context.get("notebook_path")
-
-            if kernel_id and notebook_path:
-                await executor.initialize(kernel_id, notebook_path)
-            else:
-                logger.warning("Kernel/notebook not available, using mock execution")
-
-        try:
-            result = await asyncio.wait_for(
-                executor.execute_code(
-                    code,
-                    timeout=self._timeout,
-                    add_to_notebook=self._add_to_notebook,
-                ),
-                timeout=self._timeout + 5,
-            )
-
-            # Update execution history in state
-            execution_history = state.get("execution_history", [])
-            execution_history.append(result.to_dict())
-            state["execution_history"] = execution_history
-
-            if result.success:
-                return {
-                    "success": True,
-                    "output": result.output,
-                    "execution_count": result.execution_count,
-                    "cell_index": result.cell_index,
-                    "display_data": result.display_data,
-                }
-            else:
-                # Update error state
-                state["error_count"] = state.get("error_count", 0) + 1
-                state["last_error"] = {
-                    "error_type": result.error_type,
-                    "error_message": result.error_message,
-                    "traceback": result.traceback,
-                }
-
-                return {
-                    "success": False,
-                    "error_type": result.error_type,
-                    "error": result.error_message,
-                    "traceback": result.traceback,
-                    "output": result.output,
-                }
-
-        except asyncio.TimeoutError:
-            logger.error(f"Execution timeout after {self._timeout}s")
-            state["error_count"] = state.get("error_count", 0) + 1
-            state["last_error"] = {
-                "error_type": "TimeoutError",
-                "error_message": f"Execution timed out after {self._timeout} seconds",
-            }
-
-            return {
-                "success": False,
-                "error_type": "TimeoutError",
-                "error": f"Execution timed out after {self._timeout} seconds",
-            }
-        except Exception as e:
-            logger.error(f"Execution failed: {e}")
-            state["error_count"] = state.get("error_count", 0) + 1
-            state["last_error"] = {
-                "error_type": type(e).__name__,
-                "error_message": str(e),
-            }
-
-            return {
-                "success": False,
-                "error_type": type(e).__name__,
-                "error": str(e),
-            }
-
-    async def _add_markdown_cell(
-        self,
-        tool_input: Dict[str, Any],
-        state: AgentState,
-        runtime: AgentRuntime,
-    ) -> Dict[str, Any]:
-        """Add markdown cell to notebook"""
-        content = tool_input.get("content", "")
-        if not content:
-            return {
-                "success": False,
-                "error": "No content provided",
-            }
-
-        executor = self._get_executor(runtime)
-
-        if not executor.is_initialized:
-            return {
-                "success": True,
-                "message": "[Mock] Markdown cell would be added",
-                "content": content[:100],
-            }
-
-        try:
-            cell_index = await executor.add_markdown_cell(content)
-
-            return {
-                "success": True,
-                "cell_index": cell_index,
-                "message": "Markdown cell added successfully",
-            }
-        except Exception as e:
-            logger.error(f"Failed to add markdown cell: {e}")
-            return {
-                "success": False,
-                "error": str(e),
-            }
-
-    async def _handle_final_answer(
-        self,
-        tool_input: Dict[str, Any],
-        state: AgentState,
-        runtime: AgentRuntime,
-    ) -> Dict[str, Any]:
-        """Handle final answer tool"""
-        answer = tool_input.get("answer", "")
-        summary = tool_input.get("summary", "")
-
-        # Update state
-        state["final_answer"] = answer
-        state["is_complete"] = True
-
-        return {
-            "success": True,
-            "answer": answer,
-            "summary": summary,
-            "message": "Task completed",
-            "is_complete": True,
-        }
-
-
-def create_jupyter_execution_middleware(
-    timeout: float = 60.0,
-    add_to_notebook: bool = True,
-    enabled: bool = True,
-) -> JupyterExecutionMiddleware:
-    """
-    Factory function to create Jupyter execution middleware.
-    
-    Args:
-        timeout: Execution timeout in seconds
-        add_to_notebook: Add cells to notebook
-        enabled: Whether to enable execution
-        
-    Returns:
-        Configured JupyterExecutionMiddleware instance
-    """
-    return JupyterExecutionMiddleware(
-        timeout=timeout,
-        add_to_notebook=add_to_notebook,
-        enabled=enabled,
-    )

agent_server/langchain/middleware/rag_middleware.py

@@ -1,227 +0,0 @@
-"""
-RAG Middleware
-
-Injects relevant context from the RAG system before model calls.
-Uses the existing RAG manager to retrieve relevant documentation
-and code examples based on the user's request.
-"""
-
-import logging
-from typing import Any, Dict, List, Optional
-
-from agent_server.langchain.state import AgentRuntime, AgentState
-
-logger = logging.getLogger(__name__)
-
-
-class RAGMiddleware:
-    """
-    Middleware that injects RAG context before model calls.
-    
-    This middleware:
-    1. Detects required libraries from the user request
-    2. Queries the RAG system for relevant documentation
-    3. Injects the context into the agent state
-    
-    Uses @before_model hook pattern from LangChain middleware.
-    """
-
-    def __init__(
-        self,
-        rag_manager: Any = None,
-        library_detector: Any = None,
-        max_context_length: int = 4000,
-        enabled: bool = True,
-    ):
-        """
-        Initialize RAG middleware.
-        
-        Args:
-            rag_manager: RAGManager instance for context retrieval
-            library_detector: LibraryDetector for detecting required libraries
-            max_context_length: Maximum context length to inject
-            enabled: Whether RAG injection is enabled
-        """
-        self._rag_manager = rag_manager
-        self._library_detector = library_detector
-        self._max_context_length = max_context_length
-        self._enabled = enabled
-
-    @property
-    def name(self) -> str:
-        return "RAGMiddleware"
-
-    def _get_rag_manager(self):
-        """Lazy load RAG manager if not provided"""
-        if self._rag_manager is None:
-            try:
-                from agent_server.core.rag_manager import get_rag_manager
-                self._rag_manager = get_rag_manager()
-            except ImportError:
-                logger.warning("RAG manager not available")
-                return None
-        return self._rag_manager
-
-    def _get_library_detector(self):
-        """Lazy load library detector if not provided"""
-        if self._library_detector is None:
-            try:
-                from hdsp_agent_core.knowledge.loader import get_library_detector
-                self._library_detector = get_library_detector()
-            except ImportError:
-                logger.warning("Library detector not available")
-                return None
-        return self._library_detector
-
-    def _detect_libraries(
-        self,
-        request: str,
-        imported_libraries: List[str],
-    ) -> List[str]:
-        """Detect required libraries from the request"""
-        detector = self._get_library_detector()
-        if detector is None:
-            return []
-
-        try:
-            from hdsp_agent_core.knowledge.loader import get_knowledge_base
-            knowledge_base = get_knowledge_base()
-            available = knowledge_base.list_available_libraries()
-
-            if not available:
-                return []
-
-            detected = detector.detect(
-                request=request,
-                available_libraries=available,
-                imported_libraries=imported_libraries,
-            )
-
-            return detected
-        except Exception as e:
-            logger.warning(f"Library detection failed: {e}")
-            return []
-
-    async def _get_rag_context(
-        self,
-        query: str,
-        detected_libraries: List[str],
-    ) -> Optional[str]:
-        """Get RAG context for the query"""
-        rag_manager = self._get_rag_manager()
-        if rag_manager is None or not rag_manager.is_ready:
-            return None
-
-        try:
-            context = await rag_manager.get_context_for_query(
-                query=query,
-                detected_libraries=detected_libraries,
-            )
-
-            if context and len(context) > self._max_context_length:
-                context = context[:self._max_context_length] + "\n... (truncated)"
-
-            return context
-        except Exception as e:
-            logger.warning(f"RAG context retrieval failed: {e}")
-            return None
-
-    async def before_model(
-        self,
-        state: AgentState,
-        runtime: AgentRuntime,
-    ) -> Optional[Dict[str, Any]]:
-        """
-        Hook called before each model invocation.
-        
-        Injects RAG context into the state if available.
-        
-        Args:
-            state: Current agent state
-            runtime: Agent runtime context
-            
-        Returns:
-            Updated state fields or None
-        """
-        if not self._enabled:
-            return None
-
-        # Skip if context already injected
-        if state.get("rag_context"):
-            return None
-
-        user_request = state.get("user_request", "")
-        if not user_request:
-            return None
-
-        # Detect libraries
-        notebook_context = state.get("notebook_context", {})
-        imported_libs = notebook_context.get("imported_libraries", [])
-
-        detected_libraries = self._detect_libraries(user_request, imported_libs)
-
-        if detected_libraries:
-            logger.info(f"Detected libraries: {detected_libraries}")
-
-        # Get RAG context
-        rag_context = await self._get_rag_context(user_request, detected_libraries)
-
-        if rag_context:
-            logger.info(f"RAG context injected: {len(rag_context)} chars")
-            return {
-                "rag_context": rag_context,
-                "detected_libraries": detected_libraries,
-            }
-
-        return {"detected_libraries": detected_libraries}
-
-    def format_context_for_prompt(
-        self,
-        rag_context: Optional[str],
-        detected_libraries: List[str],
-    ) -> str:
-        """
-        Format RAG context for inclusion in the prompt.
-        
-        Args:
-            rag_context: Retrieved RAG context
-            detected_libraries: List of detected libraries
-            
-        Returns:
-            Formatted context string
-        """
-        if not rag_context:
-            return ""
-
-        parts = []
-
-        if detected_libraries:
-            parts.append(f"## Detected Libraries: {', '.join(detected_libraries)}")
-
-        parts.append("## Relevant Documentation")
-        parts.append(rag_context)
-
-        return "\n\n".join(parts)
-
-
-def create_rag_middleware(
-    rag_manager: Any = None,
-    max_context_length: int = 4000,
-    enabled: bool = True,
-) -> RAGMiddleware:
-    """
-    Factory function to create RAG middleware.
-    
-    Args:
-        rag_manager: Optional RAGManager instance
-        max_context_length: Maximum context length
-        enabled: Whether to enable RAG
-        
-    Returns:
-        Configured RAGMiddleware instance
-    """
-    return RAGMiddleware(
-        rag_manager=rag_manager,
-        max_context_length=max_context_length,
-        enabled=enabled,
-    )