codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/autofix_middleware.py

@@ -0,0 +1,481 @@
+"""AutoFix Middleware - 코드 자동 수정 + 에러 복구 미들웨어.
+
+Phase 1 (awrap_tool_call): jupyter_cell 실행 전 Ruff 자동 수정 + API 안티패턴 감지
+Phase 2 (before_model): 실행 에러 → ErrorClassifier → 복구 프롬프트 주입
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from collections.abc import Awaitable, Callable
+from typing import Any, Optional, Tuple
+
+from langchain_core.messages import HumanMessage, ToolMessage
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+)
+from langgraph.prebuilt.tool_node import ToolCallRequest
+
+from codebuff_core.core.code_validator import CodeValidator, get_api_pattern_checker
+from codebuff_core.core.error_classifier import (
+    ErrorAnalysis,
+    ReplanDecision,
+    get_error_classifier,
+)
+
+from ..tools.context import NotebookContextManager
+
+logger = logging.getLogger(__name__)
+
+# HITL tools that execute code and can produce runtime errors
+HITL_ERROR_TOOLS = {
+    "jupyter_cell_tool",
+    "jupyter_cell",
+}
+
+# Marker prefix for error recovery messages (used for retry counting)
+ERROR_RECOVERY_MARKER = "[ERROR_RECOVERY]"
+
+# Error patterns in tool output
+ERROR_PATTERNS = [
+    r"(\w+Error):\s*(.+)",
+    r"(\w+Exception):\s*(.+)",
+    r"Traceback \(most recent call last\)",
+]
+
+
+class AutoFixMiddleware(AgentMiddleware):
+    """코드 자동 수정 + 에러 복구 미들웨어.
+
+    Phase 1 (awrap_tool_call): jupyter_cell 실행 전 Ruff 자동 수정 + API 안티패턴 감지
+    Phase 2 (before_model): 실행 에러 → ErrorClassifier → 복구 프롬프트 주입
+    """
+
+    def __init__(
+        self,
+        notebook_context: NotebookContextManager,
+        *,
+        max_retries: int = 2,
+        enable_ruff_fix: bool = True,
+        enable_api_check: bool = True,
+        session_id: str = "",
+    ) -> None:
+        super().__init__()
+        self._notebook_context = notebook_context
+        self._max_retries = max_retries
+        self._enable_ruff_fix = enable_ruff_fix
+        self._enable_api_check = enable_api_check
+        self._session_id = session_id
+        self._classifier = get_error_classifier()
+        self._api_checker = get_api_pattern_checker()
+        self.tools = []
+
+    # ------------------------------------------------------------------
+    # Phase 1: Pre-execution — awrap_tool_call
+    # ------------------------------------------------------------------
+
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Any]],
+    ) -> ToolMessage | Any:
+        """jupyter_cell 도구 호출 전 코드 자동 수정."""
+        tool_name = request.tool_call.get("name", "")
+
+        # jupyter_cell이 아닌 도구는 그대로 통과
+        if tool_name != "jupyter_cell":
+            return await handler(request)
+
+        args = request.tool_call.get("args", {})
+        code = args.get("code", "")
+
+        # 빈 코드는 통과
+        if not code or not code.strip():
+            return await handler(request)
+
+        fixed_code = code
+
+        # Ruff 자동 수정
+        if self._enable_ruff_fix:
+            try:
+                notebook_ctx = self._get_notebook_context_dict()
+                validator = CodeValidator(notebook_context=notebook_ctx)
+                ruff_fixed, issues = validator.check_with_ruff(code, auto_fix=True)
+                if ruff_fixed != code:
+                    fixed_code = ruff_fixed
+                    logger.info(
+                        "AutoFix: Ruff auto-fixed code (%d issues remaining)",
+                        len(issues),
+                    )
+            except Exception as e:
+                logger.debug(
+                    "AutoFix: Ruff check failed [%s], using original code",
+                    type(e).__name__,
+                    exc_info=True,
+                )
+
+        # API 안티패턴 감지 (경고 로깅만)
+        if self._enable_api_check:
+            try:
+                api_issues = self._api_checker.check(fixed_code)
+                for issue in api_issues:
+                    logger.warning("AutoFix API pattern: %s", issue.message)
+            except Exception as e:
+                logger.debug(
+                    "AutoFix: API pattern check failed [%s]",
+                    type(e).__name__,
+                    exc_info=True,
+                )
+
+        # 코드가 수정되었으면 tool_call 오버라이드
+        if fixed_code != code:
+            modified_call = {
+                **request.tool_call,
+                "args": {**args, "code": fixed_code},
+            }
+            request = request.override(tool_call=modified_call)
+
+        return await handler(request)
+
+    # ------------------------------------------------------------------
+    # Phase 2: Post-error — before_model
+    # ------------------------------------------------------------------
+
+    def before_model(
+        self, state: AgentState[Any], runtime: Any
+    ) -> dict[str, Any] | None:
+        """실행 에러 발생 시 복구 프롬프트 주입."""
+        messages = state.get("messages", [])
+        if not messages:
+            return None
+
+        last_msg = messages[-1]
+
+        # HITL 도구 에러인지 확인
+        if not _is_hitl_tool_error(last_msg):
+            return None
+
+        # 연속 복구 횟수 체크
+        recovery_count = _count_error_recoveries(messages)
+        if recovery_count >= self._max_retries:
+            logger.warning(
+                "AutoFix: max retries reached (%d/%d). Passing through.",
+                recovery_count,
+                self._max_retries,
+            )
+            return None
+
+        # 에러 파싱 및 분류
+        content = getattr(last_msg, "content", "") or ""
+        error_type, error_msg, traceback_str = _extract_error_parts(content)
+        if not error_type:
+            return None
+
+        analysis = self._classifier.classify(error_type, error_msg, traceback_str)
+
+        # ErrorClassifier 결과를 structured log에 기록
+        from .observability import _write_structured_event
+
+        _write_structured_event(
+            {
+                "event": "error_classification",
+                "session_id": self._session_id,
+                "error_type": error_type,
+                "decision": analysis.decision.value,
+                "root_cause": analysis.root_cause[:200],
+                "used_llm": analysis.used_llm,
+                "confidence": analysis.confidence,
+            }
+        )
+
+        logger.info(
+            "AutoFix: %s -> %s (attempt %d/%d)",
+            error_type,
+            analysis.decision.value,
+            recovery_count + 1,
+            self._max_retries,
+        )
+
+        # 복구 프롬프트 생성
+        prompt = _build_replan_prompt(analysis, error_type, error_msg)
+
+        return {"messages": [HumanMessage(content=prompt)]}
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _get_notebook_context_dict(self) -> dict:
+        """NotebookContextManager → CodeValidator용 dict 변환."""
+        ctx = self._notebook_context.current_notebook
+        if not ctx:
+            return {}
+
+        defined_vars: list[str] = []
+        imported_libs: list[str] = []
+
+        for cell in ctx.get_code_cells():
+            source = cell.source or ""
+            for line in source.split("\n"):
+                stripped = line.strip()
+                # import 파싱
+                if stripped.startswith("import "):
+                    parts = stripped.replace("import ", "").split(" as ")
+                    name = parts[-1].strip().split(",")[0].strip()
+                    imported_libs.append(name)
+                elif stripped.startswith("from ") and " import " in stripped:
+                    after_import = stripped.split(" import ")[1]
+                    for name in after_import.split(","):
+                        name = name.strip().split(" as ")[-1].strip()
+                        if name:
+                            imported_libs.append(name)
+                # 변수 정의 파싱 (간단한 할당)
+                elif "=" in stripped and not stripped.startswith("#"):
+                    lhs = stripped.split("=")[0].strip()
+                    # 함수 호출이나 비교 연산자 제외
+                    if (
+                        lhs.isidentifier()
+                        and not stripped.startswith("if ")
+                        and not stripped.startswith("elif ")
+                        and "==" not in stripped
+                        and "!=" not in stripped
+                        and "<=" not in stripped
+                        and ">=" not in stripped
+                    ):
+                        defined_vars.append(lhs)
+
+        return {
+            "definedVariables": defined_vars,
+            "importedLibraries": imported_libs,
+        }
+
+
+# ======================================================================
+# Module-level helper functions (ported from error_recovery.py)
+# ======================================================================
+
+
+def _is_hitl_tool_error(msg: Any) -> bool:
+    """Check if a message is a ToolMessage from a HITL tool containing an error."""
+    if getattr(msg, "type", "") != "tool":
+        return False
+
+    tool_name = getattr(msg, "name", "") or ""
+    if tool_name not in HITL_ERROR_TOOLS:
+        return False
+
+    content = getattr(msg, "content", "") or ""
+    if not isinstance(content, str):
+        return False
+
+    # Check structured execution_result first
+    try:
+        data = json.loads(content)
+        if isinstance(data, dict):
+            exec_result = data.get("execution_result")
+            if isinstance(exec_result, dict):
+                return exec_result.get("success") is False or bool(
+                    exec_result.get("error")
+                )
+    except (json.JSONDecodeError, TypeError, ValueError):
+        pass
+
+    return _has_error_indicators(content)
+
+
+def _has_error_indicators(content: str) -> bool:
+    """Check if content contains Python error indicators."""
+    for pattern in ERROR_PATTERNS:
+        if re.search(pattern, content):
+            return True
+    return False
+
+
+def _count_error_recoveries(messages: list) -> int:
+    """Count consecutive error recovery attempts since last successful HITL result."""
+    count = 0
+    for msg in reversed(messages):
+        content = getattr(msg, "content", "") or ""
+        msg_type = getattr(msg, "type", "")
+
+        if isinstance(content, str) and ERROR_RECOVERY_MARKER in content:
+            count += 1
+        elif msg_type == "tool":
+            tool_name = getattr(msg, "name", "") or ""
+            if tool_name in HITL_ERROR_TOOLS:
+                if isinstance(content, str) and _has_error_indicators(content):
+                    continue
+                else:
+                    break
+            continue
+        elif msg_type in ("ai", "human"):
+            if msg_type == "human" and isinstance(content, str):
+                if ERROR_RECOVERY_MARKER not in content:
+                    break
+            continue
+        else:
+            break
+
+    return count
+
+
+def _extract_error_parts(content: str) -> Tuple[Optional[str], str, str]:
+    """Extract error type, message, and traceback from tool output."""
+    if not content or not isinstance(content, str):
+        return None, "", ""
+
+    # Strategy 1: Structured JSON
+    try:
+        data = json.loads(content)
+        if isinstance(data, dict):
+            exec_result = data.get("execution_result")
+            if isinstance(exec_result, dict):
+                error_type = exec_result.get("error_type")
+                if error_type:
+                    error_msg = exec_result.get("error", "")
+                    prefix = f"{error_type}: "
+                    if error_msg.startswith(prefix):
+                        error_msg = error_msg[len(prefix) :]
+                    traceback_str = exec_result.get("traceback", "") or ""
+                    return error_type, error_msg, traceback_str
+    except (json.JSONDecodeError, TypeError, ValueError):
+        pass
+
+    # Strategy 2: Regex-based extraction
+    lines = content.strip().split("\n")
+
+    error_type = None
+    error_msg = ""
+
+    for line in reversed(lines):
+        line = line.strip()
+        line = re.sub(r"\x1b\[[0-9;]*m", "", line)
+
+        match = re.search(r"(\w+(?:Error|Exception)):\s*(.*)", line)
+        if match:
+            error_type = match.group(1)
+            error_msg = match.group(2)
+            break
+
+        match = re.match(r"^(\w+(?:Error|Exception))\s*$", line)
+        if match:
+            error_type = match.group(1)
+            break
+
+    traceback_str = ""
+    traceback_match = re.search(
+        r"(Traceback \(most recent call last\)[\s\S]*)", content
+    )
+    if traceback_match:
+        traceback_str = traceback_match.group(1)
+
+    return error_type, error_msg, traceback_str
+
+
+def _get_error_guidance(error_type: str, error_msg: str) -> str:
+    """Get error-specific guidance for the LLM."""
+    guidance_map = {
+        "KeyError": (
+            "Check the actual column/key names by running "
+            "`df.columns.tolist()` or `print(list(dict.keys()))` first."
+        ),
+        "FileNotFoundError": (
+            "Verify the file path by running "
+            "`import os; print(os.listdir('.'))` or checking subdirectories."
+        ),
+        "NameError": (
+            "Ensure all imports and variable definitions are in the same cell."
+        ),
+        "ModuleNotFoundError": (
+            "Install the missing package using `!pip install <package_name>` before importing."
+        ),
+        "ImportError": (
+            "Check if the package is installed and the import path is correct."
+        ),
+        "TypeError": ("Check function arguments and data types."),
+        "ValueError": ("Check input values, data types, ranges, and formats."),
+        "IndexError": ("Check array/list bounds before accessing by index."),
+        "AttributeError": (
+            "Check if the object has the expected type. Use `type(obj)` and `dir(obj)`."
+        ),
+        "SyntaxError": (
+            "Check for missing colons, brackets, parentheses, or indentation errors."
+        ),
+    }
+    return guidance_map.get(
+        error_type,
+        f"Analyze the {error_type} carefully and fix the root cause before retrying.",
+    )
+
+
+def _build_replan_prompt(
+    analysis: ErrorAnalysis,
+    error_type: str,
+    error_msg: str,
+) -> str:
+    """Build re-plan prompt based on ErrorClassifier analysis."""
+    decision = analysis.decision
+    guidance = _get_error_guidance(error_type, error_msg)
+    error_preview = error_msg[:200] if error_msg else "unknown error"
+
+    if decision == ReplanDecision.REFINE:
+        return (
+            f"{ERROR_RECOVERY_MARKER} 이전 코드 실행이 실패했습니다: "
+            f"{error_type}: {error_preview}\n\n"
+            f"근본 원인: {analysis.root_cause}\n\n"
+            f"지침:\n"
+            f"1. 에러를 주의 깊게 분석하세요\n"
+            f"2. {guidance}\n"
+            f"3. write_todos를 호출하여 수정 방안으로 계획을 업데이트하세요\n"
+            f"4. 수정된 코드로 jupyter_cell을 호출하세요\n"
+            f"5. 에러가 해결될 때까지 final_answer_tool을 호출하지 마세요"
+        )
+
+    if decision == ReplanDecision.INSERT_STEPS:
+        pkg = analysis.missing_package or "unknown"
+        return (
+            f"{ERROR_RECOVERY_MARKER} 패키지 '{pkg}'가 없어서 이전 코드 실행이 "
+            f"실패했습니다.\n\n"
+            f"지침:\n"
+            f"1. jupyter_cell을 호출하여 설치하세요: !pip install {pkg}\n"
+            f"2. 그런 다음 원래 코드를 다시 실행하세요\n"
+            f"3. write_todos를 호출하여 패키지 설치 단계를 반영하세요\n"
+            f"4. 에러가 해결될 때까지 final_answer_tool을 호출하지 마세요"
+        )
+
+    if decision == ReplanDecision.REPLACE_STEP:
+        return (
+            f"{ERROR_RECOVERY_MARKER} 이전 접근법이 완전히 실패했습니다: "
+            f"{error_type}: {error_preview}\n\n"
+            f"근본 원인: {analysis.root_cause}\n\n"
+            f"지침:\n"
+            f"1. 현재 접근법은 불가능합니다 — 완전히 다른 방법을 시도하세요\n"
+            f"2. {guidance}\n"
+            f"3. write_todos를 호출하여 현재 단계를 새로운 접근법으로 교체하세요\n"
+            f"4. 새로운 접근법으로 jupyter_cell을 호출하세요\n"
+            f"5. 에러가 해결될 때까지 final_answer_tool을 호출하지 마세요"
+        )
+
+    if decision == ReplanDecision.REPLAN_REMAINING:
+        return (
+            f"{ERROR_RECOVERY_MARKER} 치명적 에러가 발생했습니다: "
+            f"{error_type}: {error_preview}\n\n"
+            f"근본 원인: {analysis.root_cause}\n\n"
+            f"지침:\n"
+            f"1. 근본적인 문제로 남은 모든 단계를 재계획해야 합니다\n"
+            f"2. {guidance}\n"
+            f"3. write_todos를 호출하여 남은 계획을 전면 재설계하세요\n"
+            f"4. 새로운 계획으로 jupyter_cell 실행을 시작하세요\n"
+            f"5. 에러가 해결될 때까지 final_answer_tool을 호출하지 마세요"
+        )
+
+    return (
+        f"{ERROR_RECOVERY_MARKER} 에러가 발생했습니다: "
+        f"{error_type}: {error_preview}\n\n"
+        f"에러를 수정하고 계속 진행하세요. "
+        f"해결될 때까지 final_answer_tool을 호출하지 마세요."
+    )