deeptutor-memory 1.7.0__py3-none-any.whl

deeptutor/services/memory/consolidator/runs.py

@@ -0,0 +1,405 @@
+"""Persistent, cancellable consolidator runs.
+
+A "run" is one invocation of :func:`run_update` / :func:`run_audit` /
+:func:`run_dedup`. The run is owned by an asyncio task; events flow
+through a buffered ring so a disconnecting client can re-attach by
+posting ``since=<cursor>`` to the events endpoint and replay everything
+it missed.
+
+Why an in-memory manager instead of a DB:
+- Memory consolidator runs are minutes-long at most.
+- Crash / restart wipes them — that's acceptable; the docs themselves
+  are atomically written per step and the meta-id-diff still gives us
+  "what's new since last refresh" correctness on restart.
+
+Concurrency rules
+-----------------
+At most one **active** run per ``(layer, key)``. Starting a second run
+while the first is active returns ``RunBusyError``. Once a run reaches
+a terminal status (``done`` / ``cancelled`` / ``error``) it stays in
+the registry indefinitely so the UI can re-attach to see the final
+trace; older runs evict on FIFO when ``_MAX_HISTORY`` is exceeded.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+import logging
+from pathlib import Path
+from typing import Any, Awaitable, Callable, Literal
+import uuid
+
+from deeptutor.services.file_io import atomic_write_text as _atomic_write
+
+logger = logging.getLogger(__name__)
+
+RunMode = Literal["update", "audit", "dedup"]
+RunStatus = Literal["queued", "running", "cancelled", "done", "error"]
+
+_MAX_EVENTS_PER_RUN = 2000
+_MAX_HISTORY = 200
+
+
+@dataclass
+class RunEvent:
+    seq: int  # 0-based, monotonic per run
+    ts: str  # ISO-8601 UTC
+    payload: dict[str, Any]
+
+
+@dataclass
+class UndoCheckpoint:
+    id: str
+    ts: str
+    layer: str
+    key: str
+    path: str
+    existed: bool
+    previous_content: str
+    action: str
+    turn: int | None = None
+    label: str | None = None
+
+
+@dataclass
+class Run:
+    id: str
+    layer: str
+    key: str
+    mode: RunMode
+    params: dict[str, Any]
+    language: str
+    user_label: str
+    status: RunStatus = "queued"
+    started_at: str = ""
+    ended_at: str | None = None
+    error: str | None = None
+    events: list[RunEvent] = field(default_factory=list)
+    undo_stack: list[UndoCheckpoint] = field(default_factory=list)
+    _next_seq: int = field(default=0, repr=False)
+    _waiters: list[asyncio.Event] = field(default_factory=list, repr=False)
+    _task: asyncio.Task | None = field(default=None, repr=False)
+    _cancel_flag: asyncio.Event = field(default_factory=asyncio.Event, repr=False)
+
+    @property
+    def active(self) -> bool:
+        return self.status in ("queued", "running")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "layer": self.layer,
+            "key": self.key,
+            "mode": self.mode,
+            "params": self.params,
+            "language": self.language,
+            "status": self.status,
+            "started_at": self.started_at,
+            "ended_at": self.ended_at,
+            "error": self.error,
+            "event_count": len(self.events),
+            "undo_count": len(self.undo_stack),
+        }
+
+
+class RunBusyError(RuntimeError):
+    """Raised when a (layer, key) already has an active run."""
+
+
+# ContextVar holds the current Run for handlers running inside the task.
+# Mode code uses :func:`emit` indirectly via ``modes._runtime.emit``;
+# we install our own on_event that pushes into the active run too.
+_current_run: ContextVar[Run | None] = ContextVar("memory_run", default=None)
+
+
+class RunManager:
+    """Process-wide singleton — one manager owns every consolidator run.
+
+    The instance is created on first call to :func:`get_run_manager`.
+    """
+
+    def __init__(self) -> None:
+        self._runs: dict[str, Run] = {}
+        self._order: list[str] = []  # FIFO for eviction
+        self._active: dict[tuple[str, str], str] = {}
+        self._lock = asyncio.Lock()
+
+    # ── Lookup ─────────────────────────────────────────────────────────
+
+    def get(self, run_id: str) -> Run | None:
+        return self._runs.get(run_id)
+
+    def active_for(self, layer: str, key: str) -> Run | None:
+        run_id = self._active.get((layer, key))
+        if run_id is None:
+            return None
+        run = self._runs.get(run_id)
+        return run if run is not None and run.active else None
+
+    def list_for(self, layer: str | None = None, key: str | None = None) -> list[Run]:
+        out: list[Run] = []
+        for rid in self._order:
+            run = self._runs.get(rid)
+            if run is None:
+                continue
+            if layer is not None and run.layer != layer:
+                continue
+            if key is not None and run.key != key:
+                continue
+            out.append(run)
+        return out
+
+    # ── Start ──────────────────────────────────────────────────────────
+
+    async def start(
+        self,
+        *,
+        layer: str,
+        key: str,
+        mode: RunMode,
+        runner: Callable[[Callable[[dict[str, Any]], Awaitable[None]]], Awaitable[None]],
+        params: dict[str, Any] | None = None,
+        language: str = "en",
+        user_label: str = "anonymous",
+    ) -> Run:
+        """Register and kick off a new run.
+
+        ``runner`` is an awaitable factory: takes a ``on_event`` callback
+        and runs the consolidator mode. The manager wires the callback to
+        the event buffer + waiter machinery.
+        """
+        async with self._lock:
+            if self.active_for(layer, key) is not None:
+                raise RunBusyError(f"a run is already in progress for {layer}/{key}")
+            run = Run(
+                id=uuid.uuid4().hex,
+                layer=layer,
+                key=key,
+                mode=mode,
+                params=dict(params or {}),
+                language=language,
+                user_label=user_label,
+                status="queued",
+                started_at=_now_iso(),
+            )
+            self._runs[run.id] = run
+            self._order.append(run.id)
+            self._active[(layer, key)] = run.id
+            self._evict_if_needed()
+
+        run._task = asyncio.create_task(self._drive(run, runner))
+        return run
+
+    async def cancel(self, run_id: str) -> bool:
+        run = self._runs.get(run_id)
+        if run is None or not run.active:
+            return False
+        run._cancel_flag.set()
+        if run._task is not None and not run._task.done():
+            run._task.cancel()
+        return True
+
+    async def undo_last(self, run_id: str) -> RunEvent | None:
+        """Restore the document snapshot before the latest run write."""
+        run = self._runs.get(run_id)
+        if run is None:
+            raise KeyError(run_id)
+        if run.active:
+            raise RunBusyError("cancel the active run before undoing memory edits")
+        if not run.undo_stack:
+            return None
+
+        checkpoint = run.undo_stack.pop()
+        path = Path(checkpoint.path)
+        if checkpoint.existed:
+            await asyncio.to_thread(_atomic_write, path, checkpoint.previous_content)
+        else:
+            await asyncio.to_thread(_remove_if_exists, path)
+
+        return await self._emit(
+            run,
+            {
+                "stage": "undo_applied",
+                "run_id": run.id,
+                "undo_id": checkpoint.id,
+                "undo_depth": len(run.undo_stack),
+                "layer": checkpoint.layer,
+                "key": checkpoint.key,
+                "turn": checkpoint.turn,
+                "label": checkpoint.label,
+                "action": checkpoint.action,
+            },
+        )
+
+    # ── Event subscription ─────────────────────────────────────────────
+
+    async def wait_for_events(self, run: Run, *, since: int) -> list[RunEvent]:
+        """Return events since cursor; block until new ones arrive or done."""
+        if since < 0:
+            since = 0
+        # Fast path: events already buffered past cursor.
+        buffered = [event for event in run.events if event.seq >= since]
+        if buffered:
+            return buffered
+        if not run.active:
+            return []
+        waiter = asyncio.Event()
+        run._waiters.append(waiter)
+        try:
+            await waiter.wait()
+        finally:
+            try:
+                run._waiters.remove(waiter)
+            except ValueError:
+                pass
+        return [event for event in run.events if event.seq >= since]
+
+    # ── Drive ──────────────────────────────────────────────────────────
+
+    async def _drive(
+        self,
+        run: Run,
+        runner: Callable[[Callable[[dict[str, Any]], Awaitable[None]]], Awaitable[None]],
+    ) -> None:
+        token = _current_run.set(run)
+        run.status = "running"
+        await self._emit(run, {"stage": "run_started", "run_id": run.id, "mode": run.mode})
+        try:
+
+            async def on_event(evt: dict[str, Any]) -> None:
+                await self._emit(run, evt)
+
+            await runner(on_event)
+            if run.status == "running":
+                run.status = "done"
+        except asyncio.CancelledError:
+            run.status = "cancelled"
+            await self._emit(run, {"stage": "cancelled"})
+        except Exception as exc:  # noqa: BLE001
+            run.status = "error"
+            run.error = str(exc)
+            logger.warning(
+                "consolidator run failed (id=%s layer=%s key=%s mode=%s): %s",
+                run.id,
+                run.layer,
+                run.key,
+                run.mode,
+                exc,
+                exc_info=True,
+            )
+            await self._emit(run, {"stage": "error", "message": str(exc)})
+        finally:
+            run.ended_at = _now_iso()
+            await self._emit(run, {"stage": "run_ended", "status": run.status})
+            self._active.pop((run.layer, run.key), None)
+            # Wake any remaining waiters so they observe the terminal state.
+            for w in list(run._waiters):
+                w.set()
+            _current_run.reset(token)
+
+    async def _emit(self, run: Run, payload: dict[str, Any]) -> RunEvent:
+        event = RunEvent(seq=run._next_seq, ts=_now_iso(), payload=payload)
+        run._next_seq += 1
+        run.events.append(event)
+        if len(run.events) > _MAX_EVENTS_PER_RUN:
+            # Cursors are absolute. Never renumber retained events: a client
+            # may reconnect with a cursor issued before this ring wrapped.
+            run.events.pop(0)
+        for w in list(run._waiters):
+            w.set()
+        return event
+
+    def _evict_if_needed(self) -> None:
+        while len(self._order) > _MAX_HISTORY:
+            old = self._order.pop(0)
+            run = self._runs.pop(old, None)
+            if run is not None and run.active:
+                # Active runs are protected from eviction.
+                self._runs[old] = run
+                self._order.insert(0, old)
+                return
+
+
+_manager: RunManager | None = None
+
+
+def get_run_manager() -> RunManager:
+    global _manager
+    if _manager is None:
+        _manager = RunManager()
+    return _manager
+
+
+def reset_run_manager_for_tests() -> None:
+    global _manager
+    _manager = None
+
+
+def current_run() -> Run | None:
+    """Return the run currently driving the active task, if any.
+
+    Used by the LLM-IO event emitter so it can attach per-turn payloads
+    to whichever run is calling out to the model.
+    """
+    return _current_run.get()
+
+
+def push_undo_checkpoint(
+    *,
+    layer: str,
+    key: str,
+    path: Path,
+    existed: bool,
+    previous_content: str,
+    action: str,
+    turn: int | None = None,
+    label: str | None = None,
+) -> int:
+    """Register a per-write rollback snapshot on the active run."""
+    run = _current_run.get()
+    if run is None:
+        return 0
+    run.undo_stack.append(
+        UndoCheckpoint(
+            id=uuid.uuid4().hex,
+            ts=_now_iso(),
+            layer=layer,
+            key=key,
+            path=str(path),
+            existed=existed,
+            previous_content=previous_content,
+            action=action,
+            turn=turn,
+            label=label,
+        )
+    )
+    return len(run.undo_stack)
+
+
+def _now_iso() -> str:
+    return datetime.now(tz=timezone.utc).isoformat()
+
+
+def _remove_if_exists(path: Path) -> None:
+    try:
+        path.unlink()
+    except FileNotFoundError:
+        return
+
+
+__all__ = [
+    "Run",
+    "RunBusyError",
+    "RunEvent",
+    "RunManager",
+    "RunMode",
+    "RunStatus",
+    "UndoCheckpoint",
+    "current_run",
+    "get_run_manager",
+    "push_undo_checkpoint",
+    "reset_run_manager_for_tests",
+]

deeptutor/services/memory/document.py

@@ -0,0 +1,235 @@
+"""Markdown documents with footnote-style citations.
+
+Each L2/L3 file is a markdown document of the form::
+
+    # <Title>
+
+    ## <section_a>
+    - <text> [^1][^2] <!--m_xxx-->
+    - <text> [^1] <!--m_yyy-->
+
+    ## <section_b>
+    - <text> [^3] <!--m_zzz-->
+
+    ---
+
+    [^1]: notebook:abc
+    [^2]: chat:def
+    [^3]: chat:ghi
+
+Footnote labels are *integers* assigned per-document in first-appearance
+order over the bullet stream. Two entries citing the same source share a
+label, so duplicate footnote rows disappear from the rendered view.
+
+The HTML comment after each bullet (``<!--m_xxx-->``) is the entry id
+anchor. It survives round-trips and is used by audit / dedup line views
+and by ``DELETE /entry/{id}``. Parser also accepts the *legacy* format
+where the bullet ends in ``[^m_xxx]`` and footnote rows are
+``[^m_xxx]: ref1, ref2`` — this lets pre-existing docs continue working
+until the next save migrates them to the new layout.
+
+Parsing and serialization are pure functions — no I/O, no LLM. The
+round-trip ``serialize(parse(x))`` is idempotent for any document
+produced by ``serialize``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+import re
+
+_ENTRY_ID_RE = r"m_[0-9A-HJKMNP-TV-Z]{26}"
+
+_TITLE_RE = re.compile(r"^#\s+(.+?)\s*$")
+_SECTION_RE = re.compile(r"^##\s+(.+?)\s*$")
+
+# New bullet:  "- text [^1], [^3] <!--m_xxx-->"
+# Markers are optional (an entry may cite no refs); commas + whitespace
+# between markers are tolerated so the rendered superscripts read
+# ``¹, ³`` instead of the visually-merged ``¹³``.
+_NEW_BULLET_RE = re.compile(
+    rf"^\s*-\s+(?P<text>.*?)(?P<markers>(?:\s*,?\s*\[\^[^\]]+\])*)\s*<!--\s*(?P<id>{_ENTRY_ID_RE})\s*-->\s*$"
+)
+# Legacy bullet: "- text[^m_xxx]"
+_OLD_BULLET_RE = re.compile(rf"^\s*-\s+(?P<text>.*?)\[\^(?P<id>{_ENTRY_ID_RE})\]\s*$")
+
+# Legacy footnote def: "[^m_xxx]: ref1, ref2"
+_OLD_FOOTNOTE_RE = re.compile(rf"^\[\^(?P<id>{_ENTRY_ID_RE})\]:\s*(?P<refs>.*?)\s*$")
+# New footnote def: "[^1]: notebook:abc"   (label is non-m_xxx)
+_NEW_FOOTNOTE_RE = re.compile(r"^\[\^(?P<label>[^\]]+)\]:\s*(?P<ref>.*?)\s*$")
+
+_MARKER_RE = re.compile(r"\[\^([^\]]+)\]")
+
+
+@dataclass
+class Entry:
+    id: str
+    section: str
+    text: str
+    refs: list[str] = field(default_factory=list)
+
+
+@dataclass
+class Document:
+    title: str = ""
+    sections: list[tuple[str, list[Entry]]] = field(default_factory=list)
+
+    def all_entries(self) -> list[Entry]:
+        return [e for _, entries in self.sections for e in entries]
+
+    def find(self, entry_id: str) -> Entry | None:
+        for _, entries in self.sections:
+            for entry in entries:
+                if entry.id == entry_id:
+                    return entry
+        return None
+
+    def section_entries(self, name: str) -> list[Entry]:
+        """Return the entry list for ``name``, creating the section if absent."""
+        for section, entries in self.sections:
+            if section == name:
+                return entries
+        new_entries: list[Entry] = []
+        self.sections.append((name, new_entries))
+        return new_entries
+
+    def remove(self, entry_id: str) -> bool:
+        for _, entries in self.sections:
+            for i, entry in enumerate(entries):
+                if entry.id == entry_id:
+                    del entries[i]
+                    return True
+        return False
+
+
+def parse(md: str) -> Document:
+    """Parse memory md in either the new (ref-keyed) or legacy (entry-keyed) format."""
+    raw_lines = md.splitlines()
+
+    # Pass 1 — collect every footnote definition. We accept BOTH:
+    # * new ref-keyed: ``[^1]: notebook:abc``  → ref by label
+    # * old entry-keyed: ``[^m_xxx]: r1, r2``  → refs by entry id
+    refs_by_entry: dict[str, list[str]] = {}
+    ref_by_label: dict[str, str] = {}
+    for raw in raw_lines:
+        line = raw.rstrip()
+        m_old_fn = _OLD_FOOTNOTE_RE.match(line)
+        if m_old_fn:
+            refs_raw = m_old_fn.group("refs")
+            refs_by_entry[m_old_fn.group("id")] = [
+                r.strip() for r in refs_raw.split(",") if r.strip()
+            ]
+            continue
+        m_new_fn = _NEW_FOOTNOTE_RE.match(line)
+        if m_new_fn:
+            label = m_new_fn.group("label")
+            if label.startswith("m_"):
+                # Skip — that was an entry-keyed row already handled above.
+                continue
+            ref_by_label[label] = m_new_fn.group("ref").strip()
+
+    # Pass 2 — title, sections, bullets.
+    doc = Document()
+    current_entries: list[Entry] | None = None
+    current_section: str | None = None
+    for raw in raw_lines:
+        line = raw.rstrip()
+
+        if not doc.title:
+            m_title = _TITLE_RE.match(line)
+            if m_title:
+                doc.title = m_title.group(1).strip()
+                continue
+
+        m_section = _SECTION_RE.match(line)
+        if m_section:
+            current_section = m_section.group(1).strip()
+            current_entries = []
+            doc.sections.append((current_section, current_entries))
+            continue
+
+        # New format first: bullet ends with HTML-comment entry-id anchor.
+        m_new_b = _NEW_BULLET_RE.match(line)
+        if m_new_b and current_entries is not None and current_section is not None:
+            entry_id = m_new_b.group("id")
+            text = m_new_b.group("text").rstrip()
+            markers = _MARKER_RE.findall(m_new_b.group("markers") or "")
+            entry_refs: list[str] = []
+            for marker in markers:
+                ref = ref_by_label.get(marker)
+                if ref is not None and ref not in entry_refs:
+                    entry_refs.append(ref)
+            current_entries.append(
+                Entry(id=entry_id, section=current_section, text=text, refs=entry_refs)
+            )
+            continue
+
+        # Legacy bullet: refs come from refs_by_entry built in pass 1.
+        m_old_b = _OLD_BULLET_RE.match(line)
+        if m_old_b and current_entries is not None and current_section is not None:
+            entry_id = m_old_b.group("id")
+            text = m_old_b.group("text").strip()
+            current_entries.append(
+                Entry(
+                    id=entry_id,
+                    section=current_section,
+                    text=text,
+                    refs=list(refs_by_entry.get(entry_id, [])),
+                )
+            )
+            continue
+
+    return doc
+
+
+def serialize(doc: Document) -> str:
+    """Render the doc in the new consolidated, ref-keyed format.
+
+    Every unique ref across all entries gets one footnote label, assigned
+    in first-appearance order. Bullets cite their refs as ``[^1][^3]``
+    inline. The entry id is preserved as a trailing HTML comment so the
+    round-trip ``parse(serialize(d)) == d``.
+    """
+    # 1. Build the consolidated ref → label map in first-appearance order.
+    ref_order: list[str] = []
+    ref_to_label: dict[str, int] = {}
+    for entry in doc.all_entries():
+        for ref in entry.refs:
+            if ref in ref_to_label:
+                continue
+            ref_to_label[ref] = len(ref_order) + 1
+            ref_order.append(ref)
+
+    lines: list[str] = []
+    if doc.title:
+        lines.append(f"# {doc.title}")
+        lines.append("")
+
+    for section, entries in doc.sections:
+        if not entries:
+            continue
+        lines.append(f"## {section}")
+        lines.append("")
+        for entry in entries:
+            # ``, ``-separate markers so the rendered superscripts read
+            # "¹, ²" not "¹²" — important when the same bullet cites two
+            # different sources.
+            markers = ", ".join(f"[^{ref_to_label[r]}]" for r in entry.refs if r in ref_to_label)
+            text = entry.text.rstrip()
+            if markers:
+                lines.append(f"- {text} {markers} <!--{entry.id}-->")
+            else:
+                lines.append(f"- {text} <!--{entry.id}-->")
+        lines.append("")
+
+    if ref_order:
+        lines.append("---")
+        lines.append("")
+        for i, ref in enumerate(ref_order, start=1):
+            lines.append(f"[^{i}]: {ref}")
+        lines.append("")
+
+    return "\n".join(lines).rstrip() + "\n"
+
+
+__all__ = ["Document", "Entry", "parse", "serialize"]