headroom-ai 0.2.13__py3-none-any.whl

headroom/memory/store.py

@@ -0,0 +1,434 @@
+"""SQLite + FTS5 memory storage for Headroom Memory.
+
+Simple, fast, local-first storage with full-text search.
+No external dependencies - just SQLite (built into Python).
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Literal
+
+
+@dataclass
+class Memory:
+    """A single memory entry."""
+
+    content: str
+    category: Literal["preference", "fact", "context"] = "fact"
+    importance: float = 0.5
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    created_at: datetime = field(default_factory=datetime.utcnow)
+    metadata: dict = field(default_factory=dict)
+
+
+@dataclass
+class PendingExtraction:
+    """A conversation pending memory extraction."""
+
+    user_id: str
+    query: str
+    response: str
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    created_at: datetime = field(default_factory=datetime.utcnow)
+    status: Literal["pending", "processing", "done", "failed"] = "pending"
+
+
+class SQLiteMemoryStore:
+    """SQLite + FTS5 storage for memories.
+
+    Features:
+    - Full-text search via FTS5
+    - User isolation (each user_id has separate memories)
+    - Pending extractions for crash recovery
+    - Thread-safe with connection per call
+
+    Usage:
+        store = SQLiteMemoryStore("./memory.db")
+        store.save("alice", Memory(content="Prefers Python"))
+        results = store.search("alice", "python")
+    """
+
+    def __init__(self, db_path: str | Path = "headroom_memory.db"):
+        """Initialize the store.
+
+        Args:
+            db_path: Path to SQLite database file. Created if doesn't exist.
+        """
+        self.db_path = Path(db_path)
+        self._init_db()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        """Get a new connection (thread-safe pattern)."""
+        conn = sqlite3.connect(str(self.db_path))
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def _init_db(self) -> None:
+        """Initialize database schema."""
+        with self._get_conn() as conn:
+            # Main memories table
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS memories (
+                    id TEXT PRIMARY KEY,
+                    user_id TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    category TEXT NOT NULL DEFAULT 'fact',
+                    importance REAL NOT NULL DEFAULT 0.5,
+                    created_at TEXT NOT NULL,
+                    metadata TEXT NOT NULL DEFAULT '{}'
+                )
+            """)
+
+            # FTS5 virtual table for full-text search
+            conn.execute("""
+                CREATE VIRTUAL TABLE IF NOT EXISTS memories_fts USING fts5(
+                    content,
+                    content='memories',
+                    content_rowid='rowid'
+                )
+            """)
+
+            # Triggers to keep FTS in sync
+            conn.execute("""
+                CREATE TRIGGER IF NOT EXISTS memories_ai AFTER INSERT ON memories BEGIN
+                    INSERT INTO memories_fts(rowid, content)
+                    VALUES (new.rowid, new.content);
+                END
+            """)
+
+            conn.execute("""
+                CREATE TRIGGER IF NOT EXISTS memories_ad AFTER DELETE ON memories BEGIN
+                    INSERT INTO memories_fts(memories_fts, rowid, content)
+                    VALUES ('delete', old.rowid, old.content);
+                END
+            """)
+
+            conn.execute("""
+                CREATE TRIGGER IF NOT EXISTS memories_au AFTER UPDATE ON memories BEGIN
+                    INSERT INTO memories_fts(memories_fts, rowid, content)
+                    VALUES ('delete', old.rowid, old.content);
+                    INSERT INTO memories_fts(rowid, content)
+                    VALUES (new.rowid, new.content);
+                END
+            """)
+
+            # Index for user_id filtering
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_memories_user_id
+                ON memories(user_id)
+            """)
+
+            # Pending extractions table (for crash recovery)
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS pending_extractions (
+                    id TEXT PRIMARY KEY,
+                    user_id TEXT NOT NULL,
+                    query TEXT NOT NULL,
+                    response TEXT NOT NULL,
+                    created_at TEXT NOT NULL,
+                    status TEXT NOT NULL DEFAULT 'pending'
+                )
+            """)
+
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_pending_status
+                ON pending_extractions(status)
+            """)
+
+            conn.commit()
+
+    def save(self, user_id: str, memory: Memory) -> None:
+        """Save a memory for a user.
+
+        Args:
+            user_id: User identifier for isolation
+            memory: Memory to save
+        """
+        with self._get_conn() as conn:
+            conn.execute(
+                """
+                INSERT INTO memories (id, user_id, content, category, importance, created_at, metadata)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    memory.id,
+                    user_id,
+                    memory.content,
+                    memory.category,
+                    memory.importance,
+                    memory.created_at.isoformat(),
+                    json.dumps(memory.metadata),
+                ),
+            )
+            conn.commit()
+
+    def search(self, user_id: str, query: str, top_k: int = 5) -> list[Memory]:
+        """Search memories using FTS5 full-text search.
+
+        Args:
+            user_id: User identifier for isolation
+            query: Search query (auto-escaped, or use raw FTS5 syntax with prefix '_raw:')
+            top_k: Maximum number of results
+
+        Returns:
+            List of matching memories, ranked by relevance
+        """
+        # Sanitize query for FTS5 (escape special characters unless raw mode)
+        if query.startswith("_raw:"):
+            fts_query = query[5:]  # Use raw FTS5 syntax
+        else:
+            fts_query = self._sanitize_fts_query(query)
+
+        if not fts_query.strip():
+            return []
+
+        with self._get_conn() as conn:
+            # Use FTS5 MATCH with BM25 ranking, filtered by user_id
+            cursor = conn.execute(
+                """
+                SELECT m.*, bm25(memories_fts) as rank
+                FROM memories m
+                JOIN memories_fts ON m.rowid = memories_fts.rowid
+                WHERE memories_fts MATCH ? AND m.user_id = ?
+                ORDER BY rank
+                LIMIT ?
+                """,
+                (fts_query, user_id, top_k),
+            )
+
+            results = []
+            for row in cursor:
+                results.append(
+                    Memory(
+                        id=row["id"],
+                        content=row["content"],
+                        category=row["category"],
+                        importance=row["importance"],
+                        created_at=datetime.fromisoformat(row["created_at"]),
+                        metadata=json.loads(row["metadata"]),
+                    )
+                )
+            return results
+
+    def _sanitize_fts_query(self, query: str) -> str:
+        """Sanitize a query for FTS5.
+
+        Escapes special characters and converts to prefix search for better matching.
+
+        Args:
+            query: Raw user query
+
+        Returns:
+            FTS5-safe query string
+        """
+        # FTS5 special characters that need escaping
+        # We use a simple approach: extract words and use OR between them
+        import re
+
+        # Extract alphanumeric words
+        words = re.findall(r"\w+", query)
+
+        if not words:
+            return ""
+
+        # Use OR between words with prefix matching for flexibility
+        # This allows "What language" to match "Python" memories when searching
+        # by using prefix matching (word*)
+        escaped_words = []
+        for word in words:
+            # Quote each word to handle any remaining special chars
+            escaped_words.append(f'"{word}"')
+
+        return " OR ".join(escaped_words)
+
+    def get_all(self, user_id: str) -> list[Memory]:
+        """Get all memories for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            All memories for the user, ordered by creation time (newest first)
+        """
+        with self._get_conn() as conn:
+            cursor = conn.execute(
+                """
+                SELECT * FROM memories
+                WHERE user_id = ?
+                ORDER BY created_at DESC
+                """,
+                (user_id,),
+            )
+
+            return [
+                Memory(
+                    id=row["id"],
+                    content=row["content"],
+                    category=row["category"],
+                    importance=row["importance"],
+                    created_at=datetime.fromisoformat(row["created_at"]),
+                    metadata=json.loads(row["metadata"]),
+                )
+                for row in cursor
+            ]
+
+    def delete(self, user_id: str, memory_id: str) -> bool:
+        """Delete a specific memory.
+
+        Args:
+            user_id: User identifier
+            memory_id: ID of memory to delete
+
+        Returns:
+            True if deleted, False if not found
+        """
+        with self._get_conn() as conn:
+            cursor = conn.execute(
+                "DELETE FROM memories WHERE id = ? AND user_id = ?",
+                (memory_id, user_id),
+            )
+            conn.commit()
+            return cursor.rowcount > 0
+
+    def clear(self, user_id: str) -> int:
+        """Delete all memories for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            Number of memories deleted
+        """
+        with self._get_conn() as conn:
+            cursor = conn.execute(
+                "DELETE FROM memories WHERE user_id = ?",
+                (user_id,),
+            )
+            conn.commit()
+            return cursor.rowcount
+
+    def stats(self, user_id: str) -> dict:
+        """Get memory statistics for a user.
+
+        Args:
+            user_id: User identifier
+
+        Returns:
+            Dict with count, categories breakdown, etc.
+        """
+        with self._get_conn() as conn:
+            # Total count
+            total = conn.execute(
+                "SELECT COUNT(*) as count FROM memories WHERE user_id = ?",
+                (user_id,),
+            ).fetchone()["count"]
+
+            # Category breakdown
+            categories = {}
+            for row in conn.execute(
+                """
+                SELECT category, COUNT(*) as count
+                FROM memories WHERE user_id = ?
+                GROUP BY category
+                """,
+                (user_id,),
+            ):
+                categories[row["category"]] = row["count"]
+
+            return {
+                "total": total,
+                "categories": categories,
+            }
+
+    # --- Pending Extractions (for crash recovery) ---
+
+    def queue_extraction(self, pending: PendingExtraction) -> None:
+        """Queue a conversation for memory extraction.
+
+        Args:
+            pending: The pending extraction to queue
+        """
+        with self._get_conn() as conn:
+            conn.execute(
+                """
+                INSERT INTO pending_extractions (id, user_id, query, response, created_at, status)
+                VALUES (?, ?, ?, ?, ?, ?)
+                """,
+                (
+                    pending.id,
+                    pending.user_id,
+                    pending.query,
+                    pending.response,
+                    pending.created_at.isoformat(),
+                    pending.status,
+                ),
+            )
+            conn.commit()
+
+    def get_pending_extractions(
+        self, limit: int = 10, status: str = "pending"
+    ) -> list[PendingExtraction]:
+        """Get pending extractions for processing.
+
+        Args:
+            limit: Maximum number to return
+            status: Filter by status
+
+        Returns:
+            List of pending extractions
+        """
+        with self._get_conn() as conn:
+            cursor = conn.execute(
+                """
+                SELECT * FROM pending_extractions
+                WHERE status = ?
+                ORDER BY created_at ASC
+                LIMIT ?
+                """,
+                (status, limit),
+            )
+
+            return [
+                PendingExtraction(
+                    id=row["id"],
+                    user_id=row["user_id"],
+                    query=row["query"],
+                    response=row["response"],
+                    created_at=datetime.fromisoformat(row["created_at"]),
+                    status=row["status"],
+                )
+                for row in cursor
+            ]
+
+    def update_extraction_status(self, extraction_id: str, status: str) -> None:
+        """Update the status of a pending extraction.
+
+        Args:
+            extraction_id: ID of the extraction
+            status: New status
+        """
+        with self._get_conn() as conn:
+            conn.execute(
+                "UPDATE pending_extractions SET status = ? WHERE id = ?",
+                (status, extraction_id),
+            )
+            conn.commit()
+
+    def delete_extraction(self, extraction_id: str) -> None:
+        """Delete a completed extraction.
+
+        Args:
+            extraction_id: ID of the extraction to delete
+        """
+        with self._get_conn() as conn:
+            conn.execute(
+                "DELETE FROM pending_extractions WHERE id = ?",
+                (extraction_id,),
+            )
+            conn.commit()

headroom/memory/worker.py

@@ -0,0 +1,260 @@
+"""Background worker for batched memory extraction.
+
+Collects conversations in a queue and processes them in batches,
+reducing LLM calls and improving efficiency.
+"""
+
+from __future__ import annotations
+
+import atexit
+import logging
+import threading
+import time
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from headroom.memory.extractor import MemoryExtractor
+    from headroom.memory.store import SQLiteMemoryStore
+
+
+logger = logging.getLogger(__name__)
+
+
+class ExtractionWorker:
+    """Background worker that batches memory extractions.
+
+    Features:
+    - Collects conversations in a queue
+    - Processes in batches (configurable size and timeout)
+    - Persists pending work to SQLite for crash recovery
+    - Thread-safe, daemon thread (stops with main program)
+
+    Usage:
+        worker = ExtractionWorker(store, extractor)
+        worker.start()
+        worker.schedule("alice", "I prefer Python", "Great choice!")
+        # ... later, memories are extracted and saved automatically
+    """
+
+    def __init__(
+        self,
+        store: SQLiteMemoryStore,
+        extractor: MemoryExtractor,
+        batch_size: int = 10,
+        max_wait_seconds: float = 30.0,
+    ):
+        """Initialize the worker.
+
+        Args:
+            store: Memory store for saving extracted memories
+            extractor: Extractor for processing conversations
+            batch_size: Max conversations per batch
+            max_wait_seconds: Max time to wait before processing partial batch
+        """
+        self.store = store
+        self.extractor = extractor
+        self.batch_size = batch_size
+        self.max_wait_seconds = max_wait_seconds
+
+        self._queue: list[tuple[str, str, str]] = []  # (user_id, query, response)
+        self._lock = threading.Lock()
+        self._event = threading.Event()
+        self._running = False
+        self._thread: threading.Thread | None = None
+
+        # Register cleanup on exit
+        atexit.register(self._cleanup)
+
+    def start(self) -> None:
+        """Start the background worker thread."""
+        if self._running:
+            return
+
+        self._running = True
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+
+        # Process any pending extractions from previous runs (crash recovery)
+        self._recover_pending()
+
+    def stop(self, wait: bool = True, timeout: float = 5.0) -> None:
+        """Stop the worker.
+
+        Args:
+            wait: If True, process remaining queue before stopping
+            timeout: Max time to wait for remaining work
+        """
+        if not self._running:
+            return
+
+        self._running = False
+        self._event.set()  # Wake up the thread
+
+        if wait and self._thread:
+            self._thread.join(timeout=timeout)
+
+    def schedule(self, user_id: str, query: str, response: str) -> None:
+        """Schedule a conversation for memory extraction.
+
+        Non-blocking - returns immediately and extracts in background.
+
+        Args:
+            user_id: User identifier
+            query: User's message
+            response: Assistant's response
+        """
+        # Persist to SQLite first (crash recovery)
+        from headroom.memory.store import PendingExtraction
+
+        pending = PendingExtraction(
+            user_id=user_id,
+            query=query,
+            response=response,
+        )
+        self.store.queue_extraction(pending)
+
+        # Add to in-memory queue
+        with self._lock:
+            self._queue.append((user_id, query, response))
+
+            # Wake up worker if batch is full
+            if len(self._queue) >= self.batch_size:
+                self._event.set()
+
+    def flush(self, timeout: float = 60.0) -> bool:
+        """Force immediate processing of all queued extractions.
+
+        Blocks until all pending extractions are processed or timeout.
+
+        Args:
+            timeout: Max time to wait in seconds
+
+        Returns:
+            True if all extractions completed, False if timed out
+        """
+        # Signal worker to process immediately by temporarily setting max_wait to 0
+        original_max_wait = self.max_wait_seconds
+        self.max_wait_seconds = 0
+        self._event.set()
+
+        # Wait for queue to empty
+        start = time.time()
+        while time.time() - start < timeout:
+            pending = self.store.get_pending_extractions(limit=1, status="pending")
+            if not pending:
+                self.max_wait_seconds = original_max_wait
+                return True
+            time.sleep(0.5)
+
+        self.max_wait_seconds = original_max_wait
+        return False
+
+    def _run(self) -> None:
+        """Main worker loop."""
+        last_process_time = time.time()
+
+        while self._running:
+            # Wait for batch to fill or timeout
+            self._event.wait(timeout=1.0)
+            self._event.clear()
+
+            now = time.time()
+            time_since_last = now - last_process_time
+
+            with self._lock:
+                should_process = len(self._queue) >= self.batch_size or (
+                    self._queue and time_since_last >= self.max_wait_seconds
+                )
+
+                if should_process:
+                    batch = self._queue[: self.batch_size]
+                    self._queue = self._queue[self.batch_size :]
+                else:
+                    batch = []
+
+            if batch:
+                self._process_batch(batch)
+                last_process_time = time.time()
+
+        # Process remaining queue on shutdown
+        with self._lock:
+            remaining = self._queue[:]
+            self._queue = []
+
+        if remaining:
+            self._process_batch(remaining)
+
+    def _process_batch(self, batch: list[tuple[str, str, str]]) -> None:
+        """Process a batch of conversations.
+
+        Args:
+            batch: List of (user_id, query, response) tuples
+        """
+        logger.debug(f"Processing batch of {len(batch)} conversations")
+
+        try:
+            # Extract memories
+            result = self.extractor.extract_batch(batch)
+
+            # Save memories
+            for user_id, memories in result.items():
+                for memory in memories:
+                    self.store.save(user_id, memory)
+                    logger.debug(f"Saved memory for {user_id}: {memory.content[:50]}...")
+
+            # Mark pending extractions as done
+            # Note: In a production system, we'd track exact IDs
+            # For simplicity, we clear pending by matching user/query/response
+            self._mark_batch_done(batch)
+
+        except Exception as e:
+            logger.error(f"Batch extraction failed: {e}")
+            self._mark_batch_failed(batch)
+
+    def _recover_pending(self) -> None:
+        """Recover pending extractions from previous runs."""
+        pending = self.store.get_pending_extractions(limit=100, status="pending")
+
+        if not pending:
+            return
+
+        logger.info(f"Recovering {len(pending)} pending extractions")
+
+        with self._lock:
+            for p in pending:
+                self._queue.append((p.user_id, p.query, p.response))
+
+        # Trigger processing
+        self._event.set()
+
+    def _mark_batch_done(self, batch: list[tuple[str, str, str]]) -> None:
+        """Mark batch items as completed in the pending table."""
+        # Get pending extractions and mark matching ones as done
+        pending = self.store.get_pending_extractions(limit=100)
+
+        for user_id, query, response in batch:
+            for p in pending:
+                if p.user_id == user_id and p.query == query and p.response == response:
+                    self.store.delete_extraction(p.id)
+                    break
+
+    def _mark_batch_failed(self, batch: list[tuple[str, str, str]]) -> None:
+        """Mark batch items as failed in the pending table."""
+        pending = self.store.get_pending_extractions(limit=100)
+
+        for user_id, query, response in batch:
+            for p in pending:
+                if p.user_id == user_id and p.query == query and p.response == response:
+                    self.store.update_extraction_status(p.id, "failed")
+                    break
+
+    def _cleanup(self) -> None:
+        """Cleanup on program exit."""
+        if self._running:
+            self.stop(wait=True, timeout=2.0)
+
+    @property
+    def queue_size(self) -> int:
+        """Get current queue size."""
+        with self._lock:
+            return len(self._queue)