flurryx-code-memory 0.4.0__py3-none-any.whl

code_memory/embed/m3.py

@@ -0,0 +1,174 @@
+"""BGE-M3 hybrid embedder: dense + sparse from one forward pass.
+
+Opt-in backend (``EMBED_BACKEND=flagembed``). Loads m3 in-process via
+FlagEmbedding, which means each Python process pays a ~5-15 s
+cold-load. Worth it for long-lived processes (watcher, MCP server)
+that want the sparse signal; not worth it for hook-driven per-save
+CLI invocations — :class:`code_memory.embed.OllamaEmbedder` is the
+default for that reason.
+
+m3 emits three views per input:
+
+* Dense (1024-d float) — semantic similarity (cosine).
+* Sparse (token-id -> weight) — lexical/identifier signal akin to BM25
+  but learned. Used for code search where exact symbol names matter.
+* ColBERT multi-vec — not used here; cross-encoder rerank covers the
+  late-interaction case.
+
+Fusion happens server-side in Qdrant (RRF / DBSF), so both views are
+stored alongside each chunk and combined at query time.
+"""
+
+from __future__ import annotations
+
+import logging
+import platform
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any
+
+from ..config import CONFIG
+
+log = logging.getLogger(__name__)
+
+# FlagEmbedding requires a HF repo id. The legacy ``EMBED_MODEL`` env
+# var used the Ollama short name (``bge-m3``), which HF rejects, so we
+# remap it here. Other models pass through unchanged.
+_OLLAMA_TO_HF = {"bge-m3": "BAAI/bge-m3"}
+DEFAULT_MODEL = "BAAI/bge-m3"
+
+
+def _resolve_model(name: str | None) -> str:
+    raw = (name or CONFIG.embed_model or DEFAULT_MODEL).strip()
+    return _OLLAMA_TO_HF.get(raw, raw)
+
+
+@dataclass(frozen=True)
+class SparseVec:
+    """Sparse vector in Qdrant's (indices, values) layout."""
+
+    indices: list[int]
+    values: list[float]
+
+
+@dataclass(frozen=True)
+class HybridVec:
+    dense: list[float]
+    sparse: SparseVec
+
+
+def _detect_device() -> str:
+    """Best available accelerator; falls back to CPU."""
+    try:
+        import torch
+    except ImportError:
+        return "cpu"
+    if platform.system() == "Darwin" and torch.backends.mps.is_available():
+        return "mps"
+    if torch.cuda.is_available():
+        return "cuda"
+    return "cpu"
+
+
+class M3Embedder:
+    """Stateful BGE-M3 wrapper producing dense + sparse vectors.
+
+    Heavy to construct (downloads + loads ~2.3GB on first use). Cache
+    the instance for the process lifetime — see ``code_memory.embed``
+    factory below.
+    """
+
+    def __init__(
+        self,
+        model: str | None = None,
+        device: str | None = None,
+        use_fp16: bool | None = None,
+        batch_size: int = 12,
+    ) -> None:
+        from FlagEmbedding import BGEM3FlagModel
+
+        self.model_name = _resolve_model(model)
+        self.device = device or _detect_device()
+        # fp16 only safe on CUDA/MPS; CPU stays at fp32 for numerical
+        # stability + because some BLAS kernels don't support fp16.
+        if use_fp16 is None:
+            use_fp16 = self.device in ("cuda", "mps")
+        self.batch_size = batch_size
+        log.info(
+            "m3: loading %s (device=%s fp16=%s)",
+            self.model_name,
+            self.device,
+            use_fp16,
+        )
+        self._impl = BGEM3FlagModel(
+            self.model_name,
+            use_fp16=use_fp16,
+            devices=self.device,
+        )
+
+    # ----------------------------------------------------------- batch
+
+    def embed(self, texts: Sequence[str]) -> list[HybridVec]:
+        if not texts:
+            return []
+        out = self._impl.encode(
+            list(texts),
+            batch_size=self.batch_size,
+            return_dense=True,
+            return_sparse=True,
+            return_colbert_vecs=False,
+        )
+        dense = out["dense_vecs"]
+        sparse = out["lexical_weights"]
+        return [
+            HybridVec(
+                dense=list(map(float, dense[i])),
+                sparse=_to_qdrant_sparse(sparse[i]),
+            )
+            for i in range(len(texts))
+        ]
+
+    def embed_one(self, text: str) -> HybridVec:
+        return self.embed([text])[0]
+
+    # ------------------------------------------------------------ misc
+
+    def close(self) -> None:
+        # FlagEmbedding has no explicit close; drop the reference to free
+        # GPU mem on next gc cycle.
+        self._impl = None
+
+    def __enter__(self) -> M3Embedder:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+
+def _to_qdrant_sparse(weights: dict[Any, Any]) -> SparseVec:
+    """Convert m3 ``{token_id: weight}`` mapping to Qdrant sparse format.
+
+    m3 returns numpy floats keyed by string token IDs. Qdrant wants
+    plain ints and floats; the conversion is explicit so misbehaving
+    inputs (negative weights, NaN) are dropped rather than poisoning the
+    index.
+    """
+    indices: list[int] = []
+    values: list[float] = []
+    for tok, w in weights.items():
+        try:
+            idx = int(tok)
+        except (TypeError, ValueError):
+            continue
+        val = float(w)
+        if val <= 0.0 or val != val:  # drop NaN / non-positive
+            continue
+        indices.append(idx)
+        values.append(val)
+    return SparseVec(indices=indices, values=values)
+
+
+# Factory + singleton live in ``code_memory.embed.__init__`` so the
+# Ollama and M3 backends share one selection mechanism. ``M3Embedder``
+# itself remains directly constructible for tests and for users who
+# want to bypass the env-var dispatch.

code_memory/embed/ollama.py

@@ -0,0 +1,92 @@
+"""Ollama-backed dense embedder (default backend).
+
+Runs `bge-m3` (or any Ollama-served model) over HTTP. Ollama keeps the
+model loaded in its own daemon, so short-lived CLI processes (e.g.
+``code-memory reingest <file>`` invoked from a save-file hook) reuse
+the warm model instead of paying a ~5-15 s cold load every call.
+
+Trade-off vs the in-process FlagEmbedding path: Ollama only exposes the
+dense head of m3 — no sparse, no ColBERT. Sparse is returned as an
+empty :class:`SparseVec` so the Qdrant hybrid layout still upserts
+cleanly; queries through the hybrid slot then degrade to dense-only at
+RRF time. Users who want true m3 hybrid (dense + sparse from one
+forward pass) can flip ``EMBED_BACKEND=flagembed`` and accept the
+cold-load cost.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import logging
+
+import httpx
+
+from ..config import CONFIG
+from ..resilience import with_retry
+from .m3 import HybridVec, SparseVec
+
+log_ = logging.getLogger(__name__)
+
+
+class OllamaEmbedder:
+    """Thin sync wrapper over Ollama /api/embed.
+
+    Returns :class:`HybridVec` with an empty sparse component so the
+    shape matches :class:`M3Embedder`. The empty sparse vector is a
+    deliberate signal to :class:`QdrantStore` that hybrid fusion will
+    degrade to dense-only for this point.
+    """
+
+    def __init__(
+        self,
+        url: str | None = None,
+        model: str | None = None,
+        timeout: float = 300.0,
+    ) -> None:
+        self.url = (url or CONFIG.ollama_url).rstrip("/")
+        self.model = model or CONFIG.embed_model
+        self._client = httpx.Client(timeout=timeout)
+
+    def embed(self, texts: Sequence[str]) -> list[HybridVec]:
+        if not texts:
+            return []
+
+        def _call():
+            res = self._client.post(
+                f"{self.url}/api/embed",
+                json={"model": self.model, "input": list(texts)},
+            )
+            res.raise_for_status()
+            data = res.json()
+            embeddings = data.get("embeddings")
+            if embeddings is None:
+                raise RuntimeError(f"Ollama returned no embeddings: {data}")
+            return embeddings
+
+        embeddings = with_retry(
+            _call,
+            max_retries=3,
+            backoff_s=1.0,
+            on_retry=lambda attempt, exc: log_.warning(
+                "ollama embed retry %d/3 after %s", attempt, exc
+            ),
+        )
+
+        empty = SparseVec(indices=[], values=[])
+        return [
+            HybridVec(dense=[float(x) for x in vec], sparse=empty)
+            for vec in embeddings
+        ]
+
+    def embed_one(self, text: str) -> HybridVec:
+        return self.embed([text])[0]
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> OllamaEmbedder:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()

code_memory/embed/tei.py

@@ -0,0 +1,106 @@
+"""text-embeddings-inference (TEI) backend.
+
+`HuggingFace TEI <https://github.com/huggingface/text-embeddings-inference>`_
+is a purpose-built embedding server. On a Linux + NVIDIA host with the
+same ``BAAI/bge-m3`` weights, it serves embeddings at **5-10× the
+throughput** of Ollama because:
+
+* Built on ONNX Runtime / candle-rs with native CUDA batching.
+* Streams + dynamically batches concurrent requests instead of
+  serialising one-at-a-time.
+* No HTTP-to-llama.cpp daemon hop per call.
+
+For enterprise CI / staging where the cold ingest of a large monorepo
+matters, this is the way to break the ``bge-m3`` throughput floor
+without changing models or losing semantic recall.
+
+Trade-off vs ``OllamaEmbedder``:
+
+* Same shape (dense-only ``HybridVec`` with empty sparse) so callers
+  swap backends transparently.
+* TEI must be running before code-memory ingests; Ollama-style "I
+  brought my own daemon" still applies.
+* On Mac (no NVIDIA), TEI's CPU path is roughly on par with Ollama's
+  Metal path — there's no advantage. Stay on Ollama there.
+
+Activated via ``EMBED_BACKEND=tei``; URL via ``TEI_URL``.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import httpx
+
+from ..config import CONFIG
+from .m3 import HybridVec, SparseVec
+
+
+class TEIEmbedder:
+    """Sync wrapper over TEI's ``/embed`` endpoint.
+
+    Returns :class:`HybridVec` with an empty sparse component so the
+    shape matches :class:`OllamaEmbedder` and :class:`M3Embedder`.
+    Callers (pipeline, retrieve) need no branching on backend type.
+
+    TEI's request payload differs slightly from Ollama's:
+
+    * Endpoint: ``POST /embed``
+    * Body: ``{"inputs": [...]}``
+    * Response: ``[[float, ...], [float, ...]]`` (raw vector list, no
+      wrapping object).
+
+    A ``truncate=true`` flag is set so over-length chunks are silently
+    truncated to the model's max sequence length rather than failing
+    the whole batch — the same forgiving semantic Ollama applies.
+    """
+
+    def __init__(
+        self,
+        url: str | None = None,
+        timeout: float = 300.0,
+    ) -> None:
+        # TEI doesn't accept a model id at request time — the daemon
+        # is launched with a single ``--model-id`` flag — so we don't
+        # carry one through requests. ``self.model`` exists for
+        # parity with :class:`OllamaEmbedder` and is sourced from
+        # ``EMBED_MODEL`` so the cache key namespace lines up across
+        # backends pointing at the same model weights.
+        self.url = (url or CONFIG.tei_url).rstrip("/")
+        self.model = CONFIG.embed_model
+        self._client = httpx.Client(timeout=timeout)
+
+    def embed(self, texts: Sequence[str]) -> list[HybridVec]:
+        if not texts:
+            return []
+        res = self._client.post(
+            f"{self.url}/embed",
+            json={"inputs": list(texts), "truncate": True},
+        )
+        res.raise_for_status()
+        data = res.json()
+        # TEI returns ``[[float, ...], ...]`` — a bare list of
+        # vectors, one per input, in the same order. No wrapper key.
+        if not isinstance(data, list):
+            raise RuntimeError(f"TEI returned unexpected shape: {type(data).__name__}")
+        if len(data) != len(texts):
+            raise RuntimeError(
+                f"TEI returned {len(data)} vectors for {len(texts)} inputs"
+            )
+        empty = SparseVec(indices=[], values=[])
+        return [
+            HybridVec(dense=[float(x) for x in vec], sparse=empty)
+            for vec in data
+        ]
+
+    def embed_one(self, text: str) -> HybridVec:
+        return self.embed([text])[0]
+
+    def close(self) -> None:
+        self._client.close()
+
+    def __enter__(self) -> TEIEmbedder:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()

code_memory/episodic/__init__.py

@@ -0,0 +1,3 @@
+from .sqlite_store import Episode, EpisodicStore
+
+__all__ = ["Episode", "EpisodicStore"]

code_memory/episodic/sqlite_store.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import sqlite3
+import time
+import uuid
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+from ..config import CONFIG
+
+# Base table — kept minimal so a legacy DB opens without errors. Every
+# additional column lives in ``_MIGRATIONS`` so loading an old database
+# transparently catches it up to the latest schema.
+_BASE_SCHEMA = """
+CREATE TABLE IF NOT EXISTS episodes (
+    id TEXT PRIMARY KEY,
+    ts REAL NOT NULL,
+    prompt TEXT NOT NULL,
+    plan TEXT,
+    patch TEXT,
+    verdict TEXT,
+    tags TEXT,
+    meta TEXT
+);
+"""
+
+# Idempotent migrations. Each statement is run independently; failures
+# (e.g. "duplicate column" when the migration has already been applied)
+# are swallowed because that's the success path for an idempotent
+# migration. Indexes that reference migration-added columns must come
+# AFTER the corresponding ADD COLUMN, hence interleaved here.
+_MIGRATIONS = (
+    "ALTER TABLE episodes ADD COLUMN head_sha TEXT",
+    "CREATE INDEX IF NOT EXISTS idx_episodes_ts ON episodes(ts)",
+    "CREATE INDEX IF NOT EXISTS idx_episodes_verdict ON episodes(verdict)",
+    "CREATE INDEX IF NOT EXISTS idx_episodes_head_sha ON episodes(head_sha)",
+    # Content hash for dedup. Same user prompt re-asserted across turns
+    # produced one row per assertion before; now the existing row gets
+    # its ``ts`` refreshed and the new insert is a no-op. Non-unique by
+    # design so legacy rows (NULL hash) still load without conflict.
+    "ALTER TABLE episodes ADD COLUMN content_hash TEXT",
+    "CREATE INDEX IF NOT EXISTS idx_episodes_content_hash ON episodes(content_hash)",
+)
+
+
+def _content_hash(prompt: str) -> str:
+    """SHA-256 over the user prompt, normalized.
+
+    Dedup key is prompt-only on purpose: the same prompt typed twice
+    represents the same intent, regardless of which plan/patch/verdict
+    the agent eventually produced. Whitespace is normalized so a
+    trailing newline doesn't split otherwise-identical rows.
+    """
+    return hashlib.sha256(prompt.strip().encode("utf-8")).hexdigest()
+
+
+@dataclass
+class Episode:
+    prompt: str
+    plan: str | None = None
+    patch: str | None = None
+    verdict: str | None = None  # pass | fail | partial
+    tags: list[str] = field(default_factory=list)
+    meta: dict[str, Any] = field(default_factory=dict)
+    id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    ts: float = field(default_factory=time.time)
+    # Git HEAD at the moment the episode was recorded — links the
+    # agent's work back to the code state the graph was indexing then.
+    head_sha: str | None = None
+
+
+class EpisodicStore:
+    def __init__(self, path: Path | None = None) -> None:
+        self.path = path or CONFIG.episodic_db
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        self.conn = sqlite3.connect(self.path)
+        self.conn.executescript(_BASE_SCHEMA)
+        for stmt in _MIGRATIONS:
+            try:
+                self.conn.execute(stmt)
+            except sqlite3.OperationalError:
+                # column already added by a prior process — that's the
+                # success path for an idempotent migration
+                pass
+        self.conn.commit()
+
+    def add(self, ep: Episode) -> str:
+        """Insert an episode, deduping on prompt content.
+
+        If an existing row has the same ``content_hash``, refresh its
+        ``ts`` to ``ep.ts`` and fill any previously-NULL fields from
+        the new episode (plan/patch/verdict/head_sha). Tags are unioned
+        and meta is merged with new values winning on key collision.
+        Returns the existing row's id so vector upserts stay idempotent.
+        """
+        hash_ = _content_hash(ep.prompt)
+        existing = self.conn.execute(
+            "SELECT id, plan, patch, verdict, head_sha, tags, meta "
+            "FROM episodes WHERE content_hash = ? LIMIT 1",
+            (hash_,),
+        ).fetchone()
+        if existing is not None:
+            return self._refresh_existing(existing, ep)
+
+        self.conn.execute(
+            "INSERT INTO episodes(id, ts, prompt, plan, patch, verdict, tags, meta, head_sha, content_hash) "
+            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
+            (
+                ep.id,
+                ep.ts,
+                ep.prompt,
+                ep.plan,
+                ep.patch,
+                ep.verdict,
+                json.dumps(ep.tags),
+                json.dumps(ep.meta),
+                ep.head_sha,
+                hash_,
+            ),
+        )
+        self.conn.commit()
+        return ep.id
+
+    def _refresh_existing(
+        self, existing: tuple[Any, ...], ep: Episode
+    ) -> str:
+        existing_id = str(existing[0])
+        old_plan, old_patch, old_verdict, old_head = (
+            existing[1],
+            existing[2],
+            existing[3],
+            existing[4],
+        )
+        old_tags = json.loads(existing[5]) if existing[5] else []
+        old_meta = json.loads(existing[6]) if existing[6] else {}
+
+        merged_tags = list(dict.fromkeys([*old_tags, *ep.tags]))
+        merged_meta = {**old_meta, **ep.meta}
+
+        self.conn.execute(
+            "UPDATE episodes SET "
+            "  ts = ?, "
+            "  plan = COALESCE(plan, ?), "
+            "  patch = COALESCE(patch, ?), "
+            "  verdict = COALESCE(verdict, ?), "
+            "  head_sha = COALESCE(head_sha, ?), "
+            "  tags = ?, "
+            "  meta = ? "
+            "WHERE id = ?",
+            (
+                ep.ts,
+                ep.plan if ep.plan else None,
+                ep.patch if ep.patch else None,
+                ep.verdict if ep.verdict else None,
+                ep.head_sha,
+                json.dumps(merged_tags),
+                json.dumps(merged_meta),
+                existing_id,
+            ),
+        )
+        self.conn.commit()
+        return existing_id
+
+    def dedupe(self) -> dict[str, list[str]]:
+        """Compact pre-existing duplicates in the table.
+
+        For each ``content_hash`` group with >1 row, keep the row with
+        the oldest ``ts`` (first observation), update its ``ts`` to
+        ``MAX(ts)`` of the group so retrieval still surfaces it as
+        recent, and delete the rest. Returns ``{kept_id: [removed_ids]}``
+        so callers (e.g. the orchestrator) can prune matching vectors.
+
+        Backfills ``content_hash`` for legacy NULL rows on the fly.
+        """
+        null_rows = self.conn.execute(
+            "SELECT id, prompt FROM episodes WHERE content_hash IS NULL"
+        ).fetchall()
+        for ep_id, prompt in null_rows:
+            self.conn.execute(
+                "UPDATE episodes SET content_hash = ? WHERE id = ?",
+                (_content_hash(prompt), ep_id),
+            )
+        if null_rows:
+            self.conn.commit()
+
+        groups = self.conn.execute(
+            "SELECT content_hash FROM episodes "
+            "WHERE content_hash IS NOT NULL "
+            "GROUP BY content_hash HAVING COUNT(*) > 1"
+        ).fetchall()
+
+        removed: dict[str, list[str]] = {}
+        for (hash_,) in groups:
+            rows = self.conn.execute(
+                "SELECT id, ts FROM episodes WHERE content_hash = ? "
+                "ORDER BY ts ASC",
+                (hash_,),
+            ).fetchall()
+            keep_id = str(rows[0][0])
+            max_ts = max(float(r[1]) for r in rows)
+            del_ids = [str(r[0]) for r in rows[1:]]
+            self.conn.execute(
+                "UPDATE episodes SET ts = ? WHERE id = ?", (max_ts, keep_id)
+            )
+            self.conn.executemany(
+                "DELETE FROM episodes WHERE id = ?", [(d,) for d in del_ids]
+            )
+            removed[keep_id] = del_ids
+        self.conn.commit()
+        return removed
+
+    def get(self, ep_id: str) -> Episode | None:
+        row = self.conn.execute(
+            "SELECT id, ts, prompt, plan, patch, verdict, tags, meta, head_sha "
+            "FROM episodes WHERE id = ?",
+            (ep_id,),
+        ).fetchone()
+        if row is None:
+            return None
+        return _row_to_episode(row)
+
+    def recent(self, limit: int = 20) -> list[Episode]:
+        rows = self.conn.execute(
+            "SELECT id, ts, prompt, plan, patch, verdict, tags, meta, head_sha "
+            "FROM episodes ORDER BY ts DESC LIMIT ?",
+            (limit,),
+        ).fetchall()
+        return [_row_to_episode(r) for r in rows]
+
+    def by_ids(self, ids: list[str]) -> list[Episode]:
+        if not ids:
+            return []
+        placeholders = ",".join("?" for _ in ids)
+        rows = self.conn.execute(
+            f"SELECT id, ts, prompt, plan, patch, verdict, tags, meta, head_sha "
+            f"FROM episodes WHERE id IN ({placeholders})",
+            ids,
+        ).fetchall()
+        return [_row_to_episode(r) for r in rows]
+
+    def close(self) -> None:
+        self.conn.close()
+
+
+def _row_to_episode(row: tuple[Any, ...]) -> Episode:
+    return Episode(
+        id=row[0],
+        ts=row[1],
+        prompt=row[2],
+        plan=row[3],
+        patch=row[4],
+        verdict=row[5],
+        tags=json.loads(row[6]) if row[6] else [],
+        meta=json.loads(row[7]) if row[7] else {},
+        head_sha=row[8] if len(row) > 8 else None,
+    )
+
+
+def episode_text(ep: Episode) -> str:
+    """Composite text for embedding."""
+    parts = [f"PROMPT:\n{ep.prompt}"]
+    if ep.plan:
+        parts.append(f"PLAN:\n{ep.plan}")
+    if ep.patch:
+        parts.append(f"PATCH:\n{ep.patch}")
+    if ep.verdict:
+        parts.append(f"VERDICT: {ep.verdict}")
+    return "\n\n".join(parts)
+
+
+def episode_payload(ep: Episode) -> dict[str, Any]:
+    d = asdict(ep)
+    d.pop("plan", None)
+    d.pop("patch", None)
+    return d

code_memory/extractor/__init__.py

@@ -0,0 +1,3 @@
+from .treesitter import ExtractedFile, Extractor, Symbol
+
+__all__ = ["Extractor", "ExtractedFile", "Symbol"]