ragx-cli 0.1.0__py3-none-any.whl

ragx/core/query.py

@@ -0,0 +1,181 @@
+"""Query pipeline: expansion -> fan-out retrieval -> RRF -> heat traversal -> rerank -> combine."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Sequence
+
+from ragx.core.config import Config, db_path, vectors_path
+from ragx.core.errors import RagxError
+from ragx.core.expansion import expand_query
+from ragx.core.fusion import rrf
+from ragx.core.models import QueryOutput, ScoredChunk
+from ragx.core.scoring import combine
+from ragx.core.store import Store
+from ragx.core.traversal import propagate_heat
+from ragx.core.vectors import VectorIndex
+from ragx.providers.base import Embedder, Generator, Reranker
+
+RERANK_CAP = 100  # cross-encoder cost guard: rerank at most this many candidates
+
+
+@dataclass
+class QueryOptions:
+    top: int = 8
+    files_only: bool = False
+    expand: bool = True
+    graph: bool = True
+    rerank: bool = True
+    hops: int | None = None
+    explain: bool = False
+
+
+def _normalize(vec: Sequence[float]) -> list[float]:
+    norm = math.sqrt(sum(v * v for v in vec)) or 1.0
+    return [v / norm for v in vec]
+
+
+def _check_model(store: Store, embedder: Embedder) -> None:
+    built_with = store.get_meta("embedding_model")
+    if built_with != embedder.model:
+        raise RagxError(
+            f"index built with {built_with!r}, config says {embedder.model!r}; "
+            "run `ragx index` to rebuild"
+        )
+
+
+def run_query(
+    root: Path,
+    cfg: Config,
+    embedder: Embedder,
+    text: str,
+    opts: QueryOptions,
+    *,
+    generator: Generator | None = None,
+    reranker: Reranker | None = None,
+) -> QueryOutput:
+    if not text.strip():
+        raise RagxError("empty query")
+    with Store(db_path(root)) as store:
+        if store.chunk_count() == 0:
+            return QueryOutput(query=text)
+        _check_model(store, embedder)
+        index = VectorIndex.load(vectors_path(root), int(store.get_meta("embedding_dim") or 0))
+
+        # 1. expansion (optional, one LLM call)
+        variants: list[str] = []
+        if opts.expand and generator is not None:
+            exp = expand_query(
+                generator, text,
+                variants=cfg.get("expansion.variants"), hyde=cfg.get("expansion.hyde"),
+            )
+            variants = exp.variants + ([exp.hyde] if exp.hyde else [])
+
+        # 2. fan-out retrieval + RRF fusion
+        q_vecs = embedder.embed_queries([text, *variants])
+        per_top = cfg.get("fusion.per_query_top")
+        rankings = [[cid for cid, _ in index.search(v, per_top)] for v in q_vecs]
+        fused = rrf(rankings, k=cfg.get("fusion.rrf_k"))
+        if not fused:
+            return QueryOutput(query=text, variants=variants)
+
+        # 3. graph expansion via heat propagation
+        heat: dict[int, float] = dict(fused)
+        trace: dict[int, dict] = {}
+        if opts.graph and store.edge_count() > 0:
+            qv = _normalize(q_vecs[0])
+
+            def query_sim_fn(ids: Sequence[int]) -> dict[int, float]:
+                vecs = index.get_vectors(list(ids))  # stored normalized (cosine space)
+                return {i: max(0.0, sum(a * b for a, b in zip(v, qv))) for i, v in zip(ids, vecs)}
+
+            result = propagate_heat(
+                dict(fused), store.neighbors, query_sim_fn,
+                hops=opts.hops if opts.hops is not None else cfg.get("traversal.hops"),
+                decay=cfg.get("traversal.decay"),
+                query_floor=cfg.get("traversal.query_floor"),
+                max_frontier=cfg.get("traversal.max_frontier"),
+            )
+            heat, trace = result.heat, result.trace
+
+        candidates = sorted(set(fused) | set(heat))
+        alpha, beta, gamma = (
+            cfg.get("scoring.alpha_rerank"), cfg.get("scoring.beta_heat"), cfg.get("scoring.gamma_vector"),
+        )
+
+        # 4. cross-encoder rerank (capped shortlist by pre-score)
+        rerank_scores: dict[int, float] | None = None
+        if opts.rerank and reranker is not None:
+            pre = combine(candidates, fused, heat, None, alpha=alpha, beta=beta, gamma=gamma)
+            shortlist = sorted(candidates, key=lambda c: (-pre.get(c, 0.0), c))[:RERANK_CAP]
+            by_id = {c.id: c for c in store.get_chunks(shortlist)}
+            ordered = [cid for cid in shortlist if cid in by_id]
+            scores = reranker.score(text, [by_id[c].text for c in ordered])
+            rerank_scores = {cid: float(s) for cid, s in zip(ordered, scores)}
+            candidates = ordered  # unreranked stragglers would score 0 anyway
+
+        # 5. combined scoring + output
+        final = combine(candidates, fused, heat, rerank_scores, alpha=alpha, beta=beta, gamma=gamma)
+        top_ids = sorted(candidates, key=lambda c: (-final.get(c, 0.0), c))[: opts.top]
+        by_id = {c.id: c for c in store.get_chunks(top_ids)}
+        results = [
+            ScoredChunk(
+                chunk=by_id[cid],
+                score=final.get(cid, 0.0),
+                vector_score=fused.get(cid, 0.0),
+                heat=heat.get(cid, 0.0),
+                rerank_score=(rerank_scores or {}).get(cid, 0.0),
+                explain=trace.get(cid) if opts.explain else None,
+            )
+            for cid in top_ids
+            if cid in by_id
+        ]
+        return QueryOutput(query=text, variants=variants, results=results)
+
+
+def to_query_json(out: QueryOutput, *, max_chunk_chars: int = 1200) -> dict:
+    """Serialize as the versioned agent-facing schema."""
+    return {
+        "schema": "ragx.query.v1",
+        "query": out.query,
+        "variants": out.variants,
+        "results": [
+            {
+                "chunk_id": r.chunk.id,
+                "file": r.chunk.file_path,
+                "line_start": r.chunk.line_start,
+                "line_end": r.chunk.line_end,
+                "byte_start": r.chunk.byte_start,
+                "byte_end": r.chunk.byte_end,
+                "score": round(r.score, 6),
+                "breakdown": {
+                    "vector": round(r.vector_score, 6),
+                    "heat": round(r.heat, 6),
+                    "rerank": round(r.rerank_score, 6),
+                },
+                "text": r.chunk.text[:max_chunk_chars],
+                "truncated": len(r.chunk.text) > max_chunk_chars,
+                **({"explain": r.explain} if r.explain else {}),
+            }
+            for r in out.results
+        ],
+    }
+
+
+def to_files_json(out: QueryOutput) -> dict:
+    """Aggregate chunk scores per file: sum of each file's top-3 chunk scores."""
+    per_file: dict[str, list[ScoredChunk]] = {}
+    for r in out.results:
+        per_file.setdefault(r.chunk.file_path, []).append(r)
+    files = [
+        {
+            "file": path,
+            "score": round(sum(sorted((r.score for r in rs), reverse=True)[:3]), 6),
+            "chunk_ids": [r.chunk.id for r in sorted(rs, key=lambda r: r.score, reverse=True)],
+        }
+        for path, rs in per_file.items()
+    ]
+    files.sort(key=lambda f: f["score"], reverse=True)
+    return {"schema": "ragx.files.v1", "query": out.query, "variants": out.variants, "files": files}

ragx/core/scoring.py

@@ -0,0 +1,48 @@
+"""Score normalization and multi-signal combination."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+
+def normalize(scores: dict[int, float]) -> dict[int, float]:
+    """Min-max scale to [0, 1]; empty -> {}; all-equal -> all 1.0."""
+    if not scores:
+        return {}
+    lo = min(scores.values())
+    hi = max(scores.values())
+    if hi == lo:
+        return {k: 1.0 for k in scores}
+    span = hi - lo
+    return {k: (v - lo) / span for k, v in scores.items()}
+
+
+def combine(
+    candidates: Iterable[int],
+    vector: dict[int, float],
+    heat: dict[int, float],
+    rerank: dict[int, float] | None,
+    *,
+    alpha: float,
+    beta: float,
+    gamma: float,
+) -> dict[int, float]:
+    """Combine per-signal scores, min-max normalizing each component over candidates first."""
+    ids = list(candidates)
+    raw_vector = {cid: vector.get(cid, 0.0) for cid in ids}
+    raw_heat = {cid: heat.get(cid, 0.0) for cid in ids}
+    norm_vector = normalize(raw_vector)
+    norm_heat = normalize(raw_heat)
+
+    if rerank is not None:
+        raw_rerank = {cid: rerank.get(cid, 0.0) for cid in ids}
+        norm_rerank = normalize(raw_rerank)
+        return {
+            cid: alpha * norm_rerank[cid] + beta * norm_heat[cid] + gamma * norm_vector[cid]
+            for cid in ids
+        }
+
+    denom = beta + gamma
+    w_heat = beta / denom if denom else 0.0
+    w_vector = gamma / denom if denom else 0.0
+    return {cid: w_heat * norm_heat[cid] + w_vector * norm_vector[cid] for cid in ids}

ragx/core/store.py

@@ -0,0 +1,161 @@
+"""Single-file SQLite store for files, chunks, edges, and a key-value manifest."""
+
+from __future__ import annotations
+
+import sqlite3
+from pathlib import Path
+from typing import Sequence
+
+from ragx.core.models import Chunk, ChunkDraft, FileRecord
+
+SCHEMA = """
+CREATE TABLE meta(key TEXT PRIMARY KEY, value TEXT NOT NULL);
+CREATE TABLE files(path TEXT PRIMARY KEY, content_hash TEXT NOT NULL,
+                   mtime REAL NOT NULL, chunk_count INTEGER NOT NULL);
+CREATE TABLE chunks(id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    file_path TEXT NOT NULL REFERENCES files(path) ON DELETE CASCADE,
+                    text TEXT NOT NULL, byte_start INTEGER NOT NULL, byte_end INTEGER NOT NULL,
+                    line_start INTEGER NOT NULL, line_end INTEGER NOT NULL);
+CREATE INDEX chunks_file ON chunks(file_path);
+CREATE TABLE edges(src INTEGER NOT NULL, dst INTEGER NOT NULL, weight REAL NOT NULL,
+                   PRIMARY KEY(src, dst)) WITHOUT ROWID;
+CREATE INDEX edges_dst ON edges(dst);
+"""
+
+
+class Store:
+    def __init__(self, path: Path):
+        self.conn = sqlite3.connect(str(path))
+        self.conn.execute("PRAGMA journal_mode = WAL")
+        self.conn.execute("PRAGMA foreign_keys = ON")
+        user_version = self.conn.execute("PRAGMA user_version").fetchone()[0]
+        if user_version == 0:
+            self.conn.executescript(SCHEMA)
+            self.conn.execute("PRAGMA user_version = 1")
+            self.conn.commit()
+
+    def close(self) -> None:
+        self.conn.close()
+
+    def __enter__(self) -> Store:
+        return self
+
+    def __exit__(self, *exc_info) -> None:
+        self.close()
+
+    # -- manifest -----------------------------------------------------
+
+    def get_meta(self, key: str) -> str | None:
+        row = self.conn.execute("SELECT value FROM meta WHERE key = ?", (key,)).fetchone()
+        return row[0] if row else None
+
+    def set_meta(self, key: str, value: str) -> None:
+        self.conn.execute(
+            "INSERT INTO meta(key, value) VALUES (?, ?) "
+            "ON CONFLICT(key) DO UPDATE SET value = excluded.value",
+            (key, value),
+        )
+        self.conn.commit()
+
+    # -- files ----------------------------------------------------------
+
+    def get_file_hashes(self) -> dict[str, str]:
+        rows = self.conn.execute("SELECT path, content_hash FROM files").fetchall()
+        return dict(rows)
+
+    def upsert_file(self, rec: FileRecord) -> None:
+        self.conn.execute(
+            "INSERT INTO files(path, content_hash, mtime, chunk_count) VALUES (?, ?, ?, ?) "
+            "ON CONFLICT(path) DO UPDATE SET content_hash = excluded.content_hash, "
+            "mtime = excluded.mtime, chunk_count = excluded.chunk_count",
+            (rec.path, rec.content_hash, rec.mtime, rec.chunk_count),
+        )
+        self.conn.commit()
+
+    def delete_file(self, path: str) -> list[int]:
+        ids = [
+            row[0]
+            for row in self.conn.execute("SELECT id FROM chunks WHERE file_path = ?", (path,))
+        ]
+        if ids:
+            placeholders = ",".join("?" * len(ids))
+            self.conn.execute(f"DELETE FROM edges WHERE src IN ({placeholders})", ids)
+            self.conn.execute(f"DELETE FROM edges WHERE dst IN ({placeholders})", ids)
+        self.conn.execute("DELETE FROM files WHERE path = ?", (path,))
+        self.conn.commit()
+        return ids
+
+    def file_count(self) -> int:
+        return self.conn.execute("SELECT COUNT(*) FROM files").fetchone()[0]
+
+    # -- chunks ------------------------------------------------------
+
+    def insert_chunks(self, file_path: str, drafts: Sequence[ChunkDraft]) -> list[int]:
+        ids = []
+        for d in drafts:
+            cur = self.conn.execute(
+                "INSERT INTO chunks(file_path, text, byte_start, byte_end, line_start, line_end) "
+                "VALUES (?, ?, ?, ?, ?, ?)",
+                (file_path, d.text, d.byte_start, d.byte_end, d.line_start, d.line_end),
+            )
+            ids.append(cur.lastrowid)
+        self.conn.commit()
+        return ids
+
+    def get_chunks(self, ids: Sequence[int]) -> list[Chunk]:
+        if not ids:
+            return []
+        placeholders = ",".join("?" * len(ids))
+        rows = self.conn.execute(
+            f"SELECT id, file_path, text, byte_start, byte_end, line_start, line_end "
+            f"FROM chunks WHERE id IN ({placeholders})",
+            list(ids),
+        ).fetchall()
+        by_id = {row[0]: Chunk(*row) for row in rows}
+        return [by_id[i] for i in ids if i in by_id]
+
+    def chunk_ids_for_file(self, path: str) -> list[int]:
+        rows = self.conn.execute(
+            "SELECT id FROM chunks WHERE file_path = ? ORDER BY id", (path,)
+        ).fetchall()
+        return [row[0] for row in rows]
+
+    def all_chunk_ids(self) -> list[int]:
+        rows = self.conn.execute("SELECT id FROM chunks ORDER BY id").fetchall()
+        return [row[0] for row in rows]
+
+    def chunk_count(self) -> int:
+        return self.conn.execute("SELECT COUNT(*) FROM chunks").fetchone()[0]
+
+    # -- edges (undirected; stored normalized src < dst) -----------------
+
+    def replace_edges(self, src: int, neighbors: Sequence[tuple[int, float]]) -> None:
+        self.conn.execute("DELETE FROM edges WHERE src = ? OR dst = ?", (src, src))
+        for dst, weight in neighbors:
+            lo, hi = (src, dst) if src < dst else (dst, src)
+            self.conn.execute(
+                "INSERT INTO edges(src, dst, weight) VALUES (?, ?, ?) "
+                "ON CONFLICT(src, dst) DO UPDATE SET weight = excluded.weight",
+                (lo, hi, weight),
+            )
+        self.conn.commit()
+
+    def neighbors(self, chunk_id: int) -> list[tuple[int, float]]:
+        rows = self.conn.execute(
+            "SELECT dst AS other, weight FROM edges WHERE src = ? "
+            "UNION ALL "
+            "SELECT src AS other, weight FROM edges WHERE dst = ?",
+            (chunk_id, chunk_id),
+        ).fetchall()
+        return sorted(((other, weight) for other, weight in rows), key=lambda t: -t[1])
+
+    def delete_edges_incident(self, ids: Sequence[int]) -> None:
+        if not ids:
+            return
+        placeholders = ",".join("?" * len(ids))
+        self.conn.execute(f"DELETE FROM edges WHERE src IN ({placeholders})", list(ids))
+        self.conn.execute(f"DELETE FROM edges WHERE dst IN ({placeholders})", list(ids))
+        self.conn.commit()
+
+    def edge_count(self) -> int:
+        return self.conn.execute("SELECT COUNT(*) FROM edges").fetchone()[0]

ragx/core/traversal.py

@@ -0,0 +1,68 @@
+"""Heat-propagation traversal: spreads seed relevance heat across the similarity graph."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+
+@dataclass
+class TraversalResult:
+    heat: dict[int, float]  # node -> final heat
+    trace: dict[int, dict]  # node -> {"seed": int, "parent": int | None, "edge_weight": float, "hop": int}
+
+
+def propagate_heat(
+    seeds: dict[int, float],
+    neighbors_fn: Callable[[int], list[tuple[int, float]]],
+    query_sim_fn: Callable[[Sequence[int]], dict[int, float]],
+    *,
+    hops: int = 2,
+    decay: float = 0.5,
+    query_floor: float = 0.35,
+    max_frontier: int = 150,
+) -> TraversalResult:
+    heat: dict[int, float] = {}
+    trace: dict[int, dict] = {}
+    for seed_id, score in seeds.items():
+        heat[seed_id] = score
+        trace[seed_id] = {"seed": seed_id, "parent": None, "edge_weight": 0.0, "hop": 0}
+
+    sim_cache: dict[int, float] = {}
+    frontier = sorted(seeds)
+
+    for hop in range(1, hops + 1):
+        if not frontier:
+            break
+
+        edges: list[tuple[int, int, float]] = []
+        first_seen: list[int] = []
+        seen_this_hop: set[int] = set()
+        for u in frontier:
+            for v, w in neighbors_fn(u):
+                edges.append((u, v, w))
+                if v not in seeds and v not in sim_cache and v not in seen_this_hop:
+                    first_seen.append(v)
+                    seen_this_hop.add(v)
+
+        if first_seen:
+            sim_cache.update(query_sim_fn(first_seen))
+
+        improved: set[int] = set()
+        for u, v, w in edges:
+            if v not in seeds and sim_cache.get(v, 0.0) < query_floor:
+                continue
+            contribution = heat[u] * w * decay
+            if contribution > heat.get(v, float("-inf")):
+                heat[v] = contribution
+                trace[v] = {
+                    "seed": trace[u]["seed"],
+                    "parent": u,
+                    "edge_weight": w,
+                    "hop": hop,
+                }
+                improved.add(v)
+
+        frontier = sorted(improved, key=lambda n: (-heat[n], n))[:max_frontier]
+
+    return TraversalResult(heat=heat, trace=trace)

ragx/core/vectors.py

@@ -0,0 +1,105 @@
+"""HNSW-backed vector index (hnswlib), cosine space.
+
+hnswlib doesn't expose a way to enumerate/count soft-deleted labels after
+save/load, so we track the deleted-id set and capacity ourselves in a small
+JSON sidecar next to the index file.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Sequence
+
+import hnswlib
+
+from ragx.core.errors import RagxError
+
+_M = 16
+_EF_CONSTRUCTION = 200
+
+
+def _meta_path(path: Path) -> Path:
+    return path.with_name(path.name + ".meta.json")
+
+
+class VectorIndex:
+    """Wraps an hnswlib.Index with resize-on-full and soft-delete tracking."""
+
+    def __init__(self, index: hnswlib.Index, path: Path, capacity: int, deleted: set[int]):
+        self._index = index
+        self._path = path
+        self._capacity = capacity
+        self._deleted = deleted
+
+    @classmethod
+    def create(cls, path: Path, dim: int, capacity: int = 2048) -> VectorIndex:
+        index = hnswlib.Index(space="cosine", dim=dim)
+        index.init_index(max_elements=capacity, M=_M, ef_construction=_EF_CONSTRUCTION)
+        return cls(index, path, capacity, set())
+
+    @classmethod
+    def load(cls, path: Path, dim: int) -> VectorIndex:
+        if not path.exists():
+            raise RagxError(f"vector index file not found: {path}")
+        meta_path = _meta_path(path)
+        capacity = 2048
+        deleted: set[int] = set()
+        if meta_path.exists():
+            meta = json.loads(meta_path.read_text())
+            capacity = meta["capacity"]
+            deleted = set(meta["deleted"])
+        index = hnswlib.Index(space="cosine", dim=dim)
+        index.load_index(str(path), max_elements=capacity)
+        return cls(index, path, capacity, deleted)
+
+    def add(self, ids: Sequence[int], vectors: Sequence[Sequence[float]]) -> None:
+        if not ids:
+            return
+        needed = self._index.get_current_count() + len(ids)
+        while needed > self._capacity:
+            self._capacity *= 2
+        if self._capacity > self._index.get_max_elements():
+            self._index.resize_index(self._capacity)
+        self._index.add_items(list(vectors), list(ids))
+
+    def search(self, vector: Sequence[float], k: int) -> list[tuple[int, float]]:
+        active = self.count
+        if active == 0 or k <= 0:
+            return []
+        k = min(k, active)
+        ef = max(100, 2 * k)
+        self._index.set_ef(ef)
+        labels, distances = self._index.knn_query([vector], k=k)
+        results = [
+            (int(label), max(0.0, min(1.0, 1.0 - float(dist))))
+            for label, dist in zip(labels[0], distances[0])
+        ]
+        results.sort(key=lambda x: x[1], reverse=True)
+        return results
+
+    def get_vectors(self, ids: Sequence[int]) -> list[list[float]]:
+        """Stored (normalized) vectors for `ids`, order preserved. Unknown id -> RagxError."""
+        if not ids:
+            return []
+        try:
+            return [list(map(float, v)) for v in self._index.get_items(list(ids))]
+        except RuntimeError as exc:
+            raise RagxError(f"unknown vector id in {list(ids)!r}: {exc}") from exc
+
+    def mark_deleted(self, ids: Sequence[int]) -> None:
+        known_ids = set(self._index.get_ids_list())
+        for i in ids:
+            if i in known_ids and i not in self._deleted:
+                self._index.mark_deleted(i)
+                self._deleted.add(i)
+
+    def save(self) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        self._index.save_index(str(self._path))
+        meta = {"capacity": self._capacity, "deleted": sorted(self._deleted)}
+        _meta_path(self._path).write_text(json.dumps(meta))
+
+    @property
+    def count(self) -> int:
+        return self._index.get_current_count() - len(self._deleted)

ragx/providers/base.py

@@ -0,0 +1,46 @@
+"""Provider protocols. Implementations live beside this module; factories in registry.py."""
+
+from __future__ import annotations
+
+from typing import Protocol, Sequence, runtime_checkable
+
+
+@runtime_checkable
+class Embedder(Protocol):
+    """Text → vector. Implementations apply doc/query prefixes themselves."""
+
+    model: str
+
+    def dimension(self) -> int:
+        """Vector dimension. May lazily probe the backend once and cache."""
+        ...
+
+    def embed_documents(self, texts: Sequence[str]) -> list[list[float]]:
+        """Embed corpus chunks (applies doc_prefix if configured). Preserves order."""
+        ...
+
+    def embed_queries(self, texts: Sequence[str]) -> list[list[float]]:
+        """Embed search queries (applies query_prefix if configured). Preserves order."""
+        ...
+
+
+@runtime_checkable
+class Generator(Protocol):
+    """LLM text generation, used only by the optional query-expansion stage."""
+
+    model: str
+
+    def generate(self, system: str, prompt: str, *, max_tokens: int = 1024) -> str:
+        """Single completion. Returns raw text; caller parses any JSON."""
+        ...
+
+
+@runtime_checkable
+class Reranker(Protocol):
+    """Cross-encoder relevance scoring."""
+
+    model: str
+
+    def score(self, query: str, texts: Sequence[str]) -> list[float]:
+        """Relevance score per text, same order. Higher is more relevant."""
+        ...