codebuff-agent 2.1.74__py3-none-any.whl

jupyter_ext/agent/middleware.py

@@ -0,0 +1,359 @@
+"""DS 코딩 에이전트 커스텀 미들웨어.
+
+데이터 사이언스 워크플로우에 특화된 미들웨어:
+- DSMemoryMiddleware: 세션 메모리 (데이터셋, 실험 이력) 자동 추적
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import re
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+from langchain_core.messages import AIMessage, AnyMessage, SystemMessage
+
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+)
+
+from .middleware_prompts import DS_MEMORY_INJECTION_TEMPLATE
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# MemoryStore
+# ---------------------------------------------------------------------------
+
+MEMORY_CATEGORIES = (
+    "data_source",
+    "user_preference",
+    "experiment",
+    "environment",
+    "decision",
+    "finding",
+)
+
+
+@dataclass
+class MemoryEntry:
+    """세션 메모리 단일 항목."""
+
+    id: str
+    category: str
+    content: str
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+    access_count: int = 0
+    session_id: str = ""
+    tags: list[str] = field(default_factory=list)
+    persistent: bool = True  # False = 세션 종료 시 삭제, 디스크 저장 안 함
+
+
+def _entry_to_dict(entry: MemoryEntry) -> dict:
+    """MemoryEntry → JSON-serializable dict. persistent 필드는 포함하지 않음."""
+    return {
+        "id": entry.id,
+        "category": entry.category,
+        "content": entry.content,
+        "created_at": entry.created_at.isoformat(),
+        "access_count": entry.access_count,
+        "session_id": entry.session_id,
+        "tags": entry.tags,
+    }
+
+
+def _dict_to_entry(d: dict) -> MemoryEntry:
+    """dict → MemoryEntry. 디스크에서 복원된 항목은 항상 persistent=True."""
+    return MemoryEntry(
+        id=d["id"],
+        category=d["category"],
+        content=d["content"],
+        created_at=datetime.fromisoformat(d["created_at"]),
+        access_count=d.get("access_count", 0),
+        session_id=d.get("session_id", ""),
+        tags=d.get("tags", []),
+    )
+
+
+class MemoryStore:
+    """키워드 기반 메모리 저장소 (선택적 파일 영속화).
+
+    storage_path가 주어지면 persistent=True인 항목을 JSON 파일로 저장.
+    """
+
+    MAX_PERSISTENT_ENTRIES = 500
+
+    def __init__(self, storage_path: str | Path | None = None) -> None:
+        self._memories: dict[str, MemoryEntry] = {}
+        self._storage_path: Path | None = Path(storage_path) if storage_path else None
+        self._flush_timer: asyncio.TimerHandle | None = None
+        self._load()
+
+    # -- persistence helpers --------------------------------------------------
+
+    def _load(self) -> None:
+        """시작 시 디스크에서 메모리 로드."""
+        if not self._storage_path or not self._storage_path.exists():
+            return
+        try:
+            data = json.loads(self._storage_path.read_text(encoding="utf-8"))
+            for d in data:
+                entry = _dict_to_entry(d)
+                self._memories[entry.id] = entry
+            logger.info(
+                "MemoryStore: loaded %d entries from %s",
+                len(self._memories),
+                self._storage_path,
+            )
+        except (json.JSONDecodeError, KeyError, OSError) as e:
+            logger.warning("MemoryStore: failed to load memories: %s", e)
+
+    def _schedule_flush(self) -> None:
+        """디바운스 디스크 저장 예약 (0.5초 후)."""
+        if not self._storage_path:
+            return
+        if self._flush_timer is not None:
+            self._flush_timer.cancel()
+        try:
+            loop = asyncio.get_running_loop()
+            self._flush_timer = loop.call_later(
+                0.5, lambda: loop.create_task(self._flush())
+            )
+        except RuntimeError:
+            # 이벤트 루프 밖에서 호출된 경우 동기 저장
+            self._flush_sync()
+
+    def _flush_sync(self) -> None:
+        """동기적으로 persistent 항목을 디스크에 저장 (atomic write)."""
+        if not self._storage_path:
+            return
+        persistent = [
+            _entry_to_dict(e) for e in self._memories.values() if e.persistent
+        ]
+        # 최대 개수 초과 시 오래된 것부터 제거
+        if len(persistent) > self.MAX_PERSISTENT_ENTRIES:
+            persistent.sort(key=lambda d: d["created_at"])
+            persistent = persistent[-self.MAX_PERSISTENT_ENTRIES :]
+        self._storage_path.parent.mkdir(parents=True, exist_ok=True)
+        tmp = self._storage_path.with_suffix(".tmp")
+        tmp.write_text(
+            json.dumps(persistent, ensure_ascii=False, indent=2), encoding="utf-8"
+        )
+        tmp.replace(self._storage_path)  # atomic rename
+
+    async def _flush(self) -> None:
+        """비동기 디스크 저장."""
+        await asyncio.to_thread(self._flush_sync)
+
+    def flush_now(self) -> None:
+        """즉시 디스크에 저장 (세션 종료 시 호출)."""
+        if self._flush_timer is not None:
+            self._flush_timer.cancel()
+            self._flush_timer = None
+        self._flush_sync()
+
+    def clear_session(self, session_id: str) -> None:
+        """특정 세션의 transient 메모리만 제거."""
+        to_remove = [
+            k
+            for k, v in self._memories.items()
+            if not v.persistent and v.session_id == session_id
+        ]
+        for k in to_remove:
+            del self._memories[k]
+
+    # -- public API -----------------------------------------------------------
+
+    def add(
+        self,
+        content: str,
+        category: str,
+        *,
+        session_id: str = "",
+        tags: list[str] | None = None,
+        persistent: bool = True,
+    ) -> MemoryEntry:
+        """메모리 추가."""
+        entry = MemoryEntry(
+            id=str(uuid.uuid4()),
+            category=category,
+            content=content,
+            session_id=session_id,
+            tags=tags or [],
+            persistent=persistent,
+        )
+        self._memories[entry.id] = entry
+        if persistent:
+            self._schedule_flush()
+        return entry
+
+    def search(self, query: str, *, limit: int = 10) -> list[MemoryEntry]:
+        """키워드 기반 검색. 대소문자 무시."""
+        if not query.strip():
+            return []
+
+        keywords = query.lower().split()
+        scored: list[tuple[int, MemoryEntry]] = []
+
+        for entry in self._memories.values():
+            text = f"{entry.content} {entry.category} {' '.join(entry.tags)}".lower()
+            score = sum(1 for kw in keywords if kw in text)
+            if score > 0:
+                scored.append((score, entry))
+
+        scored.sort(key=lambda x: (-x[0], x[1].created_at))
+        results = [entry for _, entry in scored[:limit]]
+
+        for entry in results:
+            entry.access_count += 1
+
+        return results
+
+    def get_by_category(self, category: str, *, limit: int = 10) -> list[MemoryEntry]:
+        """카테고리별 메모리 조회."""
+        entries = [e for e in self._memories.values() if e.category == category]
+        entries.sort(key=lambda e: e.created_at, reverse=True)
+        return entries[:limit]
+
+    def all(self) -> list[MemoryEntry]:
+        """모든 메모리 조회."""
+        return list(self._memories.values())
+
+    def clear(self) -> None:
+        """모든 메모리 삭제."""
+        self._memories.clear()
+        self._schedule_flush()
+
+    @property
+    def count(self) -> int:
+        return len(self._memories)
+
+
+# ---------------------------------------------------------------------------
+# 자동 추출 규칙
+# ---------------------------------------------------------------------------
+
+_EXTRACTION_PATTERNS: list[tuple[str, str, re.Pattern[str]]] = [
+    # (카테고리, 설명 접두어, 정규식 패턴)
+    (
+        "data_source",
+        "Loaded data",
+        re.compile(
+            r"(?:read_csv|read_parquet|read_excel|read_json|load_dataset)\(['\"]([^'\"]+)['\"]\)"
+        ),
+    ),
+    ("experiment", "Model training", re.compile(r"\.(fit|train)\s*\(")),
+    (
+        "experiment",
+        "Model metric",
+        re.compile(
+            r"(?:accuracy|f1_score|rmse|auc|precision|recall|r2_score|mean_squared_error)\s*[:=]?\s*([\d.]+)"
+        ),
+    ),
+    ("finding", "Shape info", re.compile(r"\.shape\s*(?:==|is)?\s*\((\d+),\s*(\d+)\)")),
+]
+
+
+def _extract_facts_from_messages(messages: list[AnyMessage]) -> list[tuple[str, str]]:
+    """메시지에서 DS 관련 팩트를 자동 추출.
+
+    Returns:
+        (category, content) 튜플 리스트
+    """
+    facts: list[tuple[str, str]] = []
+
+    for msg in messages:
+        if not hasattr(msg, "content") or not isinstance(msg.content, str):
+            continue
+
+        text = msg.content
+
+        for category, prefix, pattern in _EXTRACTION_PATTERNS:
+            for match in pattern.finditer(text):
+                fact_content = f"{prefix}: {match.group(0)}"
+                facts.append((category, fact_content))
+
+    return facts
+
+
+# ---------------------------------------------------------------------------
+# DSMemoryMiddleware
+# ---------------------------------------------------------------------------
+
+
+class DSMemoryMiddleware(AgentMiddleware):
+    """DS 세션 메모리 미들웨어.
+
+    before_model: 관련 메모리를 SystemMessage에 주입
+    after_model: AI 응답에서 팩트를 자동 추출하여 저장
+    """
+
+    def __init__(
+        self,
+        memory_store: MemoryStore,
+        *,
+        max_memories_per_call: int = 10,
+        auto_extract: bool = True,
+    ) -> None:
+        super().__init__()
+        self.memory_store = memory_store
+        self.max_memories_per_call = max_memories_per_call
+        self.auto_extract = auto_extract
+        self.tools = []
+
+    def before_model(
+        self, state: AgentState[Any], runtime: Any
+    ) -> dict[str, Any] | None:
+        """관련 메모리를 메시지에 주입."""
+        if self.memory_store.count == 0:
+            return None
+
+        # 마지막 사용자 메시지에서 쿼리 추출
+        query = ""
+        for msg in reversed(state["messages"]):
+            if hasattr(msg, "content") and isinstance(msg.content, str):
+                if not isinstance(msg, AIMessage):
+                    query = msg.content[:200]
+                    break
+
+        if not query:
+            return None
+
+        memories = self.memory_store.search(query, limit=self.max_memories_per_call)
+        if not memories:
+            return None
+
+        memory_lines = [f"- [{m.category}] {m.content}" for m in memories]
+        memory_text = DS_MEMORY_INJECTION_TEMPLATE.format(
+            memories="\n".join(memory_lines)
+        )
+
+        return {
+            "messages": [SystemMessage(content=memory_text)],
+        }
+
+    def after_model(
+        self, state: AgentState[Any], runtime: Any
+    ) -> dict[str, Any] | None:
+        """AI 응답에서 팩트를 자동 추출."""
+        if not self.auto_extract:
+            return None
+
+        # 마지막 AI 메시지에서 팩트 추출
+        recent_messages = state["messages"][-3:]
+        facts = _extract_facts_from_messages(recent_messages)
+
+        for category, content in facts:
+            # 중복 방지: 동일 내용이 이미 있으면 스킵
+            existing = self.memory_store.search(content, limit=1)
+            if existing and existing[0].content == content:
+                continue
+            self.memory_store.add(content, category)
+
+        return None

jupyter_ext/agent/middleware_prompts.py

@@ -0,0 +1,53 @@
+"""DS 코딩 에이전트 미들웨어 프롬프트.
+
+SummarizationMiddleware, DSMemoryMiddleware, HumanInTheLoopMiddleware에서
+사용하는 프롬프트 템플릿 정의.
+"""
+
+DS_SUMMARIZATION_PROMPT = """<role>
+Data Science Session Context Extractor
+</role>
+
+<primary_objective>
+Extract the most important context from this data science conversation to preserve continuity.
+</primary_objective>
+
+<instructions>
+You must populate each section with relevant information or state "None" if nothing to report.
+
+## SESSION INTENT
+What is the user's primary data science goal? What dataset(s) are involved?
+
+## DATA CONTEXT
+- Dataset schemas: column names, dtypes, shapes (e.g., `df: (1000, 15) with columns [id, name, price, ...]`)
+- Data sources: file paths, URLs, database connections
+- Variable names in use: df, X_train, X_test, model, scaler, etc.
+- Data quality issues discovered: missing values, outliers, duplicates
+
+## ANALYSIS PHASE
+Where are we in the DS workflow? (EDA / Preprocessing / Feature Engineering / Modeling / Evaluation / Reporting)
+
+## KEY FINDINGS
+- Statistical results, correlations, distributions discovered
+- Model metrics: accuracy, F1, RMSE, AUC, hyperparameters used
+- Visualizations created and their insights
+- Failed approaches and WHY they failed (important for avoiding repetition)
+
+## CODE EXECUTION STATE
+- Last executed code and its output summary
+- Active variables and their current state
+- Installed packages or environment changes
+
+## NEXT STEPS
+What specific tasks remain? What should be done next?
+</instructions>
+
+<messages>
+Messages to summarize:
+{messages}
+</messages>"""
+
+DS_MEMORY_INJECTION_TEMPLATE = """## Relevant Session Memories
+{memories}"""
+
+DS_HITL_DESCRIPTION_PREFIX = "Agent action requires your approval"