memorytrace 0.1.0__py3-none-any.whl

engram/quality/pii.py

@@ -0,0 +1,80 @@
+"""PII (Personally Identifiable Information) detection and masking."""
+
+from __future__ import annotations
+
+import re
+from typing import Optional
+
+from engram.models.quality import PIIMatch
+
+# PII detection patterns
+_PII_PATTERNS: dict[str, re.Pattern] = {
+    "credit_card": re.compile(
+        r'\b\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}\b'
+    ),
+    "ssn": re.compile(
+        r'\b\d{3}-\d{2}-\d{4}\b'
+    ),
+    "api_key": re.compile(
+        r'\b(?:sk-|ak-|AKIA|ghp_|gho_|github_pat_)[A-Za-z0-9_\-]{20,}\b'
+    ),
+    "email": re.compile(
+        r'\b[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Z|a-z]{2,}\b'
+    ),
+    "phone_kr": re.compile(
+        r'\b0\d{1,2}-\d{3,4}-\d{4}\b'
+    ),
+    "phone_intl": re.compile(
+        r'\+\d{1,3}[-\s]?\d{1,4}[-\s]?\d{3,4}[-\s]?\d{4}\b'
+    ),
+    "rrn": re.compile(
+        # Korean Resident Registration Number (주민등록번호)
+        r'\b\d{6}[-\s]?[1-4]\d{6}\b'
+    ),
+}
+
+_REDACTED = "[REDACTED]"
+
+
+class PIIDetector:
+    """Detects and masks PII in text."""
+
+    def scan(self, text: str) -> list[PIIMatch]:
+        """Scan text for PII patterns. Returns list of matches."""
+        matches: list[PIIMatch] = []
+        for pii_type, pattern in _PII_PATTERNS.items():
+            for m in pattern.finditer(text):
+                # Store only type and length, not the original PII value
+                original_display = f"[{pii_type}:{len(m.group())}chars]"
+                matches.append(PIIMatch(
+                    pii_type=pii_type,
+                    start=m.start(),
+                    end=m.end(),
+                    original=original_display,
+                ))
+        # Sort by position (reverse for safe replacement)
+        matches.sort(key=lambda x: x.start)
+        return matches
+
+    def mask(self, text: str, matches: Optional[list[PIIMatch]] = None) -> str:
+        """Replace PII occurrences with [REDACTED].
+
+        If matches not provided, scans first.
+        """
+        if matches is None:
+            matches = self.scan(text)
+        if not matches:
+            return text
+
+        # Replace from end to start to preserve positions
+        result = text
+        for m in sorted(matches, key=lambda x: x.start, reverse=True):
+            result = result[:m.start] + _REDACTED + result[m.end:]
+        return result
+
+    def has_pii(self, text: str) -> bool:
+        """Quick check: does text contain any PII?"""
+        for pattern in _PII_PATTERNS.values():
+            if pattern.search(text):
+                return True
+        return False

engram/search/__init__.py

@@ -0,0 +1,13 @@
+"""Search backends for Engram."""
+
+from engram.search.base import SearchBackend
+from engram.search.fts5_search import FTS5Search
+from engram.search.tokenizer import tokenize, build_fts5_query, escape_fts5_query
+
+__all__ = [
+    "SearchBackend",
+    "FTS5Search",
+    "tokenize",
+    "build_fts5_query",
+    "escape_fts5_query",
+]

engram/search/base.py

@@ -0,0 +1,20 @@
+"""Search backend protocol."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from engram.models.search import SearchOptions, SearchResult
+
+
+@runtime_checkable
+class SearchBackend(Protocol):
+    """Abstract search interface — all search backends implement this."""
+
+    def search(self, options: SearchOptions) -> SearchResult: ...
+
+    def index_entity(self, entity_id: str) -> None: ...
+
+    def remove_from_index(self, entity_id: str) -> None: ...
+
+    def reindex_all(self) -> int: ...

engram/search/fts5_search.py

@@ -0,0 +1,210 @@
+"""FTS5-based BM25 search — genuine Okapi BM25 via SQLite's built-in ranking."""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import time
+from pathlib import Path
+from typing import Optional
+
+from engram.models.entity import Entity, EntityState, Tier
+from engram.models.fact import Fact, FactStatus
+from engram.models.search import SearchHit, SearchOptions, SearchResult
+from engram.models.source import Source, SourceType
+from engram.search.tokenizer import build_fts5_query
+
+
+class FTS5Search:
+    """Full-text search using SQLite FTS5's built-in BM25 ranking.
+
+    FTS5 bm25() implements genuine Okapi BM25 with:
+    - TF saturation (k1 parameter)
+    - Document length normalization (b parameter)
+    - IDF via document frequency counting
+
+    Column weights (higher = more important):
+    - entity_name: 10.0
+    - summary: 5.0
+    - fact_text: 3.0
+    - state_text: 2.0
+    - entity_type: 1.0
+    - entity_id: 0.0 (not searchable)
+    """
+
+    # Column order in FTS table: entity_id, entity_name, entity_type, fact_text, state_text, summary
+    # bm25() weights correspond to column order
+    _BM25_WEIGHTS = (0.0, 10.0, 1.0, 3.0, 2.0, 5.0)
+    _BM25_WEIGHTS_SQL = "0.0, 10.0, 1.0, 3.0, 2.0, 5.0"  # Pre-built constant for SQL
+
+    def __init__(self, db_path: Path):
+        self.db_path = Path(db_path).resolve()
+        self._conn: Optional[sqlite3.Connection] = None
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            self._conn = sqlite3.connect(str(self.db_path))
+            self._conn.execute("PRAGMA busy_timeout=5000")
+            self._conn.row_factory = sqlite3.Row
+        return self._conn
+
+    def close(self) -> None:
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+
+    def search(self, options: SearchOptions) -> SearchResult:
+        """Execute a BM25-ranked full-text search."""
+        start = time.monotonic()
+
+        fts_query = build_fts5_query(options.query)
+        if not fts_query or fts_query == '""':
+            return SearchResult(
+                query=options.query,
+                hits=[],
+                total_count=0,
+                search_time_ms=0.0,
+            )
+
+        conn = self._get_conn()
+
+        # Search with BM25 ranking and join with entities table for full data
+        # Weights are a class-level constant — never derived from user input
+        sql = f"""
+            SELECT
+                e.*,
+                bm25(memory_fts, {self._BM25_WEIGHTS_SQL}) AS rank,
+                snippet(memory_fts, 3, '>>>', '<<<', '...', 40) AS snippet
+            FROM memory_fts fts
+            JOIN entities e ON e.id = fts.entity_id
+            WHERE memory_fts MATCH ?
+        """
+        params: list = [fts_query]
+
+        # Apply filters
+        if options.entity_types:
+            placeholders = ",".join("?" for _ in options.entity_types)
+            sql += f" AND e.entity_type IN ({placeholders})"
+            params.extend(options.entity_types)
+
+        if options.tiers:
+            placeholders = ",".join("?" for _ in options.tiers)
+            sql += f" AND e.tier IN ({placeholders})"
+            params.extend(options.tiers)
+
+        if not options.include_archived:
+            sql += " AND e.tier != 'archival'"
+
+        if options.date_from:
+            sql += " AND e.updated_at >= ?"
+            params.append(options.date_from.isoformat())
+
+        if options.date_to:
+            sql += " AND e.updated_at <= ?"
+            params.append(options.date_to.isoformat())
+
+        sql += " ORDER BY rank LIMIT ?"
+        params.append(options.max_results)
+
+        try:
+            rows = conn.execute(sql, params).fetchall()
+        except sqlite3.OperationalError as e:
+            err_msg = str(e).lower()
+            # Only swallow FTS5 query syntax errors; re-raise others
+            if "fts5" in err_msg or "syntax" in err_msg or "parse" in err_msg:
+                return SearchResult(
+                    query=options.query,
+                    hits=[],
+                    total_count=0,
+                    search_time_ms=(time.monotonic() - start) * 1000,
+                )
+            raise
+
+        # Build hits
+        hits: list[SearchHit] = []
+        total_tokens = 0
+
+        for row in rows:
+            entity = self._row_to_entity(row)
+            snippet = row["snippet"] if row["snippet"] else ""
+            # Clean snippet markers
+            snippet = snippet.replace(">>>", "").replace("<<<", "")
+
+            # Load facts for this entity if within confidence threshold
+            facts = self._get_entity_facts(
+                conn, entity.id, options.min_confidence
+            )
+
+            # Estimate token count (1 token ≈ 4 chars)
+            entry_text = f"{entity.name} {snippet} {' '.join(f.raw_text for f in facts[:3])}"
+            token_count = len(entry_text) // 4
+
+            # Check token budget (max_tokens <= 0 means unlimited)
+            if options.max_tokens > 0 and total_tokens + token_count > options.max_tokens:
+                break
+
+            hits.append(SearchHit(
+                entity=entity,
+                facts=facts[:5],  # Limit facts per hit
+                relevance_score=abs(row["rank"]),  # bm25() returns negative scores
+                snippet=snippet,
+                token_count=token_count,
+            ))
+            total_tokens += token_count
+
+        elapsed = (time.monotonic() - start) * 1000
+
+        return SearchResult(
+            query=options.query,
+            hits=hits,
+            total_count=len(rows),  # Total DB matches (may exceed len(hits) due to token budget)
+            search_time_ms=round(elapsed, 2),
+            total_tokens=total_tokens,
+        )
+
+    def _get_entity_facts(
+        self, conn: sqlite3.Connection, entity_id: str, min_confidence: float
+    ) -> list[Fact]:
+        """Load current facts for an entity, filtered by confidence."""
+        rows = conn.execute(
+            """SELECT * FROM facts
+               WHERE entity_id = ? AND valid_to IS NULL AND superseded_by IS NULL
+               AND status NOT IN ('expired', 'retracted', 'conflicted')
+               AND confidence >= ?
+               ORDER BY confidence DESC, created_at DESC
+               LIMIT 10""",
+            (entity_id, min_confidence),
+        ).fetchall()
+
+        facts = []
+        for r in rows:
+            facts.append(Fact(
+                id=r["id"],
+                entity_id=r["entity_id"],
+                subject=r["subject"],
+                predicate=r["predicate"],
+                object=r["object"],
+                raw_text=r["raw_text"],
+                source=Source(
+                    type=SourceType(r["source_type"]),
+                    author=r["source_author"] or "",
+                    channel=r["source_channel"] or "",
+                ),
+                confidence=r["confidence"],
+                status=FactStatus(r["status"]),
+            ))
+        return facts
+
+    def _row_to_entity(self, row: sqlite3.Row) -> Entity:
+        """Convert a joined row to Entity model."""
+        state_data = json.loads(row["state"]) if row["state"] else {}
+        return Entity(
+            id=row["id"],
+            name=row["name"],
+            entity_type=row["entity_type"],
+            tier=Tier(row["tier"]),
+            summary=row["summary"] or "",
+            aliases=json.loads(row["aliases"]) if row["aliases"] else [],
+            state=EntityState.from_dict(state_data),
+            access_count=row["access_count"] or 0,
+        )

engram/search/hybrid.py

@@ -0,0 +1,99 @@
+"""Hybrid search — Reciprocal Rank Fusion of FTS5 + semantic search."""
+
+from __future__ import annotations
+
+import time
+
+from engram.models.search import SearchHit, SearchOptions, SearchResult
+from engram.search.fts5_search import FTS5Search
+from engram.search.semantic import SemanticSearch
+
+
+class HybridSearch:
+    """Reciprocal Rank Fusion (RRF) of keyword and semantic search.
+
+    RRF score = sum over sources of 1 / (k + rank)
+    where k=60 (standard constant).
+
+    This is simple, effective, and well-studied — no learned weights needed.
+    """
+
+    def __init__(
+        self,
+        fts: FTS5Search,
+        semantic: SemanticSearch,
+        k: int = 60,
+        fts_weight: float = 1.0,
+        semantic_weight: float = 1.0,
+    ):
+        self.fts = fts
+        self.semantic = semantic
+        self.k = k
+        self.fts_weight = fts_weight
+        self.semantic_weight = semantic_weight
+
+    def close(self) -> None:
+        self.fts.close()
+        self.semantic.close()
+
+    def search(self, options: SearchOptions) -> SearchResult:
+        """Execute hybrid search with RRF fusion."""
+        start = time.monotonic()
+
+        # Increase individual search limits to get better fusion
+        expanded = SearchOptions(
+            query=options.query,
+            max_results=options.max_results * 3,
+            max_tokens=0,  # No token limit for individual searches
+            min_confidence=options.min_confidence,
+            entity_types=options.entity_types,
+            tiers=options.tiers,
+            include_archived=options.include_archived,
+        )
+
+        # Run both searches
+        fts_result = self.fts.search(expanded)
+        sem_result = self.semantic.search(expanded)
+
+        # RRF scoring
+        scores: dict[str, float] = {}
+        hit_map: dict[str, SearchHit] = {}
+
+        for rank, hit in enumerate(fts_result.hits):
+            eid = hit.entity.id
+            scores[eid] = scores.get(eid, 0) + self.fts_weight / (self.k + rank + 1)
+            if eid not in hit_map:
+                hit_map[eid] = hit
+
+        for rank, hit in enumerate(sem_result.hits):
+            eid = hit.entity.id
+            scores[eid] = scores.get(eid, 0) + self.semantic_weight / (self.k + rank + 1)
+            if eid not in hit_map:
+                hit_map[eid] = hit
+
+        # Sort by RRF score
+        sorted_ids = sorted(scores.keys(), key=lambda eid: scores[eid], reverse=True)
+        sorted_ids = sorted_ids[:options.max_results]
+
+        # Build final hits with token budget
+        hits: list[SearchHit] = []
+        total_tokens = 0
+        for eid in sorted_ids:
+            hit = hit_map[eid]
+            hit.relevance_score = scores[eid]
+
+            token_count = len(f"{hit.entity.name} {hit.snippet}") // 4
+            if options.max_tokens > 0 and total_tokens + token_count > options.max_tokens:
+                break
+            hit.token_count = token_count
+            hits.append(hit)
+            total_tokens += token_count
+
+        elapsed = (time.monotonic() - start) * 1000
+        return SearchResult(
+            query=options.query,
+            hits=hits,
+            total_count=len(scores),
+            search_time_ms=round(elapsed, 2),
+            total_tokens=total_tokens,
+        )

engram/search/semantic.py

@@ -0,0 +1,186 @@
+"""Optional embedding-based semantic search.
+
+Requires: pip install engram[semantic]
+For embeddings API: pip install engram[llm]
+"""
+
+from __future__ import annotations
+
+import json
+import sqlite3
+import struct
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+from engram.models.entity import Entity, EntityState, Tier
+from engram.models.search import SearchHit, SearchOptions, SearchResult
+
+
+class SemanticSearch:
+    """Embedding-based semantic search using numpy cosine similarity.
+
+    Stores embeddings in a dedicated SQLite table alongside the main DB.
+    For <10k entities, in-memory numpy comparison is fast enough.
+    """
+
+    def __init__(
+        self,
+        db_path: Path,
+        embedding_provider: str = "openai",
+        embedding_model: str = "text-embedding-3-small",
+        embedding_dim: int = 1536,
+    ):
+        self.db_path = Path(db_path).resolve()
+        self.embedding_provider = embedding_provider
+        self.embedding_model = embedding_model
+        self.embedding_dim = embedding_dim
+        self._conn: Optional[sqlite3.Connection] = None
+        self._client = None
+        self._np = None
+
+        try:
+            import numpy as np
+            self._np = np
+        except ImportError:
+            raise ImportError("numpy not installed. Run: pip install engram[semantic]")
+
+        self._ensure_table()
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            self._conn = sqlite3.connect(str(self.db_path))
+            self._conn.execute("PRAGMA busy_timeout=5000")
+            self._conn.row_factory = sqlite3.Row
+        return self._conn
+
+    def _ensure_table(self) -> None:
+        conn = self._get_conn()
+        conn.execute("""
+            CREATE TABLE IF NOT EXISTS embeddings (
+                entity_id TEXT PRIMARY KEY,
+                embedding BLOB NOT NULL,
+                model TEXT NOT NULL,
+                updated_at TEXT NOT NULL
+            )
+        """)
+        conn.commit()
+
+    def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    def _get_embedding_client(self):
+        """Lazy-load embedding API client."""
+        if self._client is not None:
+            return self._client
+        if self.embedding_provider == "openai":
+            try:
+                from openai import OpenAI
+                self._client = OpenAI()
+            except ImportError:
+                raise ImportError("OpenAI SDK not installed. Run: pip install engram[llm]")
+        elif self.embedding_provider == "anthropic":
+            raise NotImplementedError("Anthropic does not provide an embedding API. Use 'openai' provider.")
+        return self._client
+
+    def embed(self, text: str) -> bytes:
+        """Compute embedding for text. Returns bytes (packed floats)."""
+        client = self._get_embedding_client()
+        response = client.embeddings.create(
+            model=self.embedding_model,
+            input=text[:8000],
+        )
+        vector = response.data[0].embedding
+        return struct.pack(f"{len(vector)}f", *vector)
+
+    def _bytes_to_vector(self, data: bytes):
+        """Convert packed bytes back to numpy array."""
+        n = len(data) // 4
+        floats = struct.unpack(f"{n}f", data)
+        return self._np.array(floats, dtype=self._np.float32)
+
+    def index_entity(self, entity_id: str, text: str) -> None:
+        """Compute and store embedding for an entity."""
+        embedding_bytes = self.embed(text)
+        conn = self._get_conn()
+        conn.execute(
+            """INSERT OR REPLACE INTO embeddings (entity_id, embedding, model, updated_at)
+               VALUES (?, ?, ?, ?)""",
+            (entity_id, embedding_bytes, self.embedding_model, datetime.now().isoformat()),
+        )
+        conn.commit()
+
+    def remove_from_index(self, entity_id: str) -> None:
+        conn = self._get_conn()
+        conn.execute("DELETE FROM embeddings WHERE entity_id = ?", (entity_id,))
+        conn.commit()
+
+    def search(self, options: SearchOptions) -> SearchResult:
+        """Semantic search using cosine similarity."""
+        start = time.monotonic()
+        np = self._np
+
+        # Embed query
+        query_bytes = self.embed(options.query)
+        query_vec = self._bytes_to_vector(query_bytes)
+        query_norm = np.linalg.norm(query_vec)
+        if query_norm == 0:
+            return SearchResult(query=options.query)
+
+        query_vec = query_vec / query_norm
+
+        # Load all embeddings
+        conn = self._get_conn()
+        rows = conn.execute("SELECT entity_id, embedding FROM embeddings").fetchall()
+
+        if not rows:
+            return SearchResult(query=options.query)
+
+        # Compute cosine similarities
+        scores: list[tuple[str, float]] = []
+        for row in rows:
+            vec = self._bytes_to_vector(row["embedding"])
+            vec_norm = np.linalg.norm(vec)
+            if vec_norm == 0:
+                continue
+            similarity = float(np.dot(query_vec, vec / vec_norm))
+            scores.append((row["entity_id"], similarity))
+
+        scores.sort(key=lambda x: x[1], reverse=True)
+        scores = scores[:options.max_results]
+
+        # Build hits
+        hits: list[SearchHit] = []
+        for entity_id, score in scores:
+            entity_row = conn.execute(
+                "SELECT * FROM entities WHERE id = ?", (entity_id,)
+            ).fetchone()
+            if not entity_row:
+                continue
+
+            state_data = json.loads(entity_row["state"]) if entity_row["state"] else {}
+            entity = Entity(
+                id=entity_row["id"],
+                name=entity_row["name"],
+                entity_type=entity_row["entity_type"],
+                tier=Tier(entity_row["tier"]),
+                summary=entity_row["summary"] or "",
+                state=EntityState.from_dict(state_data),
+            )
+
+            hits.append(SearchHit(
+                entity=entity,
+                relevance_score=score,
+                snippet=entity.summary[:200] if entity.summary else "",
+            ))
+
+        elapsed = (time.monotonic() - start) * 1000
+        return SearchResult(
+            query=options.query,
+            hits=hits,
+            total_count=len(hits),
+            search_time_ms=round(elapsed, 2),
+        )

engram/search/tokenizer.py

@@ -0,0 +1,85 @@
+"""Unicode-aware tokenizer with CJK support."""
+
+from __future__ import annotations
+
+import re
+
+# Common English stopwords (minimal set for search quality)
+_ENGLISH_STOPWORDS = frozenset({
+    "a", "an", "the", "and", "or", "but", "in", "on", "at", "to", "for",
+    "of", "with", "by", "from", "is", "was", "are", "were", "be", "been",
+    "being", "have", "has", "had", "do", "does", "did", "will", "would",
+    "could", "should", "may", "might", "shall", "can", "need", "must",
+    "not", "no", "nor", "so", "if", "then", "than", "that", "this",
+    "these", "those", "it", "its", "he", "she", "they", "we", "you",
+    "i", "me", "my", "your", "his", "her", "our", "their", "who",
+    "what", "which", "when", "where", "how", "why", "all", "each",
+    "every", "both", "few", "more", "most", "some", "any", "such",
+    "only", "own", "same", "very", "just", "about", "above", "after",
+    "again", "also", "as", "because", "before", "between", "during",
+    "into", "over", "through", "under", "until", "up", "while",
+})
+
+# Korean stopwords (particles, endings)
+_KOREAN_STOPWORDS = frozenset({
+    "은", "는", "이", "가", "을", "를", "에", "의", "로", "으로",
+    "와", "과", "도", "만", "부터", "까지", "에서", "한", "할", "하는",
+    "했다", "합니다", "있는", "있다", "없다", "되는", "된", "되다",
+    "그", "저", "이것", "그것", "저것", "여기", "거기", "저기",
+})
+
+# Token pattern: word chars + CJK unified ideographs + Hangul
+_TOKEN_PATTERN = re.compile(
+    r'[a-zA-Z0-9_]+'           # ASCII words
+    r'|[\uAC00-\uD7AF]+'       # Hangul syllables
+    r'|[\u4E00-\u9FFF]+'        # CJK unified ideographs
+    r'|[\u3040-\u309F\u30A0-\u30FF]+'  # Hiragana + Katakana
+, re.UNICODE)
+
+
+def tokenize(text: str, remove_stopwords: bool = True) -> list[str]:
+    """Tokenize text into lowercase tokens with optional stopword removal.
+
+    Handles English, Korean, Chinese, and Japanese text.
+    """
+    tokens = _TOKEN_PATTERN.findall(text.lower())
+    if not remove_stopwords:
+        return tokens
+    return [
+        t for t in tokens
+        if t not in _ENGLISH_STOPWORDS and t not in _KOREAN_STOPWORDS and len(t) > 1
+    ]
+
+
+def escape_fts5_query(query: str) -> str:
+    """Escape special FTS5 characters to prevent syntax errors.
+
+    FTS5 special chars: * " ( ) : ^
+    We wrap each token in double quotes to treat them as literals.
+    """
+    tokens = tokenize(query, remove_stopwords=False)
+    if not tokens:
+        return '""'
+    # Quote each token to escape special chars, join with implicit AND
+    escaped = []
+    for token in tokens:
+        # Remove any embedded double quotes
+        clean = token.replace('"', '')
+        if clean:
+            escaped.append(f'"{clean}"')
+    return " ".join(escaped) if escaped else '""'
+
+
+def build_fts5_query(query: str) -> str:
+    """Build an FTS5 query from natural language input.
+
+    - Removes stopwords
+    - Escapes special characters
+    - Returns FTS5-compatible query string
+    """
+    tokens = tokenize(query, remove_stopwords=True)
+    if not tokens:
+        # Fall back to raw escaped query if all words were stopwords
+        return escape_fts5_query(query)
+    # Quote each token for safety
+    return " ".join(f'"{t}"' for t in tokens)