codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/observability.py

@@ -0,0 +1,334 @@
+"""ObservabilityMiddleware — 구조화된 JSONL 이벤트 로깅.
+
+모든 model/tool call을 래핑하여 타이밍 + 메타데이터를
+extensions/logs/agent_structured.jsonl에 기록.
+"""
+
+from __future__ import annotations
+
+import datetime
+import json
+import os
+import time
+from collections.abc import Awaitable, Callable
+from contextvars import ContextVar
+from typing import Any, Optional
+
+from codebuff_core.core.exceptions import AgentError
+
+from langchain_core.messages import AIMessage
+from langgraph.errors import GraphBubbleUp
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    ModelCallResult,
+    ModelRequest,
+    ModelResponse,
+)
+
+from ..analytics_ingest import enqueue_agent_call
+from ..team_access import resolve_team
+
+# 로그 디렉토리: log_paths 공통 모듈 사용
+from ..log_paths import get_structured_log_path
+
+_STRUCTURED_LOG = get_structured_log_path()
+
+
+def _write_structured_event(event: dict) -> None:
+    """JSON 라인을 agent_structured.jsonl에 기록. ts 자동 추가.
+
+    파일 기록 뒤 RDS 로도 한 벌 보낸다 (이슈 #30). 이 한 곳이 model_call /
+    tool_call / error_classification 세 종류를 모두 덮는다.
+
+    `agent_backend` 는 codex 경로(`codex_bridge._record`)가 이미 박아 보내므로
+    여기서는 비어 있을 때만 `client` 로 채운다. 파일 기록 **전에** 찍어야 로컬
+    jsonl 과 RDS 행의 값이 같다 (이슈 #92).
+    """
+    try:
+        event["ts"] = datetime.datetime.now().isoformat(timespec="seconds")
+        event.setdefault("agent_backend", "client")
+        with open(_STRUCTURED_LOG, "a") as f:
+            f.write(json.dumps(event, ensure_ascii=False) + "\n")
+    except OSError:
+        pass
+    enqueue_agent_call(event)
+
+
+def default_user_identity() -> dict[str, str]:
+    """구조화 이벤트에 싣는 사용자 식별 — 미들웨어·턴 계측·codex 가 같은 값을 쓴다."""
+    return {
+        "team": resolve_team(),
+        "project": os.environ.get("HDSP_PRJ_ID", ""),
+        "user_id": os.environ.get("JUPYTERHUB_USER", "")
+        or os.environ.get("HDSP_BDP_ID", ""),
+    }
+
+
+def _add(total: Optional[int], value: Optional[int]) -> Optional[int]:
+    """토큰 합. 한 번도 안 잰 것은 None 으로 남긴다 — 0 은 "잰 값이 0" 이다."""
+    if value is None:
+        return total
+    return (total or 0) + value
+
+
+def _usage_of(ai_msg: Optional[AIMessage]) -> dict[str, Any]:
+    """AIMessage 의 토큰·모델명·종료 사유. 제공자가 안 주면 None (0 이 아니다)."""
+    out: dict[str, Any] = {
+        "input_tokens": None,
+        "output_tokens": None,
+        "cached_tokens": None,
+        "model_name": None,
+        "finish_reason": None,
+    }
+    if ai_msg is None:
+        return out
+    usage = getattr(ai_msg, "usage_metadata", None) or {}
+    if usage:
+        out["input_tokens"] = usage.get("input_tokens")
+        out["output_tokens"] = usage.get("output_tokens")
+        out["cached_tokens"] = (usage.get("input_token_details") or {}).get(
+            "cache_read"
+        )
+    meta = getattr(ai_msg, "response_metadata", None) or {}
+    out["model_name"] = meta.get("model_name") or meta.get("model") or None
+    out["finish_reason"] = meta.get("finish_reason") or meta.get("stop_reason") or None
+    return out
+
+
+class TurnMeter:
+    """프롬프트 1건(= 턴)의 계측. HITL 대기를 포함해 턴이 끝날 때 `turn` 행 1개를 낸다.
+
+    그래프는 캐시돼 여러 thread 가 공유하므로 미들웨어 생성자 값으로는 턴을 못 가른다.
+    요청마다 만들어 ContextVar 로 묶고(`bind_turn`), HITL 로 멈추면 runner 에 보관했다가
+    재개 요청에서 다시 묶는다 (이슈 #98).
+    """
+
+    def __init__(
+        self,
+        session_id: str,
+        ux_session_id: Optional[str] = None,
+        user_identity: Optional[dict[str, str]] = None,
+    ):
+        self.session_id = session_id
+        self.ux_session_id = ux_session_id
+        self.user_identity = (
+            user_identity if user_identity is not None else default_user_identity()
+        )
+        self.model_calls = 0
+        self.tool_calls = 0
+        self.input_tokens: Optional[int] = None
+        self.output_tokens: Optional[int] = None
+        self.cached_tokens: Optional[int] = None
+        self.model_name: Optional[str] = None
+        self.finish_reason: Optional[str] = None
+        self._t0 = time.monotonic()
+        self._emitted = False
+
+    def add_model_call(self, usage: dict[str, Any], tool_calls: int) -> None:
+        self.model_calls += 1
+        self.input_tokens = _add(self.input_tokens, usage.get("input_tokens"))
+        self.output_tokens = _add(self.output_tokens, usage.get("output_tokens"))
+        self.cached_tokens = _add(self.cached_tokens, usage.get("cached_tokens"))
+        self.model_name = usage.get("model_name") or self.model_name
+        self.finish_reason = usage.get("finish_reason") or self.finish_reason
+
+    def emit(self, status: str) -> None:
+        """turn 행을 낸다. 종료 지점이 여럿이라 두 번째부터는 무시한다."""
+        if self._emitted:
+            return
+        self._emitted = True
+        _write_structured_event(
+            {
+                "event": "turn",
+                "session_id": self.session_id,
+                "ux_session_id": self.ux_session_id,
+                **self.user_identity,
+                "call_num": self.model_calls,
+                "output_tool_calls": self.tool_calls,
+                "elapsed_ms": round((time.monotonic() - self._t0) * 1000, 1),
+                "status": status,
+                "model_name": self.model_name,
+                "input_tokens": self.input_tokens,
+                "output_tokens": self.output_tokens,
+                "cached_tokens": self.cached_tokens,
+                "finish_reason": self.finish_reason,
+            }
+        )
+
+
+_TURN: ContextVar[Optional[TurnMeter]] = ContextVar("codebuff_turn", default=None)
+
+
+def bind_turn(meter: Optional[TurnMeter]) -> None:
+    """이 컨텍스트(와 여기서 만드는 태스크)의 턴을 정한다.
+
+    `astream_events` 는 첫 iteration 에 태스크를 만들며 컨텍스트를 복사하므로 **그 전에**
+    불러야 한다.
+    """
+    _TURN.set(meter)
+
+
+def emit_parked_turn(runner: Any, status: str) -> None:
+    """HITL 대기로 runner 에 보관된 턴을 닫는다 — 취소·리셋이 runner 를 버릴 때."""
+    meter = getattr(runner, "turn_meter", None)
+    if meter is not None:
+        meter.emit(status)
+
+
+def _extract_ai_message(response: Any) -> AIMessage | None:
+    """ModelResponse/ModelCallResult에서 AIMessage 추출."""
+    if isinstance(response, AIMessage):
+        return response
+    if isinstance(response, ModelResponse):
+        for msg in reversed(response.result):
+            if isinstance(msg, AIMessage):
+                return msg
+    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
+
+
+class ObservabilityMiddleware(AgentMiddleware):
+    """model/tool call 타이밍 + 메타데이터를 JSONL로 기록.
+
+    `agent_name` 은 행이 어느 에이전트의 호출인지 가른다 — 오케스트레이터는 `main`,
+    multi 모드 서브에이전트는 자기 이름이다 (이슈 #98).
+    """
+
+    def __init__(
+        self,
+        session_id: str = "",
+        user_identity: Optional[dict[str, str]] = None,
+        agent_name: str = "main",
+    ):
+        self._session_id = session_id
+        self._user_identity = user_identity or {}
+        self._agent_name = agent_name
+        self._model_call_count = 0
+        self._tool_call_count = 0
+
+    def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> Awaitable[ModelCallResult]:
+        return self._wrap_model(request, handler)
+
+    async def _wrap_model(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        self._model_call_count += 1
+        call_num = self._model_call_count
+        input_msg_count = len(request.messages) if request.messages else 0
+
+        meter = _TURN.get()
+        t0 = time.monotonic()
+        status = "ok"
+        output_tool_calls = 0
+        usage = _usage_of(None)
+        try:
+            response = await handler(request)
+            ai_msg = _extract_ai_message(response)
+            if ai_msg and hasattr(ai_msg, "tool_calls") and ai_msg.tool_calls:
+                output_tool_calls = len(ai_msg.tool_calls)
+            usage = _usage_of(ai_msg)
+            return response
+        except Exception as exc:
+            status = (
+                f"error:{exc.error_type}"
+                if isinstance(exc, AgentError)
+                else f"error:{type(exc).__name__}"
+            )
+            raise
+        finally:
+            elapsed_ms = (time.monotonic() - t0) * 1000
+            if meter is not None:
+                meter.add_model_call(usage, output_tool_calls)
+            _write_structured_event(
+                {
+                    "event": "model_call",
+                    **self._ids(meter),
+                    "agent_name": self._agent_name,
+                    **self._user_identity,
+                    "call_num": call_num,
+                    "input_msg_count": input_msg_count,
+                    "output_tool_calls": output_tool_calls,
+                    "elapsed_ms": round(elapsed_ms, 1),
+                    "status": status,
+                    **usage,
+                }
+            )
+
+    def _ids(self, meter: Optional[TurnMeter]) -> dict[str, str]:
+        """턴이 묶여 있으면 그 thread 를 쓴다 — 생성자 값은 캐시된 그래프의 첫 thread 다."""
+        if meter is None:
+            return {"session_id": self._session_id}
+        return {"session_id": meter.session_id, "ux_session_id": meter.ux_session_id}
+
+    async def awrap_tool_call(self, request, handler):
+        meter = _TURN.get()
+        tool_name = request.tool_call.get("name", "unknown")
+
+        t0 = time.monotonic()
+        status = "ok"
+        error_msg = None
+        try:
+            result = await handler(request)
+            return result
+        except GraphBubbleUp:
+            # 서브에이전트 HITL 이 task 도구 안에서 올라온 것 — 오류도 실행도 아니다.
+            # 재개하면 도구가 다시 돌고 그때 행이 나간다 (이슈 #98).
+            status = None
+            raise
+        except Exception as exc:
+            status = (
+                f"error:{exc.error_type}"
+                if isinstance(exc, AgentError)
+                else f"error:{type(exc).__name__}"
+            )
+            error_msg = str(exc)[:200]
+            raise
+        finally:
+            if status is not None:
+                self._write_tool_call(meter, tool_name, t0, status, error_msg)
+
+    def _write_tool_call(
+        self,
+        meter: Optional[TurnMeter],
+        tool_name: str,
+        t0: float,
+        status: str,
+        error_msg: Optional[str],
+    ) -> None:
+        if meter is not None:
+            meter.tool_calls += 1
+        self._tool_call_count += 1
+        call_num = self._tool_call_count
+        elapsed_ms = (time.monotonic() - t0) * 1000
+        ev: dict[str, Any] = {
+            "event": "tool_call",
+            **self._ids(meter),
+            "agent_name": self._agent_name,
+            **self._user_identity,
+            "tool_name": tool_name,
+            "call_num": call_num,
+            "elapsed_ms": round(elapsed_ms, 1),
+            "status": status,
+        }
+        if error_msg:
+            ev["error"] = error_msg
+        _write_structured_event(ev)
+
+    def get_stats(self) -> dict:
+        """현재까지의 model_calls, tool_calls 카운트 반환."""
+        return {
+            "model_calls": self._model_call_count,
+            "tool_calls": self._tool_call_count,
+        }

jupyter_ext/agent/patch_tool_calls.py

@@ -0,0 +1,57 @@
+"""Middleware to patch dangling tool calls in the messages history."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from langchain.agents.middleware import AgentMiddleware, AgentState
+from langchain_core.messages import RemoveMessage, ToolMessage
+from langgraph.graph.message import REMOVE_ALL_MESSAGES
+from langgraph.runtime import Runtime
+
+
+class PatchToolCallsMiddleware(AgentMiddleware):
+    """Middleware to patch dangling tool calls in the messages history.
+
+    Before the agent runs, scans AI messages for tool calls that were
+    never fulfilled and creates synthetic ToolMessages to prevent agent
+    state corruption when tool calls are incomplete.
+    """
+
+    def before_agent(
+        self, state: AgentState, runtime: Runtime[Any]
+    ) -> dict[str, Any] | None:  # noqa: ARG002
+        """Handle dangling tool calls from any AIMessage."""
+        messages = state["messages"]
+        if not messages or len(messages) == 0:
+            return None
+
+        patched_messages = []
+        for i, msg in enumerate(messages):
+            patched_messages.append(msg)
+            if msg.type == "ai" and msg.tool_calls:
+                for tool_call in msg.tool_calls:
+                    corresponding_tool_msg = next(
+                        (
+                            msg
+                            for msg in messages[i:]
+                            if msg.type == "tool"
+                            and msg.tool_call_id == tool_call["id"]
+                        ),
+                        None,
+                    )
+                    if corresponding_tool_msg is None:
+                        tool_msg = (
+                            f"Tool call {tool_call['name']} with id "
+                            f"{tool_call['id']} was cancelled - another "
+                            "message came in before it could be completed."
+                        )
+                        patched_messages.append(
+                            ToolMessage(
+                                content=tool_msg,
+                                name=tool_call["name"],
+                                tool_call_id=tool_call["id"],
+                            )
+                        )
+
+        return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES), *patched_messages]}

jupyter_ext/agent/prompts.py

@@ -0,0 +1,112 @@
+"""System prompts for the agent.
+
+multi 모드 planner 는 여기 없다 — `prompts.planner` 키다 (이슈 #218).
+"""
+
+SYSTEM_PROMPTS = {
+    "default": """You are Codebuff, an AI assistant that helps users with data science and programming tasks in Jupyter notebooks.
+
+You have access to various tools to help accomplish tasks:
+- jupyter_cell: Create, read, update, delete, or execute Jupyter notebook cells
+- jupyter_markdown: Create or update markdown cells for documentation
+- file_read: Read file contents
+- file_write: Write content to files
+- file_list: List directory contents
+- search_code: Search for patterns in code files
+- web_search: Search the web for information
+
+Guidelines:
+1. Always explain what you're going to do before taking action
+2. Use tools step by step - don't try to do everything at once
+3. After executing code, check the output and handle any errors
+4. Be concise but informative in your responses
+5. If you're unsure about something, ask the user for clarification
+
+When working with notebooks:
+- Create new cells for new code
+- Update existing cells when fixing or modifying code
+- Add markdown cells to explain your work
+- Execute cells to verify your code works""",
+    "data_science": """You are Codebuff, a specialized AI assistant for data science tasks in Jupyter notebooks.
+
+You excel at:
+- Data loading and preprocessing
+- Exploratory data analysis (EDA)
+- Statistical analysis
+- Data visualization with matplotlib, seaborn, plotly
+- Machine learning with scikit-learn, pandas
+- Deep learning with PyTorch, TensorFlow
+
+Available tools:
+- jupyter_cell: Create, read, update, delete, or execute Jupyter notebook cells
+- jupyter_markdown: Create or update markdown cells for documentation
+- file_read: Read data files and scripts
+- file_write: Save processed data or results
+- file_list: Explore data directories
+- search_code: Find code patterns
+- web_search: Look up documentation or examples
+
+Best practices:
+1. Start by understanding the data structure
+2. Handle missing values and outliers appropriately
+3. Visualize data before and after transformations
+4. Document your analysis steps with markdown cells
+5. Use meaningful variable names
+6. Test code incrementally""",
+    "code_review": """You are Codebuff, a hybrid code reviewer that combines automated tool scanning with LLM deep analysis.
+
+## 3-Phase Hybrid Code Review Workflow
+
+### Phase 1: 도구 스캔 (MANDATORY - 반드시 먼저 실행)
+코드 리뷰를 시작하기 전에 반드시 security_scan 도구를 호출하세요:
+- `security_scan(scan_type="all")` 을 실행하여 자동화된 스캐너 결과를 먼저 확보
+- 이 단계를 건너뛰지 마세요. 도구 스캔 결과가 Phase 2의 입력이 됩니다.
+- 스캔 대상: SAST(bandit, semgrep), 시크릿 탐지(detect-secrets), 우회 패턴(evasion), 개인정보(PII)
+
+### Phase 2: LLM 심층 분석
+도구 스캔 결과를 바탕으로 다음을 추가 분석하세요:
+- **중복 제거**: 도구가 동일 이슈를 다른 이름으로 보고한 경우 통합
+- **비즈니스 로직 취약점**: 도구가 탐지 못하는 논리적 결함, 레이스 컨디션, 권한 상승 등
+- **한국 금융 컴플라이언스**: 개인정보보호법, 신용정보법, 전자금융거래법 관점의 준수 여부
+- **컨텍스트 분석**: 코드의 실제 의도와 사용 맥락을 고려한 위양성(false positive) 필터링
+
+### Phase 3: 통합 리포트 (5-관점 포맷)
+아래 5개 섹션으로 결과를 정리하세요. 각 발견사항에 출처 태그를 명시:
+
+1. **시큐어코딩** - 입력 검증, 인젝션, 민감 정보 노출, 크리덴셜 관리
+   - 각 항목에 [도구발견] 또는 [LLM분석] 태그 표기
+   - 심각도 표기: 🔴 Critical, 🟠 High, 🟡 Medium, 🔵 Low
+   - 하나의 셀에 여러 취약점이 있으면 각각 별도 항목으로 분리
+   - 은닉 패턴(변수 조합, 역순, ROT13, base64, 함수 기본 인자)은 패턴별로 개별 지적
+   - DataFrame PII는 모든 PII 컬럼을 나열하고 각각 마스킹 방안 제시
+2. **개인정보보호(PII)** - 한국 금융권 개인정보 패턴, 마스킹, 암호화, 로깅 안전성
+   - 관련 법률 명시 (개인정보보호법/신용정보법/전자금융거래법)
+3. **성능** - OOM 방지, 전처리·학습 속도 개선, 컴퓨팅 리소스 최적 활용. 불필요한 DataFrame 복사, 비효율적 dtypes, 중복 연산 등 DS 실무 병목 식별 (프롬프트에 실행 환경 사양이 포함된 경우 해당 환경에 맞춘 구체적 조언 제공)
+4. **코드 품질** - 구조, 일관성, 에러 처리, 타입 힌트
+5. **요약** - 종합 평가, 도구 스캔 통계, 개선 우선순위 Top 3, 개선 필요 셀의 전체 교체 코드
+
+Available tools:
+- jupyter_cell: Examine and modify code cells
+- jupyter_markdown: Add review comments and documentation
+- file_read: Read source files
+- search_code: Find related code patterns
+- web_search: Look up best practices
+- security_scan: 자동화된 보안 스캔 실행 (bandit, semgrep, detect-secrets, evasion, pii)
+
+Important formatting rules:
+- 개선 코드는 반드시 마크다운 코드 블록(```python)으로 제시하세요.
+- 절대 표(table) 형식 안에 코드를 넣지 마세요.""",
+}
+
+
+def get_system_prompt(mode: str = "default") -> str:
+    """
+    Get system prompt for the specified mode.
+
+    Args:
+        mode: Agent mode (default, data_science, code_review)
+
+    Returns:
+        System prompt string
+    """
+    return SYSTEM_PROMPTS.get(mode, SYSTEM_PROMPTS["default"])