codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/continuation_middleware.py

@@ -0,0 +1,548 @@
+"""ContinuationControl + LimitToolCalls + TodoActive 미들웨어.
+
+ContinuationControlMiddleware: 도구 완료 후 LLM에게 다음 행동 안내.
+LimitToolCallsMiddleware: LLM 응답의 tool_calls를 1개로 제한 + additional_kwargs 정리.
+TodoActiveMiddleware: wrap_tool_call로 todo_active state flag 관리 (Command 반환).
+"""
+
+from __future__ import annotations
+
+import datetime
+import json
+import logging
+from collections.abc import Awaitable, Callable
+from typing import Annotated, Any, NotRequired
+
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
+from langgraph.types import Command
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ModelCallResult,
+    ModelRequest,
+    ModelResponse,
+)
+
+from ..log_paths import get_debug_log_path
+
+logger = logging.getLogger(__name__)
+
+# Direct file-based debug logging (bypasses Python logging config)
+
+_DEBUG_LOG = get_debug_log_path()
+
+
+def _dbg(msg: str) -> None:
+    """Write debug message directly to file."""
+    try:
+        ts = datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
+        with open(_DEBUG_LOG, "a") as f:
+            f.write(f"[{ts}] CONTINUATION: {msg}\n")
+    except OSError:
+        pass
+
+
+# write_todos 후에는 "다시 write_todos 호출 금지" 별도 guidance 사용
+_WRITE_TODOS_TOOLS = {"write_todos"}
+
+
+def _get_tool_name(msg: Any) -> str:
+    """ToolMessage에서 도구 이름 추출."""
+    if hasattr(msg, "name"):
+        return msg.name or ""
+    if hasattr(msg, "additional_kwargs"):
+        return msg.additional_kwargs.get("name", "")
+    return ""
+
+
+def _extract_ai_message(response: Any) -> AIMessage | None:
+    """ModelResponse/ModelCallResult에서 AIMessage 추출."""
+    if isinstance(response, AIMessage):
+        return response
+    if hasattr(response, "message") and isinstance(response.message, AIMessage):
+        return response.message
+    if hasattr(response, "messages"):
+        for msg in reversed(response.messages):
+            if isinstance(msg, AIMessage):
+                return msg
+    return None
+
+
+def _get_todos_from_state(request: ModelRequest) -> list | None:
+    """LangGraph state에서 todos 직접 읽기. 없으면 None."""
+    state = getattr(request, "state", None)
+    if state and isinstance(state, dict):
+        todos = state.get("todos")
+        if isinstance(todos, list):
+            return todos
+    return None
+
+
+def _get_todos_from_messages(messages: list) -> list | None:
+    """메시지에서 write_todos 결과를 찾아 todos 추출. 없으면 None."""
+    for msg in reversed(messages):
+        if not hasattr(msg, "type") or msg.type != "tool":
+            continue
+        if _get_tool_name(msg) != "write_todos":
+            continue
+        content = getattr(msg, "content", "")
+        if not isinstance(content, str):
+            continue
+        try:
+            data = json.loads(content)
+            todos = data.get("todos", [])
+            if todos:
+                return todos
+        except (json.JSONDecodeError, AttributeError):
+            continue
+    return None
+
+
+def _classify_todos(todos: list | None) -> str:
+    """todos 상태 분류. 'none' | 'pending' | 'all_completed'."""
+    if not todos:
+        return "none"
+    pending = [
+        t
+        for t in todos
+        if isinstance(t, dict) and t.get("status") in ("pending", "in_progress")
+    ]
+    # dict가 아닌 경우 (Todo 객체): attribute 접근
+    if not pending:
+        pending = [
+            t
+            for t in todos
+            if hasattr(t, "status") and t.status in ("pending", "in_progress")
+        ]
+    return "pending" if pending else "all_completed"
+
+
+def _has_summary_in_messages(messages: list) -> bool:
+    """메시지에서 final_summary 도구 호출 또는 요약 텍스트 존재 여부 확인."""
+    for msg in reversed(messages[-10:]):
+        msg_type = getattr(msg, "type", "")
+        name = _get_tool_name(msg) if msg_type == "tool" else ""
+
+        # final_summary 도구 결과
+        if name in ("final_summary_tool", "final_summary"):
+            return True
+
+        # AI 메시지의 summary JSON 패턴
+        content = getattr(msg, "content", "")
+        if isinstance(content, str):
+            if ('"summary"' in content and '"next_items"' in content) or (
+                "'summary'" in content and "'next_items'" in content
+            ):
+                return True
+            # 한국어 요약 패턴
+            if any(
+                kw in content
+                for kw in ("모든 작업이 완료", "작업 요약", "다음 단계 제안")
+            ):
+                return True
+    return False
+
+
+# ---------------------------------------------------------------------------
+# TodoActiveState — state schema for todo_active flag
+# ---------------------------------------------------------------------------
+
+
+def _take_last(existing: str, new: str) -> str:
+    """Reducer: 마지막 값만 유지. 병렬 도구 호출 시 충돌 방지."""
+    return new
+
+
+def _take_last_bool(existing: bool, new: bool) -> bool:
+    """Reducer: 마지막 bool 값만 유지."""
+    return new
+
+
+class TodoActiveState(AgentState):
+    """todo_active flag를 포함하는 state schema.
+
+    - todo_active=True: write_todos가 호출된 후 (작업 진행 중)
+    - todo_active=False: final_summary_tool 호출 후 또는 초기 상태
+    - last_tool_name: 가장 최근 실행된 도구 이름 (messages[-1] 의존 제거)
+    """
+
+    todo_active: NotRequired[Annotated[bool, _take_last_bool]]
+    last_tool_name: NotRequired[Annotated[str, _take_last]]
+
+
+# ---------------------------------------------------------------------------
+# TodoActiveMiddleware — wrap_tool_call로 state flag 관리
+# ---------------------------------------------------------------------------
+
+
+class TodoActiveMiddleware(AgentMiddleware):
+    """wrap_tool_call로 todo_active, last_tool_name state flag를 관리.
+
+    - write_todos → todo_active=True, last_tool_name="write_todos"
+    - final_summary_tool → todo_active=False, last_tool_name="final_summary_tool"
+    - task → last_tool_name="task"
+    - 기타 도구 → last_tool_name=<tool_name>
+
+    Command(update={...})로 LangGraph state를 확정적으로 업데이트.
+    messages[-1] 타입에 의존하지 않아 system 메시지 끼어들기 문제 해결.
+    """
+
+    state_schema = TodoActiveState
+
+    async def awrap_tool_call(self, request, handler):
+        """도구 실행 후 state flag 업데이트."""
+        result = await handler(request)
+        tool_name = request.tool_call.get("name", "")
+
+        if tool_name == "write_todos":
+            _dbg("TodoActive: write_todos → todo_active=True")
+            return self._wrap_with_state(
+                result, todo_active=True, last_tool_name=tool_name
+            )
+        elif tool_name in ("final_summary_tool", "final_summary"):
+            _dbg(f"TodoActive: {tool_name} → todo_active=False")
+            return self._wrap_with_state(
+                result, todo_active=False, last_tool_name=tool_name
+            )
+        elif tool_name:
+            extra = ""
+            if tool_name == "task":
+                args = request.tool_call.get("args", {})
+                if isinstance(args, str):
+                    try:
+                        args = json.loads(args)
+                    except Exception:
+                        args = {}
+                agent_type = args.get("subagent_type", args.get("agent_name", ""))
+                if agent_type:
+                    extra = f":{agent_type}"
+            _dbg(f"TodoActive: {tool_name}{extra} → last_tool_name update")
+            return self._wrap_with_state(result, last_tool_name=f"{tool_name}{extra}")
+
+        return result
+
+    def _wrap_with_state(
+        self,
+        result: Any,
+        *,
+        todo_active: bool | None = None,
+        last_tool_name: str | None = None,
+    ) -> Command:
+        """도구 결과를 Command로 감싸서 state 업데이트."""
+        update: dict[str, Any] = {}
+        if todo_active is not None:
+            update["todo_active"] = todo_active
+        if last_tool_name is not None:
+            update["last_tool_name"] = last_tool_name
+
+        try:
+            if isinstance(result, Command):
+                # 기존 Command에 merge
+                existing = (
+                    result.update if hasattr(result, "update") and result.update else {}
+                )
+                merged = {**existing, **update}
+                return Command(update=merged)
+            elif isinstance(result, ToolMessage):
+                update["messages"] = [result]
+                return Command(update=update)
+            else:
+                # Unknown type → ToolMessage로 변환
+                content = str(result) if result else ""
+                tool_msg = ToolMessage(content=content, tool_call_id="")
+                update["messages"] = [tool_msg]
+                return Command(update=update)
+        except Exception as e:
+            _dbg(f"TodoActive: wrap failed: {e}")
+            return result
+
+
+# ---------------------------------------------------------------------------
+# ContinuationControlMiddleware — state flag 기반 continuation 제어
+# ---------------------------------------------------------------------------
+
+
+class ContinuationControlMiddleware(AgentMiddleware):
+    """state flag 기반으로 도구 완료 후 LLM에게 다음 행동을 안내.
+
+    TodoActiveMiddleware가 설정한 state flag를 읽어서 판단:
+    - todo_active=False + final_summary 완료 → 즉시 중단
+    - last_tool_name="task" → 코드 실행 유도 guidance 주입
+    - last_tool_name="write_todos" → 다음 작업 진행 유도
+    - pending todos → 계속 진행 유도
+    - 모든 todos 완료 → final_summary 호출 유도
+    """
+
+    def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> Awaitable[ModelCallResult]:
+        return self._impl(request, handler)
+
+    async def _impl(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        messages = request.messages
+        state = getattr(request, "state", {}) or {}
+        todo_active = state.get("todo_active", False)
+        last_tool = state.get("last_tool_name", "")
+
+        _dbg(
+            f"_impl called, msg_count={len(messages) if messages else 0}, "
+            f"todo_active={todo_active}, last_tool={last_tool}"
+        )
+
+        # ── Guard 1: final_summary 완료 후 즉시 중단 ──
+        # todo_active=False이고 메시지에 summary가 있으면 agent 중단
+        if not todo_active and last_tool in ("final_summary_tool", "final_summary"):
+            _dbg("final_summary tool + todo_active=False → stopping agent")
+            return AIMessage(content="", tool_calls=[])
+
+        # todo_active=False이면서 summary가 messages에 있으면 중단
+        if not todo_active and _has_summary_in_messages(messages):
+            _dbg("todo_active=False + summary in messages → stopping agent")
+            return AIMessage(content="", tool_calls=[])
+
+        # ── Guard 2: task 완료 후 결과 처리 유도 ──
+        if last_tool.startswith("task"):
+            subagent = last_tool.split(":", 1)[1] if ":" in last_tool else ""
+            if subagent in ("bq_agent", "athena_query", "airflow_expert"):
+                guidance = (
+                    "[SYSTEM] Sub-agent returned analysis/SQL. "
+                    "You MUST display it with jupyter_markdown() BEFORE calling write_todos. "
+                    "Do NOT skip display. Do NOT call jupyter_cell for analysis results."
+                )
+            else:
+                guidance = (
+                    "[SYSTEM] Sub-agent returned code. "
+                    "You MUST execute it with jupyter_cell() BEFORE calling write_todos. "
+                    "Do NOT skip execution. Do NOT mark the todo as completed without running the code."
+                )
+            new_messages = list(messages) + [HumanMessage(content=guidance)]
+            _dbg(f"injected post-task guidance for subagent={subagent} (state-based)")
+            request = request.override(messages=new_messages)
+            response = await handler(request)
+            return self._strip_write_todos_if_summary(response)
+
+        # ── Guard 3: todo_active=False → 간단한 작업, guidance 불필요 ──
+        if not todo_active:
+            _dbg(f"todo_active=False, last_tool={last_tool} → skip guidance")
+            response = await handler(request)
+            return self._strip_write_todos_if_summary(response)
+
+        # ── todo_active=True: continuation guidance 주입 ──
+        if last_tool:
+            is_write_todos = last_tool in _WRITE_TODOS_TOOLS
+            modified = self._inject_continuation(
+                request,
+                last_tool,
+                after_write_todos=is_write_todos,
+            )
+            if modified is None:
+                _dbg("summary exists → stopping agent")
+                return AIMessage(content="", tool_calls=[])
+            request = modified
+
+        response = await handler(request)
+        return self._strip_write_todos_if_summary(response)
+
+    def _get_last_tool_result(self, messages: list) -> dict | None:
+        """마지막 ToolMessage에서 결과를 dict로 추출. 없으면 None."""
+        for msg in reversed(messages):
+            if not isinstance(msg, ToolMessage):
+                continue
+            content = getattr(msg, "content", "")
+            if not isinstance(content, str):
+                continue
+            try:
+                return json.loads(content)
+            except (json.JSONDecodeError, TypeError):
+                # Plain text error message → wrap in dict
+                if content.strip():
+                    return {"error": content, "success": False}
+            break
+        return None
+
+    def _inject_continuation(
+        self,
+        request: ModelRequest,
+        tool_name: str,
+        *,
+        after_write_todos: bool = False,
+    ) -> ModelRequest | None:
+        """continuation 메시지 주입.
+
+        Returns:
+            수정된 request, 또는 None (에이전트 중단 신호)
+        """
+        messages = request.messages
+
+        # ── Error recovery: 마지막 도구 실패 시 복구 guidance 주입 ──
+        last_tool_result = self._get_last_tool_result(messages)
+        if last_tool_result and not last_tool_result.get("success", True):
+            error_msg = last_tool_result.get("error", "")
+            if "command not found" in error_msg or "not found" in error_msg:
+                guidance = (
+                    f"[SYSTEM] Tool '{tool_name}' failed: {error_msg[:100]}. "
+                    "Replan: if 'pip' not found, use jupyter_cell(code='!pip install ...') instead. "
+                    "Do NOT retry the same command."
+                )
+                new_messages = list(messages) + [HumanMessage(content=guidance)]
+                _dbg(f"injected error recovery guidance for {tool_name}")
+                return request.override(messages=new_messages)
+
+        # todos 가져오기: state 우선, 없으면 메시지에서 추출
+        todos = _get_todos_from_state(request) or _get_todos_from_messages(messages)
+        todos_status = _classify_todos(todos)
+        _dbg(
+            f"_inject tool={tool_name} todos_status={todos_status} after_wt={after_write_todos} todos={todos!r:.200s}"
+        )
+
+        # todos 없음 → 간단한 작업, guidance 불필요
+        if todos_status == "none":
+            _dbg("no todos → skip guidance")
+            return request
+
+        # pending todos 있음 → 계속 진행
+        if todos_status == "pending":
+            pending_preview = self._build_pending_preview(todos)
+
+            if after_write_todos:
+                guidance = (
+                    f"[SYSTEM] Todo list updated.{pending_preview} "
+                    "Now continue working on the next pending todo item. "
+                    "Use 'task' for code generation. "
+                    "Package installs → jupyter_cell(code='!pip install ...'). "
+                    "Exploration only (file check, head, pip list) → execute. "
+                    "Do NOT call write_todos again until the next task is done. "
+                    "Do NOT repeat tools that already succeeded."
+                )
+            else:
+                guidance = (
+                    f"[SYSTEM] Tool '{tool_name}' completed successfully.{pending_preview} "
+                    "Update the current todo status with write_todos, then continue "
+                    "with the next pending todo item. "
+                    "Use 'task' for code generation. "
+                    "Package installs → jupyter_cell(code='!pip install ...'). "
+                    "Exploration only (file check, head, pip list) → execute. "
+                    "Do NOT repeat tools that already succeeded."
+                )
+        else:
+            # 모든 todos 완료 → summary 이미 있으면 중단, 없으면 호출 유도
+            if _has_summary_in_messages(messages):
+                _dbg("all_completed + summary exists → stopping agent")
+                return None
+
+            guidance = (
+                "[SYSTEM] 모든 작업이 완료되었습니다. "
+                "반드시 final_summary_tool을 호출하여 작업 요약과 다음 단계를 제시하세요. "
+                "final_summary_tool(summary='완료된 작업 요약', "
+                "next_items=[{'subject': '제목', 'description': '설명'}, ...]) "
+                "(next_items 3개 이상 필수). "
+                "텍스트로 JSON을 출력하지 말고, 반드시 도구 호출로 실행하세요."
+            )
+
+        new_messages = list(messages) + [HumanMessage(content=guidance)]
+        _dbg(f"injected guidance for tool={tool_name} status={todos_status}")
+        _dbg(f"guidance: {guidance[:200]}")
+        return request.override(messages=new_messages)
+
+    def _build_pending_preview(self, todos: list | None) -> str:
+        """pending 항목 미리보기 문자열 생성."""
+        if not todos:
+            return ""
+        pending_items = [
+            t
+            for t in todos
+            if (isinstance(t, dict) and t.get("status") in ("pending", "in_progress"))
+            or (hasattr(t, "status") and t.status in ("pending", "in_progress"))
+        ]
+        if not pending_items:
+            return ""
+        previews = []
+        for t in pending_items[:3]:
+            c = (
+                t.get("content", "")
+                if isinstance(t, dict)
+                else getattr(t, "content", "")
+            )
+            previews.append(c[:40])
+        return f" Pending: {', '.join(previews)}."
+
+    def _strip_write_todos_if_summary(self, response: Any) -> Any:
+        """summary JSON이 있는 응답에서 write_todos 호출 제거."""
+        ai_msg = _extract_ai_message(response)
+        if not ai_msg or not hasattr(ai_msg, "tool_calls") or not ai_msg.tool_calls:
+            return response
+
+        content = getattr(ai_msg, "content", "")
+        has_summary = isinstance(content, str) and (
+            ('"summary"' in content and '"next_items"' in content)
+            or ("'summary'" in content and "'next_items'" in content)
+        )
+
+        if not has_summary:
+            return response
+
+        non_todo = [tc for tc in ai_msg.tool_calls if tc.get("name") != "write_todos"]
+        if len(non_todo) < len(ai_msg.tool_calls):
+            ai_msg.tool_calls = non_todo
+            logger.debug("CONTINUATION: stripped write_todos from summary response")
+
+        return response
+
+
+class LimitToolCallsMiddleware(AgentMiddleware):
+    """LLM 응답의 tool_calls를 1개로 제한 + additional_kwargs 정리.
+
+    HITL decision 수 불일치 방지:
+    - LLM이 2개 이상 tool_calls를 반환하면 첫 번째만 유지
+    - additional_kwargs['tool_calls'] 중복 제거
+    - tool_calls가 있으면 content(thinking text) 비우기
+    """
+
+    def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> Awaitable[ModelCallResult]:
+        return self._impl(request, handler)
+
+    async def _impl(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        response = await handler(request)
+
+        ai_msg = _extract_ai_message(response)
+        if not ai_msg:
+            return response
+
+        tool_calls = getattr(ai_msg, "tool_calls", None)
+        if not tool_calls:
+            return response
+
+        # tool_calls가 2개 이상이면 첫 번째만 유지
+        if len(tool_calls) > 1:
+            logger.debug(
+                "LIMIT_TOOL_CALLS: truncating %d -> 1 (kept: %s)",
+                len(tool_calls),
+                tool_calls[0].get("name", "unknown"),
+            )
+            ai_msg.tool_calls = [tool_calls[0]]
+
+        # additional_kwargs['tool_calls'] 중복 제거
+        additional_kwargs = getattr(ai_msg, "additional_kwargs", {})
+        if additional_kwargs.get("tool_calls"):
+            del additional_kwargs["tool_calls"]
+
+        # tool_calls 있으면 content 비우기 (thinking text 제거)
+        if ai_msg.tool_calls and ai_msg.content:
+            ai_msg.content = ""
+
+        return response