comodor 0.1.0__py3-none-any.whl

comodor/learning/bm25.py

@@ -0,0 +1,147 @@
+"""BM25 ranking in pure Python.
+
+SQLite's FTS5 already provides ``bm25()`` and is used when the interpreter was
+built with it. This module is the fallback for the builds that were not, and it
+doubles as the similarity function used to detect that a "new" lesson is really
+a restatement of one already stored.
+
+The corpus here is small — a few hundred lessons, not a web index — so a plain
+in-memory inverted index is both simpler and fast enough.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from collections import Counter
+from dataclasses import dataclass, field
+
+# Split on anything that is not a word character; keep identifiers like
+# `read_file` and `src/tools/fs.py` together because those are exactly the terms
+# that make a lesson findable. Trailing punctuation is stripped afterwards —
+# leaving it on turns "unittest." and "unittest" into different terms, which
+# quietly destroys both retrieval and duplicate detection.
+_TOKEN = re.compile(r"[A-Za-z0-9_][A-Za-z0-9_\-./]*")
+_TRAILING = ".-/_"
+
+STOPWORDS = frozenset("""
+a an and are as at be but by for from has have if in into is it its of on or
+that the their then there these this to was were will with you your do does
+i we they he she them our us me my than when over while after before also
+just so such not no yes can could should would may might must
+""".split())
+
+K1 = 1.5      # term-frequency saturation
+B = 0.75      # length normalisation
+
+
+def tokenize(text: str) -> list[str]:
+    tokens: list[str] = []
+    for raw in _TOKEN.findall(text or ""):
+        token = raw.lower().rstrip(_TRAILING)
+        if len(token) > 1 and token not in STOPWORDS:
+            tokens.append(token)
+    return tokens
+
+
+@dataclass
+class BM25Index:
+    """An inverted index over ``(doc_id, text)`` pairs."""
+
+    documents: dict[str, list[str]] = field(default_factory=dict)
+    frequencies: dict[str, dict[str, int]] = field(default_factory=dict)  # term -> doc -> tf
+    lengths: dict[str, int] = field(default_factory=dict)
+    _avg_length: float = 0.0
+
+    def add(self, doc_id: str, text: str) -> None:
+        tokens = tokenize(text)
+        self.remove(doc_id)
+        self.documents[doc_id] = tokens
+        self.lengths[doc_id] = len(tokens)
+        for term, count in Counter(tokens).items():
+            self.frequencies.setdefault(term, {})[doc_id] = count
+        self._avg_length = 0.0
+
+    def remove(self, doc_id: str) -> None:
+        tokens = self.documents.pop(doc_id, None)
+        if tokens is None:
+            return
+        self.lengths.pop(doc_id, None)
+        for term in set(tokens):
+            postings = self.frequencies.get(term)
+            if postings:
+                postings.pop(doc_id, None)
+                if not postings:
+                    del self.frequencies[term]
+        self._avg_length = 0.0
+
+    @property
+    def avg_length(self) -> float:
+        if not self._avg_length and self.lengths:
+            self._avg_length = sum(self.lengths.values()) / len(self.lengths)
+        return self._avg_length or 1.0
+
+    def search(self, query: str, limit: int = 10) -> list[tuple[str, float]]:
+        """Return ``(doc_id, score)`` pairs, best first."""
+        terms = tokenize(query)
+        if not terms or not self.documents:
+            return []
+
+        total_docs = len(self.documents)
+        avg = self.avg_length
+        scores: dict[str, float] = {}
+
+        for term in terms:
+            postings = self.frequencies.get(term)
+            if not postings:
+                continue
+            # Standard BM25 IDF, floored so a term in nearly every document
+            # contributes nothing rather than something negative.
+            idf = max(0.0, math.log(
+                1 + (total_docs - len(postings) + 0.5) / (len(postings) + 0.5)))
+            for doc_id, frequency in postings.items():
+                length = self.lengths.get(doc_id, 1)
+                denominator = frequency + K1 * (1 - B + B * length / avg)
+                scores[doc_id] = scores.get(doc_id, 0.0) + idf * frequency * (K1 + 1) / denominator
+
+        ranked = sorted(scores.items(), key=lambda item: item[1], reverse=True)
+        return ranked[:limit]
+
+    def __len__(self) -> int:
+        return len(self.documents)
+
+
+def similarity(left: str, right: str) -> float:
+    """How much two texts say the same thing, in ``[0, 1]``.
+
+    This uses the overlap coefficient — the intersection over the *smaller*
+    token set — rather than Jaccard. Deduplication compares a freshly written
+    lesson against a stored one, and the two are rarely the same length; Jaccard
+    punishes the longer text for its extra words and scores obvious restatements
+    far too low. Very short texts fall back to Jaccard, where the overlap
+    coefficient is too easy to satisfy by accident.
+    """
+    a, b = set(tokenize(left)), set(tokenize(right))
+    if not a or not b:
+        return 0.0
+    shared = len(a & b)
+    if min(len(a), len(b)) < 3:
+        return shared / len(a | b)
+    return shared / min(len(a), len(b))
+
+
+def coverage(query: str, document: str) -> float:
+    """The fraction of the query's terms that appear in the document.
+
+    BM25 is the right ranking function for a large corpus and the wrong one for
+    a small one: with a handful of documents the IDF of every term collapses
+    toward zero and all scores become indistinguishable noise. A brand-new brain
+    is exactly that case — and it is when good recall matters most, because the
+    user is deciding whether the learning is worth anything. Coverage stays
+    meaningful at any corpus size, so retrieval blends the two.
+    """
+    terms = set(tokenize(query))
+    if not terms:
+        return 0.0
+    present = set(tokenize(document))
+    return len(terms & present) / len(terms)

comodor/learning/hotindex.py

@@ -0,0 +1,228 @@
+"""The RAM mirror that makes memory instant.
+
+Retrieval used to go through SQL and re-tokenise every stored lesson on every
+call. That is fine at ten lessons and unacceptable at three thousand: measured
+before this module existed, deduplication took 22 ms and recall pulled the whole
+table through SQLite. Both sat directly on the path between the user pressing
+Enter and the first token arriving.
+
+The fix is boring and effective: keep the whole corpus in memory with its tokens
+already computed, plus a term to id inverted index. Lookups then touch only the
+documents that share a word with the query — a set intersection over a handful
+of candidates instead of a scan over everything.
+
+The trade is memory: a few hundred bytes per lesson, so ten thousand lessons cost
+single-digit megabytes. That is the right side of the trade for something on the
+critical path.
+"""
+
+from __future__ import annotations
+
+import threading
+from itertools import islice
+from dataclasses import dataclass
+from typing import Iterable, Iterator
+
+from .bm25 import tokenize
+
+EMPTY: frozenset[str] = frozenset()
+
+# Ceiling on how many documents a single lookup will score. Bounding the work
+# is what makes recall cost the same at ten lessons and at ten thousand.
+MAX_CANDIDATES = 400
+
+
+@dataclass(slots=True)
+class Doc:
+    """One indexed record: its id, its scope, and its tokens."""
+
+    id: int
+    kind: str                    # lesson | rule | skill
+    scope: str
+    tokens: frozenset[str]
+    text: str = ""
+
+    @property
+    def key(self) -> tuple[str, int]:
+        return (self.kind, self.id)
+
+
+class HotIndex:
+    """An in-memory inverted index, maintained incrementally.
+
+    Thread-safe because reflection writes from a background thread while the UI
+    thread prefetches recall. The lock is held only for the pointer swaps, never
+    across scoring, so a slow query cannot block a write.
+    """
+
+    def __init__(self) -> None:
+        self._docs: dict[tuple[str, int], Doc] = {}
+        self._terms: dict[str, set[tuple[str, int]]] = {}
+        self._lock = threading.RLock()
+
+    # -- maintenance ------------------------------------------------------ #
+
+    def add(self, kind: str, doc_id: int, text: str, scope: str = "global") -> Doc:
+        """Index one record, replacing any earlier version of it."""
+        tokens = frozenset(tokenize(text))
+        doc = Doc(id=doc_id, kind=kind, scope=scope, tokens=tokens, text=text)
+        with self._lock:
+            self._remove_locked(doc.key)
+            self._docs[doc.key] = doc
+            for term in tokens:
+                self._terms.setdefault(term, set()).add(doc.key)
+        return doc
+
+    def remove(self, kind: str, doc_id: int) -> bool:
+        with self._lock:
+            return self._remove_locked((kind, doc_id))
+
+    def _remove_locked(self, key: tuple[str, int]) -> bool:
+        doc = self._docs.pop(key, None)
+        if doc is None:
+            return False
+        for term in doc.tokens:
+            postings = self._terms.get(term)
+            if postings is not None:
+                postings.discard(key)
+                if not postings:
+                    del self._terms[term]
+        return True
+
+    def rebuild(self, records: Iterable[tuple[str, int, str, str]]) -> int:
+        """Replace the whole index. Used once at startup."""
+        docs: dict[tuple[str, int], Doc] = {}
+        terms: dict[str, set[tuple[str, int]]] = {}
+        for kind, doc_id, text, scope in records:
+            tokens = frozenset(tokenize(text))
+            doc = Doc(id=doc_id, kind=kind, scope=scope, tokens=tokens, text=text)
+            docs[doc.key] = doc
+            for term in tokens:
+                terms.setdefault(term, set()).add(doc.key)
+        with self._lock:
+            self._docs, self._terms = docs, terms
+        return len(docs)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._docs.clear()
+            self._terms.clear()
+
+    # -- lookup ----------------------------------------------------------- #
+
+    def candidates(self, query: str, kind: str = "",
+                   scopes: Iterable[str] | None = None,
+                   max_candidates: int = MAX_CANDIDATES) -> tuple[set[str], list[Doc]]:
+        """Documents sharing at least one term with the query.
+
+        Rare terms are expanded first and the candidate set is capped, which is
+        what keeps lookup flat as the brain grows. The reasoning is the same one
+        behind IDF: a term that appears in most of the corpus barely changes the
+        ranking but costs the most to scan, so once there are enough candidates
+        the common terms are not worth expanding. Anything they would have added
+        would have scored near the bottom anyway.
+
+        Returns the query tokens too, because every caller needs them next and
+        tokenising twice is pure waste on a hot path.
+        """
+        query_tokens = set(tokenize(query))
+        if not query_tokens:
+            return query_tokens, []
+
+        allowed = set(scopes) if scopes else None
+        with self._lock:
+            postings = sorted(
+                (self._terms.get(term, EMPTY) for term in query_tokens),
+                key=len,
+            )
+            keys: set[tuple[str, int]] = set()
+            for matches in postings:
+                room = max_candidates - len(keys)
+                if room <= 0:
+                    break
+                if len(matches) <= room:
+                    keys |= matches                        # type: ignore[arg-type]
+                else:
+                    # This term is too common to be selective, so the documents
+                    # it would add are near-random with respect to relevance.
+                    # Take a bounded slice: the genuinely informative matches
+                    # already came from the rarer terms processed before it.
+                    keys.update(islice(matches, room))
+                    break
+            docs = [self._docs[key] for key in keys if key in self._docs]
+
+        return query_tokens, [
+            doc for doc in docs
+            if (not kind or doc.kind == kind) and (allowed is None or doc.scope in allowed)
+        ]
+
+    def coverage_scan(self, query: str, kind: str = "",
+                      scopes: Iterable[str] | None = None,
+                      limit: int = 20) -> list[tuple[Doc, float]]:
+        """Rank candidates by how much of the query they cover.
+
+        Coverage rather than BM25 because this index exists to be fast and
+        stable at any corpus size; BM25's IDF is meaningless on a small brain.
+        The caller blends this with FTS ranking where that is available.
+        """
+        query_tokens, docs = self.candidates(query, kind, scopes)
+        if not query_tokens or not docs:
+            return []
+        size = len(query_tokens)
+        scored = [(doc, len(query_tokens & doc.tokens) / size) for doc in docs]
+        scored.sort(key=lambda pair: pair[1], reverse=True)
+        return scored[:limit]
+
+    def nearest(self, text: str, threshold: float = 0.55, kind: str = "",
+                scopes: Iterable[str] | None = None) -> tuple[Doc | None, float]:
+        """The closest existing record, for duplicate detection.
+
+        Uses the overlap coefficient, matching
+        :func:`comodor.learning.bm25.similarity`, so deduplication behaves the
+        same whether it runs through here or through the store.
+        """
+        query_tokens, docs = self.candidates(text, kind, scopes)
+        if not query_tokens:
+            return None, 0.0
+
+        best: Doc | None = None
+        best_score = threshold
+        for doc in docs:
+            if not doc.tokens:
+                continue
+            smaller = min(len(query_tokens), len(doc.tokens))
+            if smaller < 3:
+                score = len(query_tokens & doc.tokens) / len(query_tokens | doc.tokens)
+            else:
+                score = len(query_tokens & doc.tokens) / smaller
+            if score > best_score:
+                best, best_score = doc, score
+        return best, (best_score if best is not None else 0.0)
+
+    # -- introspection ---------------------------------------------------- #
+
+    def get(self, kind: str, doc_id: int) -> Doc | None:
+        with self._lock:
+            return self._docs.get((kind, doc_id))
+
+    def of_kind(self, kind: str) -> Iterator[Doc]:
+        with self._lock:
+            return iter([doc for doc in self._docs.values() if doc.kind == kind])
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._docs)
+
+    @property
+    def terms(self) -> int:
+        with self._lock:
+            return len(self._terms)
+
+    def stats(self) -> dict[str, int]:
+        with self._lock:
+            kinds: dict[str, int] = {}
+            for doc in self._docs.values():
+                kinds[doc.kind] = kinds.get(doc.kind, 0) + 1
+            postings = sum(len(keys) for keys in self._terms.values())
+        return {"documents": len(self._docs), "terms": len(self._terms),
+                "postings": postings, **{f"kind_{k}": v for k, v in kinds.items()}}