memplex 3.2.0__py3-none-any.whl

memnex/retrieval/dedup.py

@@ -0,0 +1,386 @@
+"""MemoryDeduplicator -- exact + semantic deduplication for Compaction.
+
+Dedup strategies (by scale, automatic fallback)::
+
+    1. FAISS IVF ANN     -- O(n log n), cross-domain, requires faiss-cpu
+    2. NumPy matrix      -- O(n*d), grouped by domain, requires numpy
+    3. Pure Python       -- O(n^2), zero dependencies, Lite fallback
+
+Usage::
+
+    dedup = MemoryDeduplicator(embedding_service)
+    result = dedup.deduplicate(memories)
+    print(result.exact_removed, result.semantic_removed)
+"""
+
+from __future__ import annotations
+
+import copy
+import hashlib
+import logging
+import math
+from enum import Enum
+from typing import Dict, List, Optional
+
+from memnex.retrieval.embedding import EmbeddingService
+from memnex.models import DedupResult, Memory
+
+logger = logging.getLogger(__name__)
+
+
+# ── Dedup strategy ────────────────────────────────────────────────────
+
+
+class DedupStrategy(Enum):
+    """Deduplication strategy."""
+
+    EXACT = "exact"        # exact content hash match
+    SEMANTIC = "semantic"  # embedding cosine similarity
+    BOTH = "both"          # exact first, then semantic
+
+
+# ── MemoryDeduplicator ────────────────────────────────────────────────
+
+
+class MemoryDeduplicator:
+    """Memory deduplication cleaner.
+
+    Used in the Compaction Pipeline's Dedup stage.
+
+    Parameters
+    ----------
+    embedding_service:
+        Used to generate embeddings for semantic dedup.
+    strategy:
+        Which dedup strategy to apply.
+    threshold:
+        Cosine similarity threshold for semantic dedup (default 0.95).
+    chunk_threshold:
+        Max memories per NumPy chunk before domain-based splitting.
+    """
+
+    def __init__(
+        self,
+        embedding_service: EmbeddingService,
+        strategy: DedupStrategy = DedupStrategy.BOTH,
+        threshold: float = 0.95,
+        chunk_threshold: int = 20000,
+    ) -> None:
+        self.embedder = embedding_service
+        self.strategy = strategy
+        self.threshold = threshold
+        self.chunk_threshold = chunk_threshold
+
+        # Counters set at the start of each ``deduplicate`` call
+        self._original_count: int = 0
+        self._exact_removed: int = 0
+        self._semantic_removed: int = 0
+
+    # ── Public API ──────────────────────────────────────────────────
+
+
+    def deduplicate(self, memories: List[Memory]) -> DedupResult:
+        """Run dedup and return cleaned memories with statistics."""
+        self._original_count = len(memories)
+        self._exact_removed = 0
+        self._semantic_removed = 0
+
+        if self.strategy in (DedupStrategy.EXACT, DedupStrategy.BOTH):
+            memories = self._exact_dedup(memories)
+
+        if self.strategy in (DedupStrategy.SEMANTIC, DedupStrategy.BOTH):
+            memories = self._semantic_dedup(memories)
+
+        return DedupResult(
+            original_count=self._original_count,
+            final_count=len(memories),
+            exact_removed=self._exact_removed,
+            semantic_removed=self._semantic_removed,
+            deduplicated=memories,
+        )
+
+    # ── Exact dedup ─────────────────────────────────────────────────
+
+
+    def _exact_dedup(self, memories: List[Memory]) -> List[Memory]:
+        """Remove exact duplicates by content hash.
+
+        When two memories share the same hash, the *better* one is kept
+        (see :meth:`_choose_better`).
+        """
+        seen: Dict[str, Memory] = {}
+        for m in memories:
+            key = self._content_hash(m)
+            if key not in seen:
+                seen[key] = m
+            else:
+                seen[key] = self._choose_better(seen[key], m)
+                self._exact_removed += 1
+        return list(seen.values())
+
+    @staticmethod
+    def _choose_better(a: Memory, b: Memory) -> Memory:
+        """Pick the higher-quality memory to keep.
+
+        Priority: newer *updated_at* > more populated fields > higher confidence.
+        """
+        # updated_at comparison (they may be str or datetime)
+        a_updated = str(getattr(a, "updated_at", "") or "")
+        b_updated = str(getattr(b, "updated_at", "") or "")
+        if a_updated > b_updated:
+            return a
+        if b_updated > a_updated:
+            return b
+
+        # field completeness
+        a_fields = sum(
+            1 for role in ("trigger", "condition", "action", "benefit")
+            if getattr(a, role, [])
+        )
+        b_fields = sum(
+            1 for role in ("trigger", "condition", "action", "benefit")
+            if getattr(b, role, [])
+        )
+        if a_fields >= b_fields:
+            return a
+        return b
+
+    # ── Semantic dedup (dispatcher) ─────────────────────────────────
+
+
+    def _semantic_dedup(self, memories: List[Memory]) -> List[Memory]:
+        """Semantic dedup with automatic backend selection.
+
+        Routing:
+        1. FAISS IVF ANN (if *faiss* is importable) -- global, cross-domain.
+        2. NumPy matrix (if *numpy* is importable, n <= chunk_threshold)
+           -- grouped by domain when n > chunk_threshold.
+        3. Pure-Python pairwise -- O(n^2), zero dependencies.
+        """
+        # 1. FAISS (best)
+        try:
+            import faiss  # noqa: F401 -- check availability
+
+            return self._semantic_dedup_faiss(memories)
+        except ImportError:
+            pass
+
+        # 2. NumPy (good for medium scale)
+        if len(memories) > self.chunk_threshold:
+            # Split by domain to keep memory bounded
+            grouped: Dict[str, List[Memory]] = {}
+            for m in memories:
+                domain = getattr(m, "domain", "unknown") or "unknown"
+                grouped.setdefault(domain, []).append(m)
+            result: List[Memory] = []
+            for group in grouped.values():
+                result.extend(self._semantic_dedup_chunk(group))
+            return result
+
+        return self._semantic_dedup_chunk(memories)
+
+    # ── FAISS ANN dedup ─────────────────────────────────────────────
+
+
+    def _semantic_dedup_faiss(self, memories: List[Memory]) -> List[Memory]:
+        """FAISS IndexFlatIP with Union-Find clustering.
+
+        Steps:
+        1. Batch embed + L2-normalise (inner product == cosine).
+        2. Build IndexFlatIP, search top-K=5 neighbours.
+        3. Union-Find groups memories above *threshold*.
+        4. Merge each cluster into one representative memory.
+        """
+        import faiss  # type: ignore
+        import numpy as np  # type: ignore
+
+        if not memories:
+            return []
+
+        texts = [self._memory_to_text(m) for m in memories]
+        raw_vecs = np.array(
+            self.embedder.embed_batch(texts), dtype="float32"
+        )
+
+        # Normalise to unit vectors so IP == cosine
+        norms = np.linalg.norm(raw_vecs, axis=1, keepdims=True)
+        vecs = raw_vecs / (norms + 1e-8)
+
+        index = faiss.IndexFlatIP(vecs.shape[1])
+        index.add(vecs)
+
+        k = min(6, len(memories))  # top-1 is self; real neighbours = k-1
+        distances, indices = index.search(vecs, k)
+
+        # Union-Find
+        parent = list(range(len(memories)))
+
+        def find(x: int) -> int:
+            while parent[x] != x:
+                parent[x] = parent[parent[x]]
+                x = parent[x]
+            return x
+
+        def union(x: int, y: int) -> None:
+            parent[find(x)] = find(y)
+
+        for i in range(len(memories)):
+            for j_idx in range(1, k):
+                j = int(indices[i][j_idx])
+                if j < 0 or j == i:
+                    continue
+                if float(distances[i][j_idx]) > self.threshold:
+                    union(i, j)
+
+        # Collect clusters
+        clusters: Dict[int, List[int]] = {}
+        for i in range(len(memories)):
+            root = find(i)
+            clusters.setdefault(root, []).append(i)
+
+        result: List[Memory] = []
+        for idxs in clusters.values():
+            cluster_mems = [memories[i] for i in idxs]
+            merged = self._merge_memories(cluster_mems)
+            self._semantic_removed += len(idxs) - 1
+            result.append(merged)
+
+        return result
+
+    # ── NumPy matrix dedup ──────────────────────────────────────────
+
+
+    def _semantic_dedup_chunk(self, memories: List[Memory]) -> List[Memory]:
+        """Dedup one chunk using NumPy cosine similarity matrix."""
+        if not memories:
+            return []
+
+        texts = [self._memory_to_text(m) for m in memories]
+        embeddings = self.embedder.embed_batch(texts)
+
+        # Try numpy; fall back to pure Python
+        try:
+            import numpy as np  # type: ignore
+        except ImportError:
+            return self._semantic_dedup_fallback(memories, embeddings)
+
+        emb_matrix = np.array(embeddings, dtype="float64")
+        norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
+        emb_normalized = emb_matrix / (norms + 1e-8)
+        sim_matrix = emb_normalized @ emb_normalized.T
+
+        result: List[Memory] = []
+        used: set = set()
+
+        for i, m in enumerate(memories):
+            mid = m.id
+            if mid in used:
+                continue
+            similar = [m]
+            used.add(mid)
+
+            for j in range(i + 1, len(memories)):
+                mjid = memories[j].id
+                if mjid in used:
+                    continue
+                if sim_matrix[i][j] > self.threshold:
+                    similar.append(memories[j])
+                    used.add(mjid)
+                    self._semantic_removed += 1
+
+            merged = self._merge_memories(similar)
+            result.append(merged)
+
+        return result
+
+    # ── Pure-Python fallback dedup ──────────────────────────────────
+
+
+    def _semantic_dedup_fallback(
+        self, memories: List[Memory], embeddings: List[list]
+    ) -> List[Memory]:
+        """O(n^2) pairwise cosine similarity, zero dependencies."""
+        result: List[Memory] = []
+        used: set = set()
+
+        for i, m in enumerate(memories):
+            mid = m.id
+            if mid in used:
+                continue
+            similar = [m]
+            used.add(mid)
+
+            for j in range(i + 1, len(memories)):
+                mjid = memories[j].id
+                if mjid in used:
+                    continue
+                sim = self._cosine_sim(embeddings[i], embeddings[j])
+                if sim > self.threshold:
+                    similar.append(memories[j])
+                    used.add(mjid)
+                    self._semantic_removed += 1
+
+            result.append(self._merge_memories(similar))
+        return result
+
+    # ── Merge helpers ───────────────────────────────────────────────
+
+
+    @staticmethod
+    def _merge_memories(memories: List[Memory]) -> Memory:
+        """Merge a list of similar memories into one.
+
+        Strategy:
+        - Keep the newest *updated_at* as the base (deep-copy to avoid
+          mutating live objects in MemoryStore).
+        - Merge *source_paragraphs* and role field-values, deduplicating
+          by ``desc`` string.
+        """
+        if len(memories) == 1:
+            return memories[0]
+
+        base = copy.deepcopy(
+            max(memories, key=lambda m: str(getattr(m, "updated_at", "") or ""))
+        )
+
+        for m in memories:
+            if m.id == base.id:
+                continue
+            # Merge source_paragraphs
+            for sp in getattr(m, "source_paragraphs", []):
+                if sp not in base.source_paragraphs:
+                    base.source_paragraphs.append(sp)
+            # Merge role field-values
+            for role in ("trigger", "condition", "action", "benefit"):
+                existing_descs = {fv.desc for fv in getattr(base, role, [])}
+                for fv in getattr(m, role, []):
+                    if fv.desc not in existing_descs:
+                        getattr(base, role).append(fv)
+                        existing_descs.add(fv.desc)
+
+        return base
+
+    # ── Utility ─────────────────────────────────────────────────────
+
+
+    @staticmethod
+    def _cosine_sim(a: list, b: list) -> float:
+        """Pure-Python cosine similarity (no numpy)."""
+        dot = sum(x * y for x, y in zip(a, b))
+        norm_a = sum(x * x for x in a) ** 0.5
+        norm_b = sum(x * x for x in b) ** 0.5
+        return dot / (norm_a * norm_b + 1e-8)
+
+    def _content_hash(self, memory: Memory) -> str:
+        """SHA-256 of the memory's text representation."""
+        content = self._memory_to_text(memory)
+        return hashlib.sha256(content.encode()).hexdigest()
+
+    @staticmethod
+    def _memory_to_text(memory: Memory) -> str:
+        """Flatten a memory into a single text for comparison."""
+        parts = [memory.name]
+        for role in ("trigger", "condition", "action", "benefit"):
+            for fv in getattr(memory, role, []):
+                parts.append(fv.desc)
+        return " ".join(parts)

memnex/retrieval/embedding.py

@@ -0,0 +1,289 @@
+"""EmbeddingService -- vector embedding generation, storage and refresh.
+
+Supports multiple embedding models, configurable dimension, batch size,
+and Contextual Retrieval (Anthropic's document-context prefix injection).
+
+Embedding strategies::
+
+    NAME_ONLY    -- function name only
+    NAME_DOMAIN  -- name + domain
+    FULL         -- name + trigger + action + benefit
+    SEMANTIC     -- concise semantic summary for search
+
+Usage::
+
+    svc = EmbeddingService(model="default", storage=store, vector_store=vs)
+    vector = svc.embed("some text")
+    svc.refresh_all()
+"""
+
+from __future__ import annotations
+
+import logging
+from enum import Enum
+from typing import List, Optional, TYPE_CHECKING
+
+from memnex.models import Function, RefreshResult
+
+if TYPE_CHECKING:
+    from memnex.storage.base import MemoryStore
+    from memnex.storage.vector import VectorStore as VectorStoreProtocol
+
+logger = logging.getLogger(__name__)
+
+# Type alias for embedding vectors
+Vector = List[float]
+
+
+# ── Embedding strategies ────────────────────────────────────────────
+
+
+class EmbeddingStrategy(Enum):
+    """Controls how a Function is converted to embeddable text."""
+
+    NAME_ONLY = "name"
+    NAME_DOMAIN = "name_domain"
+    FULL = "full"
+    SEMANTIC = "semantic"
+
+
+# ── Embedder backends ───────────────────────────────────────────────
+
+
+class _SentenceTransformerEmbedder:
+    """Wraps ``sentence_transformers.SentenceTransformer``."""
+
+    def __init__(self, model_name: str, dimension: int) -> None:
+        from sentence_transformers import SentenceTransformer  # type: ignore
+
+        self._model = SentenceTransformer(model_name)
+        self.dimension = dimension
+
+    def encode(self, text: str) -> Vector:
+        return self._model.encode([text])[0].tolist()
+
+    def encode_batch(self, texts: List[str], batch_size: int = 32) -> List[Vector]:
+        embeddings = self._model.encode(texts, batch_size=batch_size)
+        return [e.tolist() for e in embeddings]
+
+
+class _SimpleTFIDFEmbedder:
+    """Fallback embedder when sentence-transformers is unavailable.
+
+    Uses a TF-IDF-inspired bag-of-words representation.  The dimension
+    is fixed to the number of unique words seen so far (padded / truncated
+    to *dimension*).
+    """
+
+    def __init__(self, dimension: int = 384) -> None:
+        self.dimension = dimension
+        self._vocab: dict = {}  # word -> index
+        self._idf: dict = {}   # word -> idf score
+        self._doc_count: int = 0
+
+    def encode(self, text: str) -> Vector:
+        words = text.lower().split()
+        if not words:
+            return [0.0] * self.dimension
+
+        # Update vocabulary
+        self._doc_count += 1
+        unique_words = set(words)
+        for w in unique_words:
+            if w not in self._vocab:
+                self._vocab[w] = len(self._vocab)
+            self._idf[w] = self._idf.get(w, 0) + 1
+
+        # TF-IDF vector
+        vec = [0.0] * self.dimension
+        tf = {}
+        for w in words:
+            tf[w] = tf.get(w, 0) + 1
+
+        for w, count in tf.items():
+            idx = self._vocab.get(w, -1) % self.dimension
+            idf = self._doc_count / (self._idf.get(w, 1) + 1)
+            import math
+            vec[idx] = (count / len(words)) * math.log(idf + 1)
+
+        # L2 normalize
+        norm = sum(x * x for x in vec) ** 0.5
+        if norm > 0:
+            vec = [x / norm for x in vec]
+        return vec
+
+    def encode_batch(self, texts: List[str], batch_size: int = 32) -> List[Vector]:
+        return [self.encode(t) for t in texts]
+
+
+# ── EmbeddingService ─────────────────────────────────────────────────
+
+
+class EmbeddingService:
+    """Vector embedding service.
+
+    Parameters
+    ----------
+    model:
+        Embedding model name.  ``"default"`` maps to
+        ``all-MiniLM-L6-v2``.  Falls back to TF-IDF when
+        sentence-transformers is not installed.
+    dimension:
+        Embedding vector dimension.
+    storage:
+        Optional :class:`MemoryStore` for ``refresh`` / ``refresh_all``.
+    vector_store:
+        Optional :class:`VectorStore` for upsert operations.
+    """
+
+    _MODEL_MAP = {
+        "default": "all-MiniLM-L6-v2",
+        "bge-m3": "BAAI/bge-m3",
+        "bge-small": "BAAI/bge-small-en-v1.5",
+    }
+
+    def __init__(
+        self,
+        model: str = "default",
+        dimension: int = 384,
+        storage: Optional["MemoryStore"] = None,
+        vector_store: Optional["VectorStoreProtocol"] = None,
+    ) -> None:
+        self.model = model
+        self.dimension = dimension
+        self.storage = storage
+        self.vector_store = vector_store
+        self._embedder = self._create_embedder(model, dimension)
+
+    # ── Public API ──────────────────────────────────────────────────
+
+    def embed(self, text: str) -> Vector:
+        """Generate an embedding vector for a single text."""
+        return self._embedder.encode(text)
+
+    def embed_batch(
+        self, texts: List[str], batch_size: int = 32
+    ) -> List[Vector]:
+        """Batch generate embedding vectors."""
+        return self._embedder.encode_batch(texts, batch_size=batch_size)
+
+    def embed_function(
+        self,
+        func: Function,
+        source: Optional[object] = None,
+        use_contextual: bool = True,
+    ) -> Vector:
+        """Generate an embedding for a Function.
+
+        When *source* is available and *use_contextual* is True, a document
+        context prefix is prepended (Contextual Retrieval).
+        """
+        content = self.function_to_text(func)
+        if use_contextual and source is not None:
+            origin = (
+                getattr(source, "url", None)
+                or str(getattr(source, "source_path", ""))
+                or "unknown"
+            )
+            content = (
+                f"[文档: {origin} | 领域: {func.domain or '未分类'}] {content}"
+            )
+        return self.embed(content)
+
+    def function_to_text(
+        self,
+        func: Function,
+        strategy: EmbeddingStrategy = EmbeddingStrategy.FULL,
+    ) -> str:
+        """Convert a Function to embeddable text per *strategy*."""
+        if strategy == EmbeddingStrategy.NAME_ONLY:
+            return func.name
+
+        if strategy == EmbeddingStrategy.NAME_DOMAIN:
+            return f"{func.name} {func.domain or ''}"
+
+        if strategy == EmbeddingStrategy.SEMANTIC:
+            parts = [func.name, func.domain or ""]
+            if func.trigger:
+                parts.append(
+                    f"触发: {'; '.join(fv.desc for fv in func.trigger[:2])}"
+                )
+            if func.action:
+                parts.append(
+                    f"动作: {'; '.join(fv.desc for fv in func.action[:2])}"
+                )
+            return " ".join(parts)
+
+        # FULL
+        parts = [func.name, func.domain or ""]
+        for fv in func.trigger:
+            parts.append(fv.desc)
+        for fv in func.action:
+            parts.append(fv.desc)
+        for fv in func.benefit:
+            parts.append(fv.desc)
+        return " ".join(parts)
+
+    def refresh(self, func_id: str) -> None:
+        """Re-embed a single Function and upsert into the vector store."""
+        if self.storage is None or self.vector_store is None:
+            logger.warning("Cannot refresh: storage or vector_store not configured")
+            return
+        func = self.storage.get(func_id)
+        if func is None:
+            logger.warning("Function %s not found for refresh", func_id)
+            return
+        vector = self.embed_function(func)
+        self.vector_store.upsert(func_id, vector)
+
+    def refresh_all(self, batch_size: int = 100) -> RefreshResult:
+        """Re-embed all Functions in batches and upsert into the vector store.
+
+        Used after model or strategy changes.
+        """
+        if self.storage is None or self.vector_store is None:
+            logger.warning("Cannot refresh_all: storage or vector_store not configured")
+            return RefreshResult(total=0, refreshed=0)
+
+        refreshed = 0
+        offset = 0
+
+        while True:
+            batch = self.storage.list_functions(
+                offset=offset, limit=batch_size
+            )
+            if not batch:
+                break
+
+            texts = [self.function_to_text(f) for f in batch]
+            vectors = self.embed_batch(texts)
+            self.vector_store.upsert_batch(
+                {f.id: v for f, v in zip(batch, vectors)}
+            )
+            refreshed += len(batch)
+            offset += batch_size
+
+        return RefreshResult(total=refreshed, refreshed=refreshed)
+
+    # ── Private ─────────────────────────────────────────────────────
+
+    def _create_embedder(self, model: str, dimension: int):
+        """Create the appropriate embedder backend."""
+        model_name = self._MODEL_MAP.get(model, model)
+
+        try:
+            return _SentenceTransformerEmbedder(model_name, dimension)
+        except ImportError:
+            logger.info(
+                "sentence-transformers not available, "
+                "falling back to TF-IDF embedder"
+            )
+            return _SimpleTFIDFEmbedder(dimension=dimension)
+        except Exception as exc:
+            logger.warning(
+                "Failed to load sentence-transformers model %s: %s. "
+                "Falling back to TF-IDF embedder",
+                model_name,
+                exc,
+            )
+            return _SimpleTFIDFEmbedder(dimension=dimension)