Clean up nested duplicate: v2/v2/engine/memory.py

v2/v2/engine/memory.py

@@ -1,410 +0,0 @@
-"""Rivet Memory — persistent, per-user memory over SQLite.
-
-Same shape as Ember's semantic memory (flat records + relevance recall)
-but scoped for a code assistant and dependency-free: instead of an
-embedding model, recall scores stored memories against the query with
-TF-IDF-style keyword overlap, weighted by recency and how often a memory
-has recurred. Good enough for "what has this user hit before" — not a
-vector store.
-
-Rivet decides what to remember (calls remember_*). Memory decides what
-to surface (recall / session_brief ranks and returns it). Memory never
-inspects code or talks to the model — it only stores and ranks text
-Rivet hands it.
-
-One SQLite file per deployment. Each user gets their own table
-(memories_<user_id>) so one user's history never leaks into another's
-recall results.
-"""
-
-import argparse
-import hashlib
-import math
-import os
-import re
-import sqlite3
-from collections import Counter
-from dataclasses import dataclass
-from datetime import datetime, timezone
-from pathlib import Path
-
-
-DEFAULT_DB_PATH = Path(
-    os.environ.get("RIVET_MEMORY_DB")
-    or (Path(__file__).resolve().parent.parent / "data" / "rivet_memory.db")
-)
-
-KIND_FILE = "file"                # a file the user works on
-KIND_QA = "qa"                    # a question asked + the answer given
-KIND_PATTERN = "pattern"          # a recurring behavior (e.g. keeps proposing destructive migrations)
-KIND_GATE = "gate_correction"     # a discipline-gate flag that fired on this user's work
-KIND_CONTEXT = "context"          # single-value facts: current_branch, recent_pr, ...
-
-_STOPWORDS = {
-    "the", "a", "an", "is", "are", "was", "were", "to", "of", "in", "on",
-    "for", "and", "or", "it", "this", "that", "with", "at", "as", "be",
-    "by", "from", "how", "what", "why", "do", "does", "did", "i", "you",
-    "we", "my", "me", "can", "will", "should", "please",
-}
-
-# Vocabulary from DESIGN.md's architecture map — matched against file
-# paths so file memories carry cross-cutting tags ("auth", "migration")
-# even when the path itself doesn't spell them out.
-_DOMAIN_TAGS = {
-    "auth": ("auth", "jwt", "session", "middleware"),
-    "migration": ("migration", "migrations", "schema"),
-    "webhook": ("webhook", "webhooks", "stripe"),
-    "socket": ("socket", "sockets", "realtime", "pubsub"),
-    "store": ("store", "stores", "zustand"),
-    "route": ("route", "routes", "api"),
-    "job": ("job", "jobs", "pg-boss", "queue"),
-}
-
-
-def _now() -> str:
-    return datetime.now(timezone.utc).isoformat()
-
-
-def _sanitize_user_id(user_id: str) -> str:
-    slug = re.sub(r"[^a-zA-Z0-9_]", "_", user_id.strip())
-    if not slug:
-        raise ValueError("user_id must contain at least one alphanumeric character")
-    return slug
-
-
-def _tokenize(text: str) -> list:
-    words = re.findall(r"[a-z0-9_./-]+", text.lower())
-    return [w for w in words if w not in _STOPWORDS and len(w) > 1]
-
-
-def _content_key(text: str) -> str:
-    return hashlib.sha1(text.strip().lower().encode("utf-8")).hexdigest()[:16]
-
-
-def _path_tags(file_path: str) -> list:
-    lowered = file_path.lower()
-    tags = set()
-    for tag, needles in _DOMAIN_TAGS.items():
-        if any(n in lowered for n in needles):
-            tags.add(tag)
-    return sorted(tags)
-
-
-def _age_days(iso_timestamp: str, now: datetime) -> float:
-    try:
-        ts = datetime.fromisoformat(iso_timestamp)
-    except ValueError:
-        return 0.0
-    if ts.tzinfo is None:
-        ts = ts.replace(tzinfo=timezone.utc)
-    return max(0.0, (now - ts).total_seconds() / 86400.0)
-
-
-@dataclass
-class MemoryEntry:
-    id: int
-    kind: str
-    key: str
-    content: str
-    tags: str
-    weight: float
-    hit_count: int
-    created_at: str
-    updated_at: str
-    score: float = 0.0
-
-    @property
-    def tag_list(self) -> list:
-        return [t for t in self.tags.split(",") if t] if self.tags else []
-
-
-def _to_entry(row: sqlite3.Row, score: float = 0.0) -> MemoryEntry:
-    return MemoryEntry(
-        id=row["id"], kind=row["kind"], key=row["key"], content=row["content"],
-        tags=row["tags"] or "", weight=row["weight"], hit_count=row["hit_count"],
-        created_at=row["created_at"], updated_at=row["updated_at"], score=score,
-    )
-
-
-class MemoryStore:
-    """Per-user persistent memory backed by SQLite. One table per user."""
-
-    def __init__(self, db_path=DEFAULT_DB_PATH):
-        self.db_path = Path(db_path)
-        self.db_path.parent.mkdir(parents=True, exist_ok=True)
-        self._conn = sqlite3.connect(str(self.db_path))
-        self._conn.row_factory = sqlite3.Row
-        self._conn.execute("PRAGMA journal_mode=WAL")
-        self._conn.execute("""
-            CREATE TABLE IF NOT EXISTS users (
-                user_id TEXT PRIMARY KEY,
-                table_name TEXT NOT NULL,
-                first_seen TEXT NOT NULL,
-                last_seen TEXT NOT NULL
-            )
-        """)
-        self._conn.commit()
-
-    def close(self):
-        self._conn.close()
-
-    def __enter__(self):
-        return self
-
-    def __exit__(self, *exc):
-        self.close()
-
-    # -- per-user table management ------------------------------------
-
-    def _table_for(self, user_id: str) -> str:
-        table = f"memories_{_sanitize_user_id(user_id)}"
-        now = _now()
-        row = self._conn.execute(
-            "SELECT table_name FROM users WHERE user_id = ?", (user_id,)
-        ).fetchone()
-        if row is None:
-            self._conn.execute(
-                "INSERT INTO users (user_id, table_name, first_seen, last_seen) VALUES (?, ?, ?, ?)",
-                (user_id, table, now, now),
-            )
-            self._conn.execute(f"""
-                CREATE TABLE IF NOT EXISTS "{table}" (
-                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    kind TEXT NOT NULL,
-                    key TEXT NOT NULL,
-                    content TEXT NOT NULL,
-                    tags TEXT DEFAULT '',
-                    weight REAL NOT NULL DEFAULT 1.0,
-                    hit_count INTEGER NOT NULL DEFAULT 1,
-                    created_at TEXT NOT NULL,
-                    updated_at TEXT NOT NULL,
-                    UNIQUE(kind, key)
-                )
-            """)
-            self._conn.commit()
-        else:
-            self._conn.execute("UPDATE users SET last_seen = ? WHERE user_id = ?", (now, user_id))
-            self._conn.commit()
-            table = row["table_name"]
-        return table
-
-    # -- writing (Rivet decides what to remember) ----------------------
-
-    def remember(self, user_id: str, kind: str, content: str, key: str = None,
-                 tags=None, weight: float = 1.0) -> None:
-        """Store or reinforce a memory. Same (kind, key) upserts: content is
-        replaced with the latest version, weight and hit_count accumulate so
-        recurring evidence outranks one-off mentions."""
-        table = self._table_for(user_id)
-        key = key or _content_key(content)
-        tag_str = ",".join(sorted(set(tags))) if tags else ""
-        now = _now()
-        self._conn.execute(f"""
-            INSERT INTO "{table}" (kind, key, content, tags, weight, hit_count, created_at, updated_at)
-            VALUES (?, ?, ?, ?, ?, 1, ?, ?)
-            ON CONFLICT(kind, key) DO UPDATE SET
-                content = excluded.content,
-                tags = CASE WHEN excluded.tags != '' THEN excluded.tags ELSE tags END,
-                weight = weight + excluded.weight,
-                hit_count = hit_count + 1,
-                updated_at = excluded.updated_at
-        """, (kind, key, content, tag_str, weight, now, now))
-        self._conn.commit()
-
-    def remember_file(self, user_id: str, file_path: str, note: str = "") -> None:
-        """A file the user touched, optionally with what they were doing there."""
-        self.remember(user_id, KIND_FILE, note or file_path, key=file_path,
-                      tags=_path_tags(file_path), weight=1.0)
-
-    def remember_qa(self, user_id: str, intent: str, outcome: str, tags=None) -> None:
-        """A question + answer, compressed to intent/outcome (not full transcript)."""
-        content = f"Q: {intent.strip()}\nA: {outcome.strip()}"
-        self.remember(user_id, KIND_QA, content, key=_content_key(intent), tags=tags, weight=1.0)
-
-    def remember_pattern(self, user_id: str, pattern: str, detail: str = "", tags=None) -> None:
-        """A recurring behavior worth flagging preemptively next time it shows up."""
-        self.remember(user_id, KIND_PATTERN, detail or pattern, key=pattern,
-                      tags=tags, weight=2.0)
-
-    def remember_gate_correction(self, user_id: str, rule: str, message: str,
-                                  severity: str = "WARN") -> None:
-        """A discipline-gate flag that fired against this user's suggestion."""
-        self.remember(user_id, KIND_GATE, message, key=f"{severity}:{rule}",
-                      tags=[rule.lower(), severity.lower()], weight=1.5)
-
-    def set_context(self, user_id: str, key: str, value: str) -> None:
-        """Single-value session facts: current_branch, recent_pr, etc."""
-        self.remember(user_id, KIND_CONTEXT, str(value), key=key, tags=[key], weight=1.0)
-
-    def get_context(self, user_id: str, key: str, default=None):
-        table = self._table_for(user_id)
-        row = self._conn.execute(
-            f'SELECT content FROM "{table}" WHERE kind = ? AND key = ?', (KIND_CONTEXT, key)
-        ).fetchone()
-        return row["content"] if row else default
-
-    # -- reading (memory decides what to surface) -----------------------
-
-    def recall(self, user_id: str, query: str, top_n: int = 8, kinds=None) -> list:
-        """Rank stored memories against a free-text query.
-
-        Scoring is TF-IDF-style keyword overlap (computed on the fly over
-        this user's memories — no external embedding model), boosted by
-        recency and by how often the memory has recurred.
-        """
-        table = self._table_for(user_id)
-        sql = f'SELECT * FROM "{table}"'
-        params = []
-        if kinds:
-            sql += f" WHERE kind IN ({','.join('?' for _ in kinds)})"
-            params.extend(kinds)
-        rows = self._conn.execute(sql, params).fetchall()
-        if not rows:
-            return []
-
-        query_terms = _tokenize(query)
-        if not query_terms:
-            return self._rank_by_importance(rows, top_n)
-
-        doc_freq = Counter()
-        doc_terms = []
-        for row in rows:
-            terms = Counter(_tokenize(f"{row['key']} {row['content']} {row['tags']}"))
-            doc_terms.append(terms)
-            doc_freq.update(terms.keys())
-
-        n_docs = len(rows)
-        now = datetime.now(timezone.utc)
-        scored = []
-        for row, terms in zip(rows, doc_terms):
-            overlap = sum(
-                terms[qt] * math.log(1 + n_docs / doc_freq[qt])
-                for qt in query_terms if terms.get(qt)
-            )
-            if overlap <= 0:
-                continue
-            recency = 1.0 / (1.0 + _age_days(row["updated_at"], now) / 14.0)
-            importance = math.log(1 + row["hit_count"]) * row["weight"]
-            score = overlap * (0.6 + 0.4 * recency) * (1.0 + 0.3 * importance)
-            scored.append(_to_entry(row, score))
-
-        scored.sort(key=lambda e: e.score, reverse=True)
-        return scored[:top_n]
-
-    def _rank_by_importance(self, rows, top_n: int) -> list:
-        now = datetime.now(timezone.utc)
-        scored = []
-        for row in rows:
-            recency = 1.0 / (1.0 + _age_days(row["updated_at"], now) / 14.0)
-            importance = math.log(1 + row["hit_count"]) * row["weight"]
-            scored.append(_to_entry(row, importance * recency))
-        scored.sort(key=lambda e: e.score, reverse=True)
-        return scored[:top_n]
-
-    def top_files(self, user_id: str, n: int = 5) -> list:
-        table = self._table_for(user_id)
-        rows = self._conn.execute(
-            f'SELECT * FROM "{table}" WHERE kind = ? ORDER BY hit_count DESC, updated_at DESC LIMIT ?',
-            (KIND_FILE, n),
-        ).fetchall()
-        return [_to_entry(r) for r in rows]
-
-    def recent_patterns(self, user_id: str, n: int = 5) -> list:
-        table = self._table_for(user_id)
-        rows = self._conn.execute(
-            f'SELECT * FROM "{table}" WHERE kind = ? ORDER BY weight DESC, updated_at DESC LIMIT ?',
-            (KIND_PATTERN, n),
-        ).fetchall()
-        return [_to_entry(r) for r in rows]
-
-    def recent_gate_corrections(self, user_id: str, n: int = 5) -> list:
-        table = self._table_for(user_id)
-        rows = self._conn.execute(
-            f'SELECT * FROM "{table}" WHERE kind = ? ORDER BY updated_at DESC LIMIT ?',
-            (KIND_GATE, n),
-        ).fetchall()
-        return [_to_entry(r) for r in rows]
-
-    def session_brief(self, user_id: str, first_message: str = "", top_n: int = 8) -> dict:
-        """Everything Rivet should know when a session starts: memories
-        relevant to the first message, plus standing context that should
-        always be visible (frequent files, recurring patterns, gate history,
-        branch/PR state) regardless of what was asked."""
-        return {
-            "relevant": self.recall(user_id, first_message, top_n=top_n) if first_message else [],
-            "top_files": self.top_files(user_id),
-            "known_patterns": self.recent_patterns(user_id),
-            "gate_history": self.recent_gate_corrections(user_id),
-            "current_branch": self.get_context(user_id, "current_branch"),
-            "recent_pr": self.get_context(user_id, "recent_pr"),
-        }
-
-    def to_system_block(self, user_id: str, first_message: str = "") -> str:
-        """Format session_brief as a text block ready to inject into the
-        model's system context, in the same spirit as RivetContext in
-        context_loader.py."""
-        brief = self.session_brief(user_id, first_message)
-        parts = ["# USER MEMORY\n"]
-
-        if brief["current_branch"] or brief["recent_pr"]:
-            parts.append("## Current Work")
-            if brief["current_branch"]:
-                parts.append(f"- Branch: {brief['current_branch']}")
-            if brief["recent_pr"]:
-                parts.append(f"- Recent PR: {brief['recent_pr']}")
-            parts.append("")
-
-        if brief["top_files"]:
-            parts.append("## Frequently Touched Files")
-            parts.extend(f"- {e.key}: {e.content}" for e in brief["top_files"])
-            parts.append("")
-
-        if brief["known_patterns"]:
-            parts.append("## Recurring Patterns (flag preemptively)")
-            parts.extend(f"- {e.key}: {e.content}" for e in brief["known_patterns"])
-            parts.append("")
-
-        if brief["gate_history"]:
-            parts.append("## Discipline Gate History")
-            parts.extend(f"- {e.content}" for e in brief["gate_history"])
-            parts.append("")
-
-        if brief["relevant"]:
-            parts.append("## Relevant To This Question")
-            parts.extend(f"- {e.content}" for e in brief["relevant"])
-            parts.append("")
-
-        return "\n".join(parts).strip()
-
-
-def _cli():
-    parser = argparse.ArgumentParser(description="Rivet memory database utility")
-    parser.add_argument("command", choices=["init", "demo"])
-    parser.add_argument("--db", default=str(DEFAULT_DB_PATH))
-    args = parser.parse_args()
-
-    if args.command == "init":
-        store = MemoryStore(args.db)
-        print(f"Memory database ready at {store.db_path}")
-        store.close()
-        return
-
-    # demo: quick smoke test exercising the write/recall paths
-    store = MemoryStore(args.db)
-    user = "demo_user"
-    store.remember_file(user, "src/routes/auth.ts", "adding refresh-token rotation")
-    store.remember_qa(user, "why does the webhook route 500 on missing signature",
-                       "STRIPE_SECRET falls back to '' — Audit C2, reject instead of skip")
-    store.remember_pattern(user, "proposes_destructive_migration",
-                            "suggested DROP COLUMN twice — always steer to additive migrations")
-    store.remember_gate_correction(user, "webhook_bypass",
-                                    "Webhook signature bypass pattern (Audit C2) flagged and blocked",
-                                    severity="BLOCK")
-    store.set_context(user, "current_branch", "feat/refresh-token-rotation")
-
-    print(store.to_system_block(user, first_message="can I touch the webhook signature check?"))
-    store.close()
-
-
-if __name__ == "__main__":
-    _cli()