routing-memory 0.1.0__py3-none-any.whl

rm/__init__.py

@@ -0,0 +1,28 @@
+"""Routing Memory — Lightweight long-term memory via vector-quantized routing."""
+
+from .memory import RoutingMemory
+from .codebook import Codebook, CentroidBucket
+from .retrieval import L1Retriever, L1Result, RetrievalResult
+from .filtering import filter_by_score, filter_top_n
+from .drift import DriftMonitor, DriftAlarm
+from .embeddings import LocalEmbeddings, EmbeddingBackend
+from .storage import RMSQLiteBackend, RMStorageBackend
+
+__all__ = [
+    "RoutingMemory",
+    "Codebook",
+    "CentroidBucket",
+    "L1Retriever",
+    "L1Result",
+    "RetrievalResult",
+    "filter_by_score",
+    "filter_top_n",
+    "DriftMonitor",
+    "DriftAlarm",
+    "LocalEmbeddings",
+    "EmbeddingBackend",
+    "RMSQLiteBackend",
+    "RMStorageBackend",
+]
+
+__version__ = "0.1.0"

rm/codebook.py

@@ -0,0 +1,337 @@
+"""VQ Codebook: MiniBatchKMeans with adaptive K + online adaptation.
+
+Features:
+- Adaptive K: K = ceil(N / B_target)
+- EMA centroid update: centroids drift toward recent data
+- Split: high-variance buckets split into 2
+- Prune: idle centroids get reassigned and removed
+- K range guard: K stays in [K_MIN, K_MAX]
+"""
+
+import logging
+import time
+import numpy as np
+from typing import Dict, List, Optional, Tuple
+from dataclasses import dataclass, field
+
+logger = logging.getLogger("rm.codebook")
+
+
+@dataclass
+class CentroidBucket:
+    """Items assigned to a centroid."""
+    item_ids: List[str] = field(default_factory=list)
+    embeddings: List[np.ndarray] = field(default_factory=list)
+    last_accessed: float = field(default_factory=time.time)
+    access_count: int = 0
+
+
+class Codebook:
+    """Vector Quantization codebook using MiniBatchKMeans.
+
+    Adaptive K: K = ceil(N / B_target) where B_target is items per bucket.
+    """
+
+    B_TARGET = 20  # target items per bucket
+    K_MIN = 2
+    K_MAX = 512
+
+    def __init__(self, dim: int = 384, seed: int = 42):
+        self.dim = dim
+        self.seed = seed
+        self.centroids: Optional[np.ndarray] = None  # (K, dim)
+        self.buckets: Dict[int, CentroidBucket] = {}
+        self._n_items = 0
+        self._fitted = False
+
+    @property
+    def K(self) -> int:
+        if self.centroids is None:
+            return 0
+        return self.centroids.shape[0]
+
+    @property
+    def fitted(self) -> bool:
+        return self._fitted and self.centroids is not None
+
+    def _adaptive_k(self, n: int) -> int:
+        k = max(self.K_MIN, int(np.ceil(n / self.B_TARGET)))
+        return min(k, self.K_MAX)
+
+    def fit(self, embeddings: np.ndarray, item_ids: List[str]) -> None:
+        """Fit codebook from scratch using MiniBatchKMeans."""
+        n = embeddings.shape[0]
+        if n == 0:
+            return
+
+        k = self._adaptive_k(n)
+        k = min(k, n)  # can't have more centroids than items
+
+        from sklearn.cluster import MiniBatchKMeans
+        kmeans = MiniBatchKMeans(
+            n_clusters=k,
+            random_state=self.seed,
+            batch_size=min(256, n),
+            n_init=3,
+        )
+        labels = kmeans.fit_predict(embeddings)
+
+        # Normalize centroids
+        norms = np.linalg.norm(kmeans.cluster_centers_, axis=1, keepdims=True)
+        norms = np.maximum(norms, 1e-10)
+        self.centroids = (kmeans.cluster_centers_ / norms).astype(np.float32)
+
+        # Build buckets
+        self.buckets = {}
+        for idx, label in enumerate(labels):
+            label = int(label)
+            if label not in self.buckets:
+                self.buckets[label] = CentroidBucket()
+            self.buckets[label].item_ids.append(item_ids[idx])
+            self.buckets[label].embeddings.append(embeddings[idx])
+
+        self._n_items = n
+        self._fitted = True
+
+    def encode(self, x: np.ndarray) -> Tuple[int, float]:
+        """Find nearest centroid. Returns (centroid_id, quantization_error)."""
+        if not self.fitted:
+            raise RuntimeError("Codebook not fitted. Call fit() or add items first.")
+        dots = x @ self.centroids.T  # (K,)
+        best = int(np.argmax(dots))
+        qerr = 1.0 - float(dots[best])
+        return best, qerr
+
+    def conf(self, x: np.ndarray) -> float:
+        """Confidence: 1 / (1 + qerr)."""
+        _, qerr = self.encode(x)
+        return 1.0 / (1.0 + qerr)
+
+    def margin(self, x: np.ndarray) -> float:
+        """Gap between top-2 centroid similarities."""
+        if not self.fitted or self.K < 2:
+            return 0.0
+        dots = x @ self.centroids.T
+        top2 = np.partition(dots, -2)[-2:]
+        return float(top2.max() - top2.min())
+
+    def top_centroids(self, x: np.ndarray, n: int = 3) -> List[int]:
+        """Return top-n centroid IDs by similarity."""
+        if not self.fitted:
+            return []
+        dots = x @ self.centroids.T
+        n = min(n, self.K)
+        return list(np.argsort(dots)[-n:][::-1])
+
+    def add(self, item_id: str, embedding: np.ndarray) -> int:
+        """Assign item to nearest centroid. Returns centroid ID."""
+        if not self.fitted:
+            # Bootstrap: collect items and fit when we have enough
+            if 0 not in self.buckets:
+                self.buckets[0] = CentroidBucket()
+            self.buckets[0].item_ids.append(item_id)
+            self.buckets[0].embeddings.append(embedding)
+            self._n_items += 1
+
+            if self._n_items >= self.K_MIN:
+                self._refit()
+            return 0
+
+        cid, _ = self.encode(embedding)
+        if cid not in self.buckets:
+            self.buckets[cid] = CentroidBucket()
+        self.buckets[cid].item_ids.append(item_id)
+        self.buckets[cid].embeddings.append(embedding)
+        self._n_items += 1
+
+        # Check if we need to refit (bucket too large)
+        if self._n_items % 100 == 0:
+            ideal_k = self._adaptive_k(self._n_items)
+            if ideal_k > self.K * 1.5:
+                self._refit()
+
+        return cid
+
+    def _refit(self) -> None:
+        """Refit codebook from all stored embeddings."""
+        all_ids = []
+        all_embs = []
+        for bucket in self.buckets.values():
+            all_ids.extend(bucket.item_ids)
+            all_embs.extend(bucket.embeddings)
+        if len(all_embs) >= self.K_MIN:
+            emb_matrix = np.array(all_embs, dtype=np.float32)
+            self.fit(emb_matrix, all_ids)
+
+    def get_bucket_items(self, centroid_id: int) -> CentroidBucket:
+        return self.buckets.get(centroid_id, CentroidBucket())
+
+    # ─── Online Adaptation ────────────────────────────────────────────────
+
+    EMA_ETA = 0.01          # EMA learning rate
+    SPLIT_VARIANCE_THRESHOLD = 0.15  # bucket variance above this → split
+    PRUNE_IDLE_SECONDS = 3600        # 1 hour idle → prune candidate
+    UPDATE_INTERVAL = 50             # run adaptation every N queries
+
+    def ema_update(self, centroid_id: int, new_embedding: np.ndarray,
+                   eta: float = None) -> None:
+        """EMA update: centroid drifts toward new data point."""
+        if not self.fitted or centroid_id >= self.K:
+            return
+        eta = eta or self.EMA_ETA
+        old = self.centroids[centroid_id]
+        updated = (1 - eta) * old + eta * new_embedding
+        # Renormalize
+        norm = np.linalg.norm(updated)
+        if norm > 1e-10:
+            updated = updated / norm
+        self.centroids[centroid_id] = updated.astype(np.float32)
+
+        # Track access
+        if centroid_id in self.buckets:
+            self.buckets[centroid_id].last_accessed = time.time()
+            self.buckets[centroid_id].access_count += 1
+
+    def bucket_variance(self, centroid_id: int) -> float:
+        """Compute intra-bucket variance (mean squared distance to centroid)."""
+        bucket = self.buckets.get(centroid_id)
+        if not bucket or len(bucket.embeddings) < 2:
+            return 0.0
+        centroid = self.centroids[centroid_id]
+        embs = np.array(bucket.embeddings, dtype=np.float32)
+        dists = 1.0 - embs @ centroid  # cosine distance
+        return float(np.mean(dists ** 2))
+
+    def maybe_split(self, centroid_id: int) -> bool:
+        """Split a high-variance bucket into 2 if K < K_MAX.
+
+        Returns True if split occurred.
+        """
+        if self.K >= self.K_MAX:
+            return False
+        variance = self.bucket_variance(centroid_id)
+        if variance < self.SPLIT_VARIANCE_THRESHOLD:
+            return False
+
+        bucket = self.buckets.get(centroid_id)
+        if not bucket or len(bucket.embeddings) < 4:
+            return False
+
+        # Mini k-means on bucket items
+        from sklearn.cluster import KMeans
+        embs = np.array(bucket.embeddings, dtype=np.float32)
+        km = KMeans(n_clusters=2, random_state=42, n_init=3)
+        labels = km.fit_predict(embs)
+
+        # Create new centroids
+        c0 = km.cluster_centers_[0]
+        c1 = km.cluster_centers_[1]
+        c0 = c0 / max(np.linalg.norm(c0), 1e-10)
+        c1 = c1 / max(np.linalg.norm(c1), 1e-10)
+
+        # Replace original centroid with c0, add c1 as new
+        self.centroids[centroid_id] = c0.astype(np.float32)
+        new_cid = self.K  # new centroid ID
+        self.centroids = np.vstack([self.centroids, c1.astype(np.float32).reshape(1, -1)])
+
+        # Redistribute items
+        new_bucket_0 = CentroidBucket()
+        new_bucket_1 = CentroidBucket()
+        for i, label in enumerate(labels):
+            if label == 0:
+                new_bucket_0.item_ids.append(bucket.item_ids[i])
+                new_bucket_0.embeddings.append(bucket.embeddings[i])
+            else:
+                new_bucket_1.item_ids.append(bucket.item_ids[i])
+                new_bucket_1.embeddings.append(bucket.embeddings[i])
+
+        self.buckets[centroid_id] = new_bucket_0
+        self.buckets[new_cid] = new_bucket_1
+
+        logger.info(f"Split centroid {centroid_id}: variance={variance:.4f}, "
+                     f"K={self.K}, sizes={len(new_bucket_0.item_ids)}/{len(new_bucket_1.item_ids)}")
+        return True
+
+    def maybe_prune(self, centroid_id: int) -> bool:
+        """Prune idle centroid: reassign items to nearest neighbors.
+
+        Returns True if pruned.
+        """
+        if self.K <= self.K_MIN:
+            return False
+
+        bucket = self.buckets.get(centroid_id)
+        if not bucket:
+            return False
+
+        idle_time = time.time() - bucket.last_accessed
+        if idle_time < self.PRUNE_IDLE_SECONDS:
+            return False
+
+        # Reassign items to nearest remaining centroids
+        for i, emb in enumerate(bucket.embeddings):
+            dots = emb @ self.centroids.T
+            # Exclude the centroid being pruned
+            dots[centroid_id] = -np.inf
+            new_cid = int(np.argmax(dots))
+            if new_cid not in self.buckets:
+                self.buckets[new_cid] = CentroidBucket()
+            self.buckets[new_cid].item_ids.append(bucket.item_ids[i])
+            self.buckets[new_cid].embeddings.append(emb)
+
+        # Remove centroid
+        del self.buckets[centroid_id]
+        # Rebuild centroids array without the pruned one
+        mask = np.ones(self.K, dtype=bool)
+        mask[centroid_id] = False
+        self.centroids = self.centroids[mask]
+        # Remap bucket keys
+        self._remap_buckets_after_prune(centroid_id)
+
+        logger.info(f"Pruned centroid {centroid_id}: idle={idle_time:.0f}s, K={self.K}")
+        return True
+
+    def _remap_buckets_after_prune(self, removed_id: int) -> None:
+        """Remap bucket IDs after removing a centroid."""
+        new_buckets = {}
+        for old_id, bucket in sorted(self.buckets.items()):
+            if old_id < removed_id:
+                new_buckets[old_id] = bucket
+            elif old_id > removed_id:
+                new_buckets[old_id - 1] = bucket
+        self.buckets = new_buckets
+
+    def run_adaptation(self, force: bool = False) -> Dict:
+        """Run a full adaptation cycle: EMA already applied per-query,
+        this handles splits and prunes.
+
+        Returns dict of operations performed.
+        """
+        ops = {"splits": 0, "prunes": 0, "k_before": self.K}
+
+        if not self.fitted:
+            ops["k_after"] = self.K
+            return ops
+
+        # Check splits on high-variance buckets
+        cids_to_check = list(self.buckets.keys())
+        for cid in cids_to_check:
+            if cid < self.K and self.maybe_split(cid):
+                ops["splits"] += 1
+
+        # Check prunes on idle buckets (only during explicit adaptation)
+        if force:
+            cids_to_check = list(self.buckets.keys())
+            for cid in cids_to_check:
+                if cid < self.K and self.maybe_prune(cid):
+                    ops["prunes"] += 1
+
+        ops["k_after"] = self.K
+        return ops
+
+    def get_state(self) -> Tuple[Optional[np.ndarray], Dict]:
+        """For persistence: return (centroids, {cid: [item_ids]})."""
+        assignments = {}
+        for cid, bucket in self.buckets.items():
+            assignments[cid] = bucket.item_ids
+        return self.centroids, assignments

rm/drift.py

@@ -0,0 +1,145 @@
+"""Drift Monitor — Detects distribution shift via quantization error tracking.
+
+Monitors rolling qerr and margin to detect when incoming queries diverge
+from the codebook's trained distribution. When drift is detected:
+- Triggers aggressive codebook adaptation (splits/updates)
+- Logs alarm for observability
+"""
+
+import logging
+import time
+from collections import deque
+from dataclasses import dataclass, field
+from typing import Deque, List, Optional
+
+logger = logging.getLogger("rm.drift")
+
+
+@dataclass
+class DriftAlarm:
+    """Record of a drift event."""
+    timestamp: float
+    ratio: float        # qerr_recent / qerr_baseline
+    episode: int
+    message: str
+
+
+class DriftMonitor:
+    """Tracks rolling quantization error to detect distribution drift.
+
+    Alarm triggers when recent qerr exceeds baseline by a configurable factor.
+    """
+
+    WINDOW_SIZE = 20          # rolling window for recent stats
+    BASELINE_WINDOW = 50      # initial baseline window
+    ALARM_RATIO = 1.3         # qerr ratio threshold for alarm
+    COOLDOWN_QUERIES = 30     # min queries between alarms
+
+    def __init__(self):
+        self._qerr_history: Deque[float] = deque(maxlen=2000)
+        self._margin_history: Deque[float] = deque(maxlen=2000)
+        self._baseline_qerr: Optional[float] = None
+        self._baseline_margin: Optional[float] = None
+        self._episode = 0
+        self._last_alarm_episode = -self.COOLDOWN_QUERIES
+        self._alarms: List[DriftAlarm] = []
+        self._in_drift = False
+
+    @property
+    def episode_count(self) -> int:
+        return self._episode
+
+    @property
+    def in_drift(self) -> bool:
+        return self._in_drift
+
+    @property
+    def alarms(self) -> List[DriftAlarm]:
+        return list(self._alarms)
+
+    @property
+    def baseline_qerr(self) -> Optional[float]:
+        return self._baseline_qerr
+
+    @property
+    def recent_qerr(self) -> float:
+        if len(self._qerr_history) < 5:
+            return 0.0
+        recent = list(self._qerr_history)[-self.WINDOW_SIZE:]
+        return sum(recent) / len(recent)
+
+    @property
+    def drift_ratio(self) -> float:
+        if self._baseline_qerr is None or self._baseline_qerr < 1e-10:
+            return 0.0
+        return self.recent_qerr / self._baseline_qerr
+
+    def record(self, qerr: float, margin: float = 0.0) -> Optional[DriftAlarm]:
+        """Record a query's quantization error and margin.
+
+        Returns DriftAlarm if alarm was triggered, None otherwise.
+        """
+        self._episode += 1
+        self._qerr_history.append(qerr)
+        self._margin_history.append(margin)
+
+        # Establish baseline from first N queries
+        if self._baseline_qerr is None and len(self._qerr_history) >= self.BASELINE_WINDOW:
+            baseline_data = list(self._qerr_history)[:self.BASELINE_WINDOW]
+            self._baseline_qerr = sum(baseline_data) / len(baseline_data)
+            baseline_margins = list(self._margin_history)[:self.BASELINE_WINDOW]
+            self._baseline_margin = sum(baseline_margins) / len(baseline_margins)
+            return None
+
+        # Check for drift
+        return self._check_alarm()
+
+    def _check_alarm(self) -> Optional[DriftAlarm]:
+        """Check if current qerr indicates drift."""
+        if self._baseline_qerr is None:
+            return None
+        if len(self._qerr_history) < self.WINDOW_SIZE:
+            return None
+
+        ratio = self.drift_ratio
+
+        if ratio >= self.ALARM_RATIO:
+            if not self._in_drift:
+                self._in_drift = True
+            # Rate-limit alarms
+            if self._episode - self._last_alarm_episode >= self.COOLDOWN_QUERIES:
+                alarm = DriftAlarm(
+                    timestamp=time.time(),
+                    ratio=ratio,
+                    episode=self._episode,
+                    message=f"DRIFT_ALARM: qerr {ratio:.1f}x baseline at episode {self._episode}",
+                )
+                self._alarms.append(alarm)
+                self._last_alarm_episode = self._episode
+                logger.warning(f"[rm] {alarm.message}")
+                return alarm
+        else:
+            if self._in_drift:
+                self._in_drift = False
+                logger.info(f"[rm] Drift recovered at episode {self._episode}, "
+                            f"ratio={ratio:.2f}")
+
+        return None
+
+    def reset_baseline(self) -> None:
+        """Reset baseline using recent data (after adaptation)."""
+        if len(self._qerr_history) >= self.WINDOW_SIZE:
+            recent = list(self._qerr_history)[-self.WINDOW_SIZE:]
+            self._baseline_qerr = sum(recent) / len(recent)
+            self._in_drift = False
+            logger.info(f"[rm] Baseline reset: new baseline_qerr={self._baseline_qerr:.4f}")
+
+    def get_stats(self) -> dict:
+        return {
+            "episode": self._episode,
+            "baseline_qerr": self._baseline_qerr,
+            "recent_qerr": self.recent_qerr,
+            "drift_ratio": self.drift_ratio,
+            "in_drift": self._in_drift,
+            "alarm_count": len(self._alarms),
+        }

rm/embeddings/__init__.py

@@ -0,0 +1,4 @@
+from .local import LocalEmbeddings
+from .base import EmbeddingBackend
+
+__all__ = ["LocalEmbeddings", "EmbeddingBackend"]

rm/embeddings/base.py

@@ -0,0 +1,22 @@
+"""Abstract embedding backend interface."""
+
+from abc import ABC, abstractmethod
+from typing import List
+import numpy as np
+
+
+class EmbeddingBackend(ABC):
+    """Abstract base for embedding providers."""
+
+    @property
+    @abstractmethod
+    def dim(self) -> int:
+        """Embedding dimensionality."""
+
+    @abstractmethod
+    def encode(self, text: str) -> np.ndarray:
+        """Encode a single text to a normalized vector."""
+
+    @abstractmethod
+    def encode_batch(self, texts: List[str]) -> np.ndarray:
+        """Encode a batch of texts. Returns (N, dim) matrix."""

rm/embeddings/local.py

@@ -0,0 +1,56 @@
+"""Local sentence-transformers embedding backend."""
+
+from typing import Dict, List
+import numpy as np
+
+from .base import EmbeddingBackend
+
+_model_instance = None
+
+
+def _get_model(model_name: str):
+    """Lazy singleton for SentenceTransformer model."""
+    global _model_instance
+    if _model_instance is None:
+        from sentence_transformers import SentenceTransformer
+        _model_instance = SentenceTransformer(model_name)
+    return _model_instance
+
+
+class LocalEmbeddings(EmbeddingBackend):
+    """sentence-transformers wrapper with caching. Default: all-MiniLM-L6-v2 (d=384)."""
+
+    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
+        self.model_name = model_name
+        self._cache: Dict[str, np.ndarray] = {}
+        self._dim: int = 384  # MiniLM default
+
+    @property
+    def dim(self) -> int:
+        return self._dim
+
+    def encode(self, text: str) -> np.ndarray:
+        if text in self._cache:
+            return self._cache[text]
+        model = _get_model(self.model_name)
+        vec = model.encode(text, normalize_embeddings=True)
+        vec = np.asarray(vec, dtype=np.float32)
+        self._dim = vec.shape[0]
+        self._cache[text] = vec
+        return vec
+
+    def encode_batch(self, texts: List[str]) -> np.ndarray:
+        # Check cache for all
+        uncached = [(i, t) for i, t in enumerate(texts) if t not in self._cache]
+
+        if uncached:
+            model = _get_model(self.model_name)
+            uncached_texts = [t for _, t in uncached]
+            vecs = model.encode(uncached_texts, normalize_embeddings=True, batch_size=64)
+            vecs = np.asarray(vecs, dtype=np.float32)
+            self._dim = vecs.shape[1]
+            for (_, text), vec in zip(uncached, vecs):
+                self._cache[text] = vec
+
+        result = np.array([self._cache[t] for t in texts], dtype=np.float32)
+        return result

rm/filtering.py

@@ -0,0 +1,15 @@
+"""Score-based token filtering (τ threshold)."""
+
+from typing import List
+from .retrieval import RetrievalResult
+
+
+def filter_by_score(results: List[RetrievalResult], tau: float = 0.5) -> List[RetrievalResult]:
+    """Filter retrieval results below score threshold τ."""
+    return [r for r in results if r.score >= tau]
+
+
+def filter_top_n(results: List[RetrievalResult], n: int = 3) -> List[RetrievalResult]:
+    """Keep only top-n results by score."""
+    sorted_results = sorted(results, key=lambda r: r.score, reverse=True)
+    return sorted_results[:n]