differix 0.1.0a1__py3-none-any.whl

memory/manager.py

@@ -0,0 +1,186 @@
+"""
+Unified MemoryManager: orchestrates SQLiteMemoryStorage + A-Score ranker,
+exposes the feedback loop (trigger / helpful / violation / promote).
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+from memory.ranker import rank_memory_nodes_scored
+from memory.schemas import (
+    MemoryNode,
+    MemoryScope,
+    MemoryStorageError,
+    RankingWeights,
+    ScoredMemoryNode,
+)
+from memory.storage import SQLiteMemoryStorage
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryManager:
+    """High-level façade for lesson / memory node management."""
+
+    def __init__(
+        self,
+        storage: Optional[SQLiteMemoryStorage] = None,
+        db_path: str = ":memory:",
+        expected_dim: Optional[int] = None,
+        default_weights: Optional[RankingWeights] = None,
+    ) -> None:
+        if storage is not None:
+            if not isinstance(storage, SQLiteMemoryStorage):
+                raise MemoryStorageError(
+                    f"storage must be a SQLiteMemoryStorage, "
+                    f"got {type(storage).__name__}"
+                )
+            self.storage: SQLiteMemoryStorage = storage
+        else:
+            self.storage = SQLiteMemoryStorage(
+                db_path=db_path, expected_dim=expected_dim
+            )
+        self.weights: RankingWeights = default_weights or RankingWeights()
+
+    # ------------------------------------------------------------------ #
+    # Writes
+    # ------------------------------------------------------------------ #
+    def add(self, node: MemoryNode) -> None:
+        self.storage.insert(node)
+
+    def add_batch(self, nodes: List[MemoryNode]) -> None:
+        self.storage.insert_batch(nodes)
+
+    def create(
+        self,
+        content: str,
+        vector: List[float],
+        metadata: Optional[Dict[str, Any]] = None,
+        node_id: Optional[str] = None,
+        base_trust: float = 0.3,
+        is_critical: bool = False,
+        scope_type: MemoryScope = MemoryScope.GLOBAL,
+        scope_value: str = "*",
+        created_at: Optional[float] = None,
+    ) -> MemoryNode:
+        now = time.time()
+        node = MemoryNode(
+            node_id=node_id if node_id is not None else str(uuid.uuid4()),
+            vector=vector,
+            content=content,
+            metadata=metadata or {},
+            base_trust=base_trust,
+            is_critical=is_critical,
+            scope_type=scope_type,
+            scope_value=scope_value,
+            created_at=created_at if created_at is not None else now,
+            last_used_at=created_at if created_at is not None else now,
+        )
+        self.storage.insert(node)
+        return node
+
+    # ------------------------------------------------------------------ #
+    # Reads
+    # ------------------------------------------------------------------ #
+    def get(self, node_id: str) -> MemoryNode:
+        return self.storage.get(node_id)
+
+    def list_all(self) -> List[MemoryNode]:
+        return self.storage.list_all()
+
+    def count(self) -> int:
+        return self.storage.count()
+
+    # ------------------------------------------------------------------ #
+    # Retrieval
+    # ------------------------------------------------------------------ #
+    def search(
+        self,
+        query_vector: List[float],
+        limit: int = 10,
+        query_scope: Optional[MemoryScope] = None,
+        query_scope_value: Optional[str] = None,
+        weights: Optional[RankingWeights] = None,
+        candidate_pool: Optional[int] = None,
+        record_trigger: bool = True,
+        reference_time: Optional[float] = None,
+    ) -> List[ScoredMemoryNode]:
+        if limit <= 0:
+            return []
+        pool = candidate_pool if candidate_pool is not None else max(limit * 4, 100)
+        candidates = self.storage.list_all()
+        if not candidates:
+            return []
+
+        eff_weights = weights or self.weights
+        scored = rank_memory_nodes_scored(
+            nodes=candidates,
+            query_vector=query_vector,
+            weights=eff_weights,
+            query_scope=query_scope,
+            query_scope_value=query_scope_value,
+            top_k=pool,
+            reference_time=reference_time,
+        )
+        top = scored[:limit]
+
+        if record_trigger and top:
+            ids = [s.node.node_id for s in top]
+            ts = time.time() if reference_time is None else float(reference_time)
+            try:
+                self.storage.record_trigger(ids, touched_at=ts)
+                for s in top:
+                    s.node.trigger_count += 1
+                    s.node.last_used_at = ts
+            except MemoryStorageError as exc:
+                logger.warning(
+                    "Failed to persist trigger feedback for %d nodes: %s",
+                    len(ids), exc,
+                )
+        return top
+
+    # ------------------------------------------------------------------ #
+    # Feedback loop
+    # ------------------------------------------------------------------ #
+    def record_trigger(
+        self, node_ids: List[str], touched_at: Optional[float] = None
+    ) -> int:
+        """Increment T (trigger_count) and refresh last_used_at for each node."""
+        return self.storage.record_trigger(node_ids, touched_at=touched_at)
+
+    def record_helpful(self, node_ids: List[str]) -> int:
+        """Increment H (helpful_count) for each node."""
+        return self.storage.record_helpful(node_ids)
+
+    def record_violation(self, node_ids: List[str]) -> int:
+        """Increment V (violation_count) for each node."""
+        return self.storage.record_violation(node_ids)
+
+    def promote(self, node_id: str, base_trust: float = 1.0) -> None:
+        """Human promotion: raise B and mark critical (decay-resistant)."""
+        self.storage.promote(node_id, base_trust=base_trust)
+
+    # ------------------------------------------------------------------ #
+    # Deletes / lifecycle
+    # ------------------------------------------------------------------ #
+    def delete(self, node_id: str) -> bool:
+        return self.storage.delete(node_id)
+
+    def clear(self) -> None:
+        self.storage.clear()
+
+    def close(self) -> None:
+        self.storage.close()
+
+    def __enter__(self) -> "MemoryManager":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.close()
+
+
+__all__ = ["MemoryManager"]

memory/migrator.py

@@ -0,0 +1,253 @@
+"""
+Vector migration across embedding provider changes.
+
+When the active provider's signature differs from what's recorded in the
+database (or when no signature is recorded but rows already exist), existing
+vectors are incompatible with new queries. This module re-embeds every
+stored node from its `content` field and swaps them atomically, preserving
+all A-Score signals (base_trust, trigger_count, helpful_count,
+violation_count, is_critical, scope, timestamps).
+
+Public API:
+    check_and_migrate(storage, provider) -> MigrationReport
+
+Design:
+  - Re-embedding happens BEFORE any write, so a failing provider leaves
+    the DB untouched (the caller can retry with a different provider).
+  - The swap uses `storage.replace_all()` which is transactional at the
+    SQLite layer.
+  - If there are existing rows and no signature, we MUST re-embed — the
+    vector dimensions of pre-signature rows are unknown. This closes the
+    failure mode where a fresh schema_meta table caused old vectors to be
+    silently accepted as if they matched a new provider.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from dataclasses import dataclass
+from typing import List, Optional, Tuple
+
+from memory.embeddings import BaseEmbeddingProvider
+from memory.schemas import MemoryNode, MemoryStorageError
+from memory.storage import SQLiteMemoryStorage
+
+logger = logging.getLogger(__name__)
+
+
+_SIGNATURE_KEY = "embedding_signature"
+
+
+@dataclass
+class MigrationReport:
+    action: str  # "init" | "none" | "migrated" | "aborted"
+    rows_migrated: int = 0
+    old_signature: Optional[Tuple[str, str, int]] = None
+    new_signature: Optional[Tuple[str, str, int]] = None
+    elapsed_ms: float = 0.0
+    error: Optional[str] = None
+
+    def __str__(self) -> str:
+        if self.action == "none":
+            return f"none (already {self.new_signature})"
+        if self.action == "init":
+            return f"init → {self.new_signature}"
+        if self.action == "migrated":
+            return (
+                f"migrated {self.rows_migrated} rows: "
+                f"{self.old_signature} → {self.new_signature} "
+                f"in {self.elapsed_ms:.0f}ms"
+            )
+        return f"aborted: {self.error}"
+
+
+def _serialize_signature(sig: Tuple[str, str, int]) -> str:
+    return f"{sig[0]}::{sig[1]}::{sig[2]}"
+
+
+def _deserialize_signature(raw: str) -> Optional[Tuple[str, str, int]]:
+    try:
+        parts = raw.split("::")
+        if len(parts) != 3:
+            return None
+        return (parts[0], parts[1], int(parts[2]))
+    except (ValueError, IndexError):
+        return None
+
+
+def check_and_migrate(
+    storage: SQLiteMemoryStorage,
+    provider: BaseEmbeddingProvider,
+) -> MigrationReport:
+    """
+    Ensure the storage contains vectors compatible with `provider`.
+
+    Four outcomes:
+      init      — fresh DB with no rows, signature recorded, no re-embed
+      none      — DB signature matches active provider, no-op
+      migrated  — rows re-embedded and swapped atomically
+      aborted   — re-embedding failed; DB untouched, next call retries
+    """
+    t0 = time.perf_counter()
+    new_sig = provider.signature
+
+    raw = storage.get_meta(_SIGNATURE_KEY)
+    stored_sig = _deserialize_signature(raw) if raw else None
+
+    # ---- Case 1: fresh DB with no rows and no signature ------------------
+    if stored_sig is None:
+        try:
+            row_count = storage.count()
+        except MemoryStorageError as exc:
+            return MigrationReport(
+                action="aborted",
+                new_signature=new_sig,
+                error=f"Failed to count rows: {exc}",
+                elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+            )
+
+        if row_count == 0:
+            storage.set_meta(_SIGNATURE_KEY, _serialize_signature(new_sig))
+            return MigrationReport(
+                action="init",
+                new_signature=new_sig,
+                elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+            )
+
+        # ---- Case 2: rows exist but no signature was ever recorded --------
+        # The rows' vector dimensions are unknown; we cannot trust them to
+        # match the new provider. Re-embed all of them.
+        logger.warning(
+            "Found %d pre-existing rows with no recorded signature. "
+            "Re-embedding all rows with the active provider (%s).",
+            row_count,
+            new_sig,
+        )
+        return _reembed_and_swap(
+            storage=storage,
+            provider=provider,
+            new_sig=new_sig,
+            old_sig=None,
+            reason="no-signature-with-rows",
+            start_time=t0,
+        )
+
+    # ---- Case 3: signature matches — no-op ------------------------------
+    if stored_sig == new_sig:
+        return MigrationReport(
+            action="none",
+            old_signature=stored_sig,
+            new_signature=new_sig,
+            elapsed_ms=(time.perf_counter() - t0) * 1000.0,
+        )
+
+    # ---- Case 4: signature mismatch — re-embed --------------------------
+    logger.warning(
+        "Embedding signature changed: %s → %s. Re-embedding all rows.",
+        stored_sig,
+        new_sig,
+    )
+    return _reembed_and_swap(
+        storage=storage,
+        provider=provider,
+        new_sig=new_sig,
+        old_sig=stored_sig,
+        reason="signature-mismatch",
+        start_time=t0,
+    )
+
+
+def _reembed_and_swap(
+    storage: SQLiteMemoryStorage,
+    provider: BaseEmbeddingProvider,
+    new_sig: Tuple[str, str, int],
+    old_sig: Optional[Tuple[str, str, int]],
+    reason: str,
+    start_time: float,
+) -> MigrationReport:
+    """
+    Shared re-embedding logic for both mismatch and pre-existing-rows cases.
+    """
+    try:
+        existing_nodes: List[MemoryNode] = storage.list_all()
+    except MemoryStorageError as exc:
+        return MigrationReport(
+            action="aborted",
+            old_signature=old_sig,
+            new_signature=new_sig,
+            error=f"Failed to read existing nodes: {exc}",
+            elapsed_ms=(time.perf_counter() - start_time) * 1000.0,
+        )
+
+    if not existing_nodes:
+        # Rows disappeared between count and list — treat as init.
+        storage.set_meta(_SIGNATURE_KEY, _serialize_signature(new_sig))
+        return MigrationReport(
+            action="init",
+            old_signature=old_sig,
+            new_signature=new_sig,
+            elapsed_ms=(time.perf_counter() - start_time) * 1000.0,
+        )
+
+    # Re-embed everything BEFORE touching the DB. A failing provider must
+    # not leave the storage in a half-migrated state.
+    try:
+        contents = [n.content for n in existing_nodes]
+        new_vectors = provider.embed_batch(contents)
+        if len(new_vectors) != len(existing_nodes):
+            raise RuntimeError(
+                f"Provider returned {len(new_vectors)} vectors "
+                f"for {len(existing_nodes)} inputs"
+            )
+        # Validate each returned vector has the expected dimension
+        for i, vec in enumerate(new_vectors):
+            if not isinstance(vec, list) or len(vec) != provider.dim:
+                raise RuntimeError(
+                    f"Provider returned vector {i} with shape "
+                    f"{len(vec) if isinstance(vec, list) else type(vec).__name__}, "
+                    f"expected {provider.dim}"
+                )
+    except Exception as exc:
+        logger.error("Re-embedding failed (%s): %s", reason, exc)
+        return MigrationReport(
+            action="aborted",
+            old_signature=old_sig,
+            new_signature=new_sig,
+            error=f"Re-embedding failed: {exc}",
+            elapsed_ms=(time.perf_counter() - start_time) * 1000.0,
+        )
+
+    # Build the new node set — only `vector` changes.
+    migrated_nodes: List[MemoryNode] = []
+    for node, new_vec in zip(existing_nodes, new_vectors):
+        migrated_nodes.append(
+            node.model_copy(update={"vector": list(new_vec)})
+        )
+
+    # Atomic swap. Update the storage's dimension constraint first so
+    # `_validate_vector` accepts the new length.
+    storage.set_expected_dim(provider.dim)
+    try:
+        storage.replace_all(migrated_nodes)
+        storage.set_meta(_SIGNATURE_KEY, _serialize_signature(new_sig))
+    except MemoryStorageError as exc:
+        # Best effort: keep the DB as-is. The next run retries.
+        return MigrationReport(
+            action="aborted",
+            old_signature=old_sig,
+            new_signature=new_sig,
+            error=f"Atomic swap failed: {exc}",
+            elapsed_ms=(time.perf_counter() - start_time) * 1000.0,
+        )
+
+    return MigrationReport(
+        action="migrated",
+        rows_migrated=len(migrated_nodes),
+        old_signature=old_sig,
+        new_signature=new_sig,
+        elapsed_ms=(time.perf_counter() - start_time) * 1000.0,
+    )
+
+
+__all__ = ["check_and_migrate", "MigrationReport"]

memory/ranker.py

@@ -0,0 +1,203 @@
+"""
+Activation Score (A-Score) ranking for MemoryNode candidates.
+
+    A = [ (B · (1 + H) · Sim^α · ScopeMatch) / (T + 2 + V) ]
+        · e^(−λ_eff · Δt)  +  ε
+
+    λ_eff = λ_base / (1 + β · 1_critical)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import math
+import time
+from typing import List, Optional, Sequence, Tuple
+
+from memory.schemas import (
+    DimensionMismatchError,
+    MemoryNode,
+    MemoryScope,
+    MemoryStorageError,
+    RankingWeights,
+    ScoredMemoryNode,
+)
+
+
+def _validate_vector(vector: Sequence[float], label: str) -> None:
+    if not vector:
+        raise MemoryStorageError(f"{label} must be non-empty.")
+    for idx, value in enumerate(vector):
+        if not isinstance(value, (int, float)) or isinstance(value, bool):
+            raise MemoryStorageError(
+                f"{label} element at index {idx} is not numeric: {value!r}"
+            )
+        if math.isnan(value) or math.isinf(value):
+            raise MemoryStorageError(
+                f"{label} contains NaN/Inf at index {idx}: {value!r}"
+            )
+
+
+def cosine_similarity(a: Sequence[float], b: Sequence[float]) -> float:
+    if len(a) != len(b):
+        raise DimensionMismatchError(len(a), len(b))
+    dot = 0.0
+    na = 0.0
+    nb = 0.0
+    for x, y in zip(a, b):
+        dot += x * y
+        na += x * x
+        nb += y * y
+    if na == 0.0 or nb == 0.0:
+        return 0.0
+    raw = dot / (math.sqrt(na) * math.sqrt(nb))
+    return max(-1.0, min(1.0, raw))
+
+
+def _normalized_similarity(a: Sequence[float], b: Sequence[float]) -> float:
+    return (cosine_similarity(a, b) + 1.0) / 2.0
+
+
+def _deterministic_jitter(node_id: str, scale: float) -> float:
+    if scale <= 0.0:
+        return 0.0
+    digest = hashlib.sha256(node_id.encode("utf-8")).digest()
+    bucket = int.from_bytes(digest[:4], "big") / 0xFFFFFFFF
+    return scale * (2.0 * bucket - 1.0)
+
+
+def _scope_match(
+    node_scope: MemoryScope,
+    node_value: str,
+    query_scope: Optional[MemoryScope],
+    query_value: Optional[str],
+) -> float:
+    if node_scope == MemoryScope.GLOBAL:
+        return 1.0
+    if query_scope is None:
+        return 1.0
+    if node_scope == query_scope and node_value == query_value:
+        return 1.0
+    return 0.5
+
+
+def activation_score(
+    node: MemoryNode,
+    query_vector: Sequence[float],
+    weights: Optional[RankingWeights] = None,
+    query_scope: Optional[MemoryScope] = None,
+    query_scope_value: Optional[str] = None,
+    reference_time: Optional[float] = None,
+) -> Tuple[float, float, float, float, float]:
+    if weights is None:
+        weights = RankingWeights()
+    if len(node.vector) != len(query_vector):
+        raise DimensionMismatchError(len(node.vector), len(query_vector))
+
+    sim = _normalized_similarity(query_vector, node.vector)
+    sim_term = sim ** float(weights.alpha)
+    sm = _scope_match(node.scope_type, node.scope_value, query_scope, query_scope_value)
+
+    numerator = (
+        float(node.base_trust)
+        * (1.0 + float(node.helpful_count))
+        * sim_term
+        * sm
+    )
+    denominator = float(node.trigger_count) + 2.0 + float(node.violation_count)
+
+    now = time.time() if reference_time is None else float(reference_time)
+    delta_hours = max(0.0, (now - float(node.last_used_at)) / 3600.0)
+    protection = 1.0 + weights.beta_protection * (1.0 if node.is_critical else 0.0)
+    lambda_eff = float(weights.lambda_base) / protection
+    decay = math.exp(-lambda_eff * delta_hours)
+
+    jitter = _deterministic_jitter(node.node_id, float(weights.epsilon))
+    score = max(0.0, (numerator / denominator) * decay + jitter)
+    return score, sim, sm, decay, jitter
+
+
+def _score_and_sort(
+    nodes: Sequence[MemoryNode],
+    query_vector: Sequence[float],
+    weights: RankingWeights,
+    query_scope: Optional[MemoryScope],
+    query_scope_value: Optional[str],
+    reference_time: Optional[float],
+) -> List[ScoredMemoryNode]:
+    scored: List[ScoredMemoryNode] = []
+    for node in nodes:
+        if not isinstance(node, MemoryNode):
+            raise MemoryStorageError(
+                f"ranker received non-MemoryNode: {type(node).__name__}"
+            )
+        score, sim, sm, decay, jitter = activation_score(
+            node=node,
+            query_vector=query_vector,
+            weights=weights,
+            query_scope=query_scope,
+            query_scope_value=query_scope_value,
+            reference_time=reference_time,
+        )
+        scored.append(
+            ScoredMemoryNode(
+                node=node,
+                score=score,
+                similarity=sim,
+                scope_match=sm,
+                decay=decay,
+                jitter=jitter,
+            )
+        )
+    scored.sort(key=lambda s: (-s.score, s.node.node_id))
+    return scored
+
+
+def rank_memory_nodes(
+    nodes: List[MemoryNode],
+    query_vector: List[float],
+    weights: Optional[RankingWeights] = None,
+    query_scope: Optional[MemoryScope] = None,
+    query_scope_value: Optional[str] = None,
+    top_k: Optional[int] = None,
+    reference_time: Optional[float] = None,
+) -> List[MemoryNode]:
+    if not nodes:
+        return []
+    _validate_vector(query_vector, "query_vector")
+    eff = weights or RankingWeights()
+    scored = _score_and_sort(
+        nodes, query_vector, eff, query_scope, query_scope_value, reference_time
+    )
+    if top_k is not None:
+        scored = scored[: max(0, top_k)]
+    return [s.node for s in scored]
+
+
+def rank_memory_nodes_scored(
+    nodes: List[MemoryNode],
+    query_vector: List[float],
+    weights: Optional[RankingWeights] = None,
+    query_scope: Optional[MemoryScope] = None,
+    query_scope_value: Optional[str] = None,
+    top_k: Optional[int] = None,
+    reference_time: Optional[float] = None,
+) -> List[ScoredMemoryNode]:
+    if not nodes:
+        return []
+    _validate_vector(query_vector, "query_vector")
+    eff = weights or RankingWeights()
+    scored = _score_and_sort(
+        nodes, query_vector, eff, query_scope, query_scope_value, reference_time
+    )
+    if top_k is not None:
+        scored = scored[: max(0, top_k)]
+    return scored
+
+
+__all__ = [
+    "cosine_similarity",
+    "activation_score",
+    "rank_memory_nodes",
+    "rank_memory_nodes_scored",
+]