imprint-memory 0.1.0__py3-none-any.whl

imprint_memory/__init__.py

@@ -0,0 +1 @@
+# imprint-memory: Persistent memory system for Claude Code

imprint_memory/bus.py

@@ -0,0 +1,53 @@
+"""
+Message bus.
+Shared log for messages sent/received across different sources.
+"""
+
+import os
+from .db import _get_db, now_str
+
+MESSAGE_BUS_LIMIT = int(os.environ.get("MESSAGE_BUS_LIMIT", 40))
+
+
+def bus_post(source: str, direction: str, content: str) -> None:
+    """Write a message to the bus. Auto-prunes old messages beyond limit.
+    source: free-form label (e.g. cc, chat, api, webhook)
+    direction: in (received) / out (sent)
+    content: message content (auto-truncated to 200 chars)"""
+    if len(content) > 200:
+        content = content[:197] + "..."
+
+    db = _get_db()
+    db.execute(
+        "INSERT INTO message_bus (source, direction, content, created_at) VALUES (?, ?, ?, ?)",
+        (source, direction, content, now_str()),
+    )
+    db.execute(
+        "DELETE FROM message_bus WHERE id NOT IN (SELECT id FROM message_bus ORDER BY id DESC LIMIT ?)",
+        (MESSAGE_BUS_LIMIT,),
+    )
+    db.commit()
+    db.close()
+
+
+def bus_read(limit: int = 20) -> list[dict]:
+    """Read recent bus messages."""
+    db = _get_db()
+    rows = db.execute(
+        "SELECT source, direction, content, created_at FROM message_bus ORDER BY id DESC LIMIT ?",
+        (limit,),
+    ).fetchall()
+    db.close()
+    return [dict(r) for r in reversed(rows)]
+
+
+def bus_format(limit: int = 20) -> str:
+    """Format bus messages for context injection."""
+    messages = bus_read(limit)
+    if not messages:
+        return "(No recent messages)"
+    lines = ["# Recent Messages\n"]
+    for m in messages:
+        arrow = "\u2192" if m["direction"] == "out" else "\u2190"
+        lines.append(f"[{m['created_at']}] [{m['source']}] {arrow} {m['content']}")
+    return "\n".join(lines)

imprint_memory/chat_sync_receiver.py

@@ -0,0 +1,390 @@
+"""
+Chat Sync Receiver — HTTP endpoint that ingests conversations from the
+imprint-chat-sync browser extension (or any client posting to /api/ingest).
+
+Run with:
+    imprint-memory-receiver               # listens on 127.0.0.1:8001
+    PORT=9001 imprint-memory-receiver     # custom port
+
+Pipeline per ingest call:
+    POST /api/ingest                  (browser extension)
+      ↓
+    log_message()                     (writes to conversation_log)
+      ↓ background task
+    embed_new_messages()              (Gemini Embedding 2 per message)
+      ↓
+    detect_topic_shifts()             (adjacent user-msg cosine → topic edges)
+      ↓
+    incremental_chunk_update()        (chunk → summarize → embed → top-K edges)
+
+The receiver shares the same SQLite database as the main imprint-memory server
+(configured via IMPRINT_DATA_DIR / IMPRINT_DB).
+"""
+
+import math
+import os
+import re
+import struct
+import threading
+import time
+
+from starlette.applications import Starlette
+from starlette.background import BackgroundTask
+from starlette.middleware.cors import CORSMiddleware
+from starlette.responses import JSONResponse
+from starlette.routing import Route
+
+from .conversation import log_message, get_recent
+from .db import _get_db as _get_app_db
+
+DEFAULT_PORT = 8001
+THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)
+EMBED_DELAY = float(os.environ.get("IMPRINT_RECEIVER_EMBED_DELAY", "0.7"))
+SHIFT_THRESHOLD = float(os.environ.get("IMPRINT_RECEIVER_SHIFT_THRESHOLD", "0.50"))
+RECEIVER_HOST_ENV = os.environ.get("IMPRINT_RECEIVER_HOST", os.environ.get("HOST", "127.0.0.1"))
+RECEIVER_PORT_ENV = int(os.environ.get("IMPRINT_RECEIVER_PORT", os.environ.get("PORT", DEFAULT_PORT)))
+RECEIVER_CORS_ORIGIN_REGEX = os.environ.get(
+    "IMPRINT_RECEIVER_CORS_ORIGIN_REGEX",
+    r"^chrome-extension://.*$",
+)
+
+
+def _blob_to_vec(blob):
+    """Decode a float32 embedding blob."""
+    return list(struct.unpack(f"{len(blob)//4}f", blob))
+
+
+def _cosine_sim(a, b):
+    """Return cosine similarity for same-length vectors."""
+    if len(a) != len(b):
+        return 0.0
+    dot = sum(x * y for x, y in zip(a, b))
+    na = math.sqrt(sum(x * x for x in a))
+    nb = math.sqrt(sum(x * x for x in b))
+    return dot / (na * nb) if na and nb else 0.0
+
+
+def _get_db():
+    """Open the receiver database connection and ensure receiver-owned tables."""
+    db = _get_app_db()
+    db.execute("""CREATE TABLE IF NOT EXISTS conversation_vectors (
+        msg_id INTEGER PRIMARY KEY,
+        embedding BLOB NOT NULL,
+        model TEXT DEFAULT 'gemini-embedding-2'
+    )""")
+    db.execute("""CREATE TABLE IF NOT EXISTS conversation_edges (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        msg_before INTEGER NOT NULL,
+        msg_after INTEGER NOT NULL,
+        session_id TEXT DEFAULT '',
+        similarity REAL NOT NULL,
+        strength REAL DEFAULT 1.0,
+        surfaced_count INTEGER DEFAULT 0,
+        last_surfaced_at TEXT,
+        status TEXT DEFAULT 'active',
+        created_at TEXT NOT NULL,
+        UNIQUE(msg_before, msg_after)
+    )""")
+    return db
+
+
+def embed_new_messages(msg_ids):
+    """Background: embed newly ingested messages with the configured embedder."""
+    if not msg_ids:
+        return
+    try:
+        from .memory_manager import EMBED_MODEL, _embed
+    except Exception:
+        return
+
+    db = _get_db()
+    try:
+        for msg_id in msg_ids:
+            existing = db.execute(
+                "SELECT 1 FROM conversation_vectors WHERE msg_id=?", (msg_id,)
+            ).fetchone()
+            if existing:
+                continue
+
+            row = db.execute(
+                "SELECT content, direction FROM conversation_log WHERE id=?", (msg_id,)
+            ).fetchone()
+            if not row:
+                continue
+
+            content = row["content"]
+            if row["direction"] == "out":
+                content = THINK_RE.sub("", content).strip()
+            if len(content) < 10:
+                continue
+
+            try:
+                vec = _embed(content[:2000])
+                if vec:
+                    db.execute(
+                        "INSERT OR IGNORE INTO conversation_vectors (msg_id, embedding, model) VALUES (?, ?, ?)",
+                        (msg_id, struct.pack(f"{len(vec)}f", *vec), EMBED_MODEL),
+                    )
+                    db.commit()
+            except Exception:
+                pass
+
+            time.sleep(EMBED_DELAY)
+
+        detect_topic_shifts(db, msg_ids)
+    except Exception:
+        pass
+    finally:
+        db.close()
+
+    # Incremental chunk processing: split → summarize → embed → Top-K edges
+    try:
+        from .conversation_chunker import incremental_chunk_update
+        result = incremental_chunk_update(batch_size=200)
+        if result.get("chunks_created", 0) > 0:
+            print(
+                f"[chunk] Created {result['chunks_created']} chunks, "
+                f"{result.get('edges_created', 0)} edges",
+                flush=True,
+            )
+    except Exception as e:
+        print(f"[chunk] Error: {e}", flush=True)
+
+
+def detect_topic_shifts(db, msg_ids):
+    """Detect topic shifts by cosine similarity between adjacent user messages."""
+    if not msg_ids:
+        return
+
+    placeholders = ",".join("?" * len(msg_ids))
+    sessions = db.execute(
+        f"SELECT DISTINCT session_id FROM conversation_log WHERE id IN ({placeholders})",
+        msg_ids,
+    ).fetchall()
+
+    shifts = 0
+    for (session_id,) in sessions:
+        if not session_id:
+            continue
+
+        rows = db.execute(
+            """SELECT c.id, v.embedding
+               FROM conversation_log c
+               JOIN conversation_vectors v ON c.id = v.msg_id
+               WHERE c.session_id = ? AND c.direction = 'in'
+               ORDER BY c.id""",
+            (session_id,),
+        ).fetchall()
+
+        if len(rows) < 2:
+            continue
+
+        for i in range(1, len(rows)):
+            vec_a = _blob_to_vec(rows[i - 1]["embedding"])
+            vec_b = _blob_to_vec(rows[i]["embedding"])
+            sim = _cosine_sim(vec_a, vec_b)
+
+            if sim < SHIFT_THRESHOLD:
+                try:
+                    db.execute(
+                        """INSERT OR IGNORE INTO conversation_edges
+                           (msg_before, msg_after, session_id, similarity, created_at)
+                           VALUES (?, ?, ?, ?, datetime('now'))""",
+                        (rows[i - 1]["id"], rows[i]["id"], session_id, sim),
+                    )
+                    shifts += 1
+                except Exception:
+                    pass
+
+    if shifts:
+        db.commit()
+
+
+async def ingest(request):
+    """Ingest a batch of conversation messages from the browser extension."""
+    try:
+        data = await request.json()
+    except Exception:
+        return JSONResponse({"ok": False, "error": "invalid JSON"}, status_code=400)
+
+    conversation_id = data.get("conversation_id", "")
+    messages = data.get("messages", [])
+
+    if not messages:
+        return JSONResponse({"ok": True, "ingested": 0, "skipped": 0, "errors": 0})
+
+    results = {"ingested": 0, "skipped": 0, "errors": 0}
+    ingested_ids = []
+
+    for msg in messages:
+        content = (msg.get("content") or "").strip()
+        if not content:
+            results["errors"] += 1
+            continue
+
+        result = log_message(
+            platform=data.get("platform", "claude.ai"),
+            direction=msg.get("direction", "in"),
+            content=content,
+            speaker=msg.get("speaker", ""),
+            session_id=conversation_id,
+            entrypoint="browser_extension",
+            created_at=msg.get("created_at", ""),
+            summary=msg.get("summary", ""),
+            model=msg.get("model", data.get("model", "")),
+            external_id=msg.get("external_id") or msg.get("uuid", ""),
+        )
+
+        if result.get("skipped"):
+            results["skipped"] += 1
+        elif result.get("ok"):
+            results["ingested"] += 1
+            if result.get("id"):
+                ingested_ids.append(result["id"])
+        else:
+            results["errors"] += 1
+
+    return JSONResponse(
+        {"ok": True, **results},
+        background=BackgroundTask(embed_new_messages, ingested_ids),
+    )
+
+
+async def health(request):
+    """Return a lightweight liveness response."""
+    return JSONResponse({"ok": True, "service": "imprint-chat-sync-receiver"})
+
+
+async def status(request):
+    """Return recent ingest and embedding status for the extension popup."""
+    recent = get_recent(platform="claude.ai", limit=5)
+    db = _get_db()
+    try:
+        vec_count = db.execute("SELECT count(*) FROM conversation_vectors").fetchone()[0]
+    finally:
+        db.close()
+    return JSONResponse({
+        "ok": True,
+        "recent_count": len(recent),
+        "last_message": recent[-1]["created_at"] if recent else None,
+        "vectors": vec_count,
+    })
+
+
+app = Starlette(
+    routes=[
+        Route("/api/ingest", ingest, methods=["POST"]),
+        Route("/api/health", health, methods=["GET"]),
+        Route("/api/status", status, methods=["GET"]),
+    ],
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origin_regex=RECEIVER_CORS_ORIGIN_REGEX,
+    allow_methods=["GET", "POST", "OPTIONS"],
+    allow_headers=["Content-Type"],
+)
+
+
+def _backfill_on_startup():
+    """Background-embed any messages that don't have a vector yet."""
+
+    def _run():
+        db = _get_db()
+        try:
+            unembedded = db.execute("""
+                SELECT COUNT(*) as c FROM conversation_log cl
+                LEFT JOIN conversation_vectors cv ON cl.id = cv.msg_id
+                WHERE cv.msg_id IS NULL AND cl.platform NOT IN ('cc')
+                AND cl.content IS NOT NULL AND length(cl.content) > 0
+            """).fetchone()[0]
+        finally:
+            db.close()
+
+        if unembedded == 0:
+            print("[backfill] No unembedded messages, skipping", flush=True)
+            return
+
+        print(
+            f"[backfill] Found {unembedded} unembedded messages, processing...",
+            flush=True,
+        )
+
+        db = _get_db()
+        try:
+            rows = db.execute("""
+                SELECT cl.id FROM conversation_log cl
+                LEFT JOIN conversation_vectors cv ON cl.id = cv.msg_id
+                WHERE cv.msg_id IS NULL AND cl.platform NOT IN ('cc')
+                AND cl.content IS NOT NULL AND length(cl.content) > 0
+                ORDER BY cl.id
+            """).fetchall()
+            msg_ids = [r[0] for r in rows]
+        finally:
+            db.close()
+
+        batch_size = 100
+        for i in range(0, len(msg_ids), batch_size):
+            batch = msg_ids[i:i + batch_size]
+            embed_new_messages(batch)
+            print(
+                f"[backfill] Embedded {min(i + batch_size, len(msg_ids))}/{len(msg_ids)}",
+                flush=True,
+            )
+
+        print("[backfill] Embedding complete. Running incremental chunk update...", flush=True)
+        try:
+            from .conversation_chunker import incremental_chunk_update
+            result = incremental_chunk_update(batch_size=200)
+            print(f"[backfill] Chunk update: {result}", flush=True)
+        except Exception as e:
+            print(f"[backfill] Chunk update error: {e}", flush=True)
+
+    t = threading.Thread(target=_run, daemon=True)
+    t.start()
+
+
+def main():
+    """Run the chat-sync receiver HTTP service."""
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Imprint chat sync receiver")
+    parser.add_argument(
+        "--host",
+        default=RECEIVER_HOST_ENV,
+        help="Bind address (default: 127.0.0.1)",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=RECEIVER_PORT_ENV,
+        help=f"Port (default: {DEFAULT_PORT})",
+    )
+    parser.add_argument(
+        "--no-backfill",
+        action="store_true",
+        help="Skip the startup backfill pass over unembedded messages",
+    )
+    args = parser.parse_args()
+
+    try:
+        import uvicorn
+    except ImportError:
+        print(
+            "uvicorn is required. Install with: pip install 'imprint-memory[receiver]'",
+            flush=True,
+        )
+        raise SystemExit(1)
+
+    print(
+        f"imprint-memory chat-sync receiver listening on {args.host}:{args.port}",
+        flush=True,
+    )
+    if not args.no_backfill:
+        _backfill_on_startup()
+    uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+
+
+if __name__ == "__main__":
+    main()

imprint_memory/compress.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""
+Compress old messages in a context file using a local Ollama model.
+
+Can be called standalone or imported as a library.
+Designed to keep a rolling context file (e.g. recent_context.md) from growing
+unbounded by summarizing older messages with a local LLM.
+
+Usage:
+    python3 -m imprint_memory.compress /path/to/recent_context.md
+
+Environment variables:
+    OLLAMA_URL       — Ollama API endpoint (default: http://localhost:11434)
+    COMPRESS_MODEL   — Model to use (default: qwen3:8b)
+    COMPRESS_KEEP    — Number of recent lines to keep as-is (default: 30)
+    COMPRESS_THRESHOLD — Line count that triggers compression (default: 50)
+"""
+
+import json
+import os
+import sys
+import urllib.request
+from pathlib import Path
+
+OLLAMA_URL = os.environ.get("OLLAMA_URL", "http://localhost:11434")
+COMPRESS_MODEL = os.environ.get("COMPRESS_MODEL", "qwen3:8b")
+KEEP_RECENT = int(os.environ.get("COMPRESS_KEEP", 30))
+THRESHOLD = int(os.environ.get("COMPRESS_THRESHOLD", 50))
+
+SYSTEM_PROMPT = (
+    "You are a log compressor. "
+    "Compress the following log into 3-5 summary lines. "
+    "Preserve all topics faithfully. "
+    "Keep names, timestamps, decisions, and context. "
+    "Output ONLY the summary lines, nothing else."
+)
+
+
+def compress_messages(messages: list[str], model: str = "", ollama_url: str = "") -> str | None:
+    """Call local Ollama to summarize messages into 3-5 lines.
+
+    Returns the summary string, or None if Ollama is unavailable.
+    """
+    text = "\n".join(messages)
+    url = ollama_url or OLLAMA_URL
+    mdl = model or COMPRESS_MODEL
+
+    try:
+        req = urllib.request.Request(
+            f"{url}/api/chat",
+            data=json.dumps({
+                "model": mdl,
+                "messages": [
+                    {"role": "system", "content": SYSTEM_PROMPT},
+                    {"role": "user", "content": f"Compress this log:\n\n{text}"},
+                ],
+                "stream": False,
+                "think": False,
+                "options": {"temperature": 0.3, "num_predict": 500},
+            }).encode(),
+            headers={"Content-Type": "application/json"},
+        )
+        resp = json.loads(urllib.request.urlopen(req, timeout=60).read())
+        result = resp.get("message", {}).get("content", "").strip()
+        lines = [l.strip() for l in result.splitlines() if l.strip()]
+        return "\n".join(lines) if lines else None
+    except Exception as e:
+        print(f"Ollama compression failed: {e}", file=sys.stderr)
+        return None
+
+
+def compress_file(context_file: Path, keep: int = 0, threshold: int = 0) -> bool:
+    """Compress a context file in-place. Returns True if compression happened."""
+    keep = keep or KEEP_RECENT
+    threshold = threshold or THRESHOLD
+
+    if not context_file.exists():
+        return False
+
+    content = context_file.read_text(encoding="utf-8")
+    lines = content.splitlines()
+
+    # Separate header comments from message lines
+    header_lines = []
+    message_lines = []
+    for line in lines:
+        if line.startswith("<!--") or not line.strip():
+            header_lines.append(line)
+        else:
+            message_lines.append(line)
+
+    if len(message_lines) <= threshold:
+        print(f"Only {len(message_lines)} messages, below threshold {threshold}", file=sys.stderr)
+        return False
+
+    old_messages = message_lines[:-keep]
+    recent_messages = message_lines[-keep:]
+
+    print(f"Compressing {len(old_messages)} old messages, keeping {len(recent_messages)} recent", file=sys.stderr)
+
+    summary = compress_messages(old_messages)
+    if summary is None:
+        print("Compression failed, truncating to recent messages only", file=sys.stderr)
+        new_content = "\n".join(header_lines + [""] + recent_messages) + "\n"
+    else:
+        new_content = "\n".join(header_lines + [""] + [summary, ""] + recent_messages) + "\n"
+
+    # Atomic write
+    tmp = context_file.with_suffix(".tmp")
+    tmp.write_text(new_content, encoding="utf-8")
+    os.replace(str(tmp), str(context_file))
+
+    n = len(summary.splitlines()) if summary else 0
+    print(f"Compressed: {len(old_messages)} -> {n} lines", file=sys.stderr)
+    return True
+
+
+def main():
+    """CLI entry point for context-file compression."""
+    if len(sys.argv) < 2:
+        print("Usage: python3 -m imprint_memory.compress <context_file>", file=sys.stderr)
+        sys.exit(1)
+    compress_file(Path(sys.argv[1]))
+
+
+if __name__ == "__main__":
+    main()