coderag-ai 0.3.0__py3-none-any.whl

coderag/embeddings/pipeline.py

@@ -0,0 +1,87 @@
+"""Embedding pipeline: embed a repository's symbols, cached by source hash.
+
+Unchanged code is never re-embedded: an embedding is (re)computed only when no
+row exists for (symbol, model, version) or the symbol's ``source_hash`` changed.
+Embedding inputs carry structural context (repo/file/symbol/type/signature/
+docstring/code), not just the raw body (spec §9).
+"""
+
+from __future__ import annotations
+
+from sqlalchemy import delete, select
+from sqlalchemy.orm import Session
+
+from coderag.core.logging import get_logger
+from coderag.db.models import Repository, Symbol, SymbolEmbedding
+from coderag.embeddings.base import EmbeddingProvider, build_embedding_input
+
+log = get_logger("embeddings")
+
+
+class EmbeddingPipeline:
+    def __init__(self, session: Session, provider: EmbeddingProvider) -> None:
+        self.session = session
+        self.provider = provider
+
+    def embed_repository(self, repo: Repository) -> int:
+        model = self.provider.model_name
+        version = self.provider.model_version
+
+        existing: dict[int, str] = {
+            row[0]: row[1]
+            for row in self.session.execute(
+                select(SymbolEmbedding.symbol_id, SymbolEmbedding.source_hash).where(
+                    SymbolEmbedding.repository_id == repo.id,
+                    SymbolEmbedding.embedding_model == model,
+                    SymbolEmbedding.embedding_version == version,
+                )
+            ).all()
+        }
+
+        symbols = list(
+            self.session.scalars(
+                select(Symbol).where(Symbol.repository_id == repo.id)
+            )
+        )
+        stale: list[Symbol] = [
+            s for s in symbols if existing.get(s.id) != s.source_hash
+        ]
+        if not stale:
+            return 0
+
+        # Drop any prior (stale) rows for these symbols before re-inserting.
+        stale_ids = [s.id for s in stale]
+        self.session.execute(
+            delete(SymbolEmbedding).where(
+                SymbolEmbedding.symbol_id.in_(stale_ids),
+                SymbolEmbedding.embedding_model == model,
+                SymbolEmbedding.embedding_version == version,
+            )
+        )
+
+        inputs = [
+            build_embedding_input(
+                repo.name, s.file_path, s.qualified_name, s.symbol_type,
+                s.signature, s.docstring, s.source_code,
+            )
+            for s in stale
+        ]
+        vectors = self.provider.embed_documents(inputs)
+        for s, vec in zip(stale, vectors, strict=True):
+            self.session.add(
+                SymbolEmbedding(
+                    repository_id=repo.id,
+                    symbol_id=s.id,
+                    source_hash=s.source_hash,
+                    embedding_model=model,
+                    embedding_version=version,
+                    embedding_dimension=self.provider.dimension,
+                    embedding=vec,
+                )
+            )
+        self.session.flush()
+        log.info(
+            "embed_repository.done", repository=repo.name, model=model,
+            embedded=len(stale), skipped=len(symbols) - len(stale),
+        )
+        return len(stale)

coderag/embeddings/registry.py

@@ -0,0 +1,38 @@
+"""Embedding provider selection + a tiny cache keyed by configuration."""
+
+from __future__ import annotations
+
+from coderag.core.config import Settings, get_settings
+from coderag.embeddings.base import EmbeddingProvider
+
+_cache: dict[tuple, EmbeddingProvider] = {}
+
+
+def get_embedding_provider(settings: Settings | None = None) -> EmbeddingProvider:
+    settings = settings or get_settings()
+    key = (
+        settings.embedding_provider,
+        settings.embedding_model,
+        settings.embedding_device,
+        settings.embedding_dimension,
+    )
+    if key in _cache:
+        return _cache[key]
+
+    if settings.embedding_provider == "sentence_transformer":
+        from coderag.embeddings.sentence_transformer import (
+            SentenceTransformerEmbeddingProvider,
+        )
+
+        provider: EmbeddingProvider = SentenceTransformerEmbeddingProvider(
+            model_name=settings.embedding_model,
+            device=settings.embedding_device,
+            batch_size=settings.embedding_batch_size,
+        )
+    else:
+        from coderag.embeddings.hashing import HashingEmbeddingProvider
+
+        provider = HashingEmbeddingProvider(dimension=settings.embedding_dimension)
+
+    _cache[key] = provider
+    return provider

coderag/embeddings/sentence_transformer.py

@@ -0,0 +1,47 @@
+"""Local SentenceTransformers embedding provider (optional `embeddings` extra).
+
+Runs entirely on-box — no source code leaves your infrastructure. The model,
+device, and batch size come from configuration.
+"""
+
+from __future__ import annotations
+
+from coderag.embeddings.base import EmbeddingProvider
+
+
+class SentenceTransformerEmbeddingProvider(EmbeddingProvider):
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        device: str = "cpu",
+        batch_size: int = 32,
+    ) -> None:
+        try:
+            from sentence_transformers import SentenceTransformer  # noqa: PLC0415
+        except ImportError as exc:  # pragma: no cover - optional dependency
+            raise RuntimeError(
+                "sentence-transformers is not installed. Install the 'embeddings' "
+                "extra: pip install 'coderag[embeddings]'"
+            ) from exc
+        self._model = SentenceTransformer(model_name, device=device)
+        self.model_name = model_name
+        self.model_version = "st1"
+        self.batch_size = batch_size
+        self.dimension = int(self._model.get_sentence_embedding_dimension())
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        vectors = self._model.encode(
+            texts,
+            batch_size=self.batch_size,
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return [v.tolist() for v in vectors]
+
+    def embed_query(self, text: str) -> list[float]:
+        vector = self._model.encode(
+            [text], normalize_embeddings=True, convert_to_numpy=True,
+            show_progress_bar=False,
+        )[0]
+        return vector.tolist()

coderag/evaluation/datasets.py

@@ -0,0 +1,36 @@
+"""Evaluation dataset format + loader.
+
+A dataset is a JSON list of cases; each case has a ``question`` and the
+``expected_symbols`` that should be retrieved. Expected symbols may be full
+qualified names or suffixes (see matching in ``harness``). Point the eval at your
+own dataset to measure retrieval on your codebase.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from pathlib import Path
+
+DEMO_DATASET = str(
+    Path(__file__).resolve().parents[3] / "examples" / "eval" / "eval_dataset.json"
+)
+
+
+@dataclass
+class EvalCase:
+    question: str
+    expected_symbols: list[str]
+
+
+def load_dataset(path: str) -> list[EvalCase]:
+    data = json.loads(Path(path).read_text())
+    cases = []
+    for item in data:
+        cases.append(
+            EvalCase(
+                question=item["question"],
+                expected_symbols=list(item["expected_symbols"]),
+            )
+        )
+    return cases

coderag/evaluation/harness.py

@@ -0,0 +1,246 @@
+"""Evaluation harness.
+
+Retrieval quality (Recall@K, MRR) + token accounting, and a measured comparison
+of a naive baseline ("send whole files") against the budgeted Code-RAG context.
+We report measured numbers only — no unverified savings claims (spec §18/§19).
+"""
+
+from __future__ import annotations
+
+import statistics
+import time
+from dataclasses import dataclass, field
+
+from sqlalchemy import select
+from sqlalchemy.orm import Session
+
+from coderag.context.builder import ContextBuilder
+from coderag.core.config import Settings, get_settings
+from coderag.core.tokens import get_token_counter
+from coderag.db.models import EvaluationRun, Repository, SourceFile
+from coderag.evaluation.datasets import EvalCase
+from coderag.git.repo import GitRepo
+from coderag.retrieval.engine import build_engine
+
+
+def _matches(qualified: str, expected: str) -> bool:
+    q, e = qualified.lower(), expected.lower()
+    return q == e or q.endswith("." + e) or e.endswith("." + q)
+
+
+def _first_hit_rank(candidate_quals: list[str], expected: list[str]) -> int | None:
+    for rank, qual in enumerate(candidate_quals, start=1):
+        if any(_matches(qual, e) for e in expected):
+            return rank
+    return None
+
+
+@dataclass
+class RetrievalMetrics:
+    n_cases: int = 0
+    recall_at: dict[int, float] = field(default_factory=dict)
+    mrr: float = 0.0
+    avg_candidate_tokens: float = 0.0
+    avg_context_tokens: float = 0.0
+    avg_retrieval_latency_ms: float = 0.0
+    per_case: list[dict] = field(default_factory=list)
+
+
+def evaluate_retrieval(
+    session: Session,
+    repo: Repository,
+    cases: list[EvalCase],
+    settings: Settings | None = None,
+    ks: tuple[int, ...] = (1, 3, 5, 10),
+    *,
+    semantic: bool = True,
+    graph: bool = True,
+    rerank: bool = False,
+) -> RetrievalMetrics:
+    settings = settings or get_settings()
+    engine = build_engine(
+        settings,
+        embedding_provider=_provider(settings) if semantic else None,
+        with_graph=graph,
+        with_reranker=rerank,
+    )
+    builder = ContextBuilder(session, settings=settings)
+
+    hits: dict[int, list[int]] = {k: [] for k in ks}
+    rrs: list[float] = []
+    cand_tokens: list[int] = []
+    ctx_tokens: list[int] = []
+    latencies: list[float] = []
+    per_case: list[dict] = []
+
+    for case in cases:
+        outcome = engine.search(session, repo.id, case.question, top_n=None)
+        quals = [c.qualified_name for c in outcome.candidates]
+        rank = _first_hit_rank(quals, case.expected_symbols)
+        for k in ks:
+            hits[k].append(1 if (rank is not None and rank <= k) else 0)
+        rrs.append(1.0 / rank if rank else 0.0)
+        latencies.append(outcome.latency_ms)
+        cand_tokens.append(sum(c.token_count for c in outcome.candidates))
+        pkg = builder.build(case.question, outcome.candidates, repo)
+        ctx_tokens.append(pkg.accounting.context_tokens)
+        per_case.append({
+            "question": case.question, "first_hit_rank": rank,
+            "candidate_tokens": cand_tokens[-1], "context_tokens": ctx_tokens[-1],
+        })
+
+    n = len(cases)
+    return RetrievalMetrics(
+        n_cases=n,
+        recall_at={k: (sum(hits[k]) / n if n else 0.0) for k in ks},
+        mrr=(sum(rrs) / n if n else 0.0),
+        avg_candidate_tokens=(statistics.mean(cand_tokens) if cand_tokens else 0.0),
+        avg_context_tokens=(statistics.mean(ctx_tokens) if ctx_tokens else 0.0),
+        avg_retrieval_latency_ms=(statistics.mean(latencies) if latencies else 0.0),
+        per_case=per_case,
+    )
+
+
+@dataclass
+class BaselineComparison:
+    n_cases: int
+    baseline_kind: str
+    avg_baseline_tokens: float
+    avg_topfiles_tokens: float
+    avg_rag_context_tokens: float
+    avg_rag_prompt_tokens: float
+    token_reduction_percent: float
+
+
+def compare_baseline(
+    session: Session,
+    repo: Repository,
+    cases: list[EvalCase],
+    settings: Settings | None = None,
+    topfiles: int = 3,
+) -> BaselineComparison:
+    """Compare Code-RAG context against naive baselines.
+
+    Headline baseline is **broad repository context** (send every source file) —
+    the thing you do without retrieval. We also report the "top-N whole files"
+    baseline. Savings scale with repository size: on a tiny demo the whole-repo
+    number is modest, but the *ratio* (RAG context vs the code you'd otherwise
+    paste) is what generalises.
+    """
+    settings = settings or get_settings()
+    engine = build_engine(settings, embedding_provider=_provider(settings), with_graph=True)
+    builder = ContextBuilder(session, settings=settings)
+    counter = get_token_counter(settings)
+    git = GitRepo(repo.local_path)
+
+    file_cache: dict[str, int] = {}
+
+    def file_tokens(path: str) -> int:
+        if path not in file_cache:
+            file_cache[path] = counter.count(git.read_text(path) or "")
+        return file_cache[path]
+
+    all_files = list(
+        session.scalars(
+            select(SourceFile.path).where(SourceFile.repository_id == repo.id)
+        )
+    )
+    whole_repo_tokens = sum(file_tokens(f) for f in all_files)
+
+    topfiles_tok: list[int] = []
+    rag_ctx: list[int] = []
+    rag_prompt: list[int] = []
+    for case in cases:
+        outcome = engine.search(session, repo.id, case.question, top_n=None)
+        files: list[str] = []
+        for c in outcome.candidates:
+            if c.file_path not in files:
+                files.append(c.file_path)
+            if len(files) >= topfiles:
+                break
+        topfiles_tok.append(sum(file_tokens(f) for f in files))
+        pkg = builder.build(case.question, outcome.candidates, repo)
+        rag_ctx.append(pkg.accounting.context_tokens)
+        rag_prompt.append(pkg.accounting.final_prompt_tokens)
+
+    avg_ctx = statistics.mean(rag_ctx) if rag_ctx else 0.0
+    reduction = (
+        round(100.0 * (1.0 - avg_ctx / whole_repo_tokens), 1) if whole_repo_tokens else 0.0
+    )
+    return BaselineComparison(
+        n_cases=len(cases), baseline_kind="whole_repository",
+        avg_baseline_tokens=float(whole_repo_tokens),
+        avg_topfiles_tokens=(statistics.mean(topfiles_tok) if topfiles_tok else 0.0),
+        avg_rag_context_tokens=avg_ctx,
+        avg_rag_prompt_tokens=(statistics.mean(rag_prompt) if rag_prompt else 0.0),
+        token_reduction_percent=reduction,
+    )
+
+
+@dataclass
+class LatencyBenchmark:
+    n: int
+    p50_ms: float
+    p95_ms: float
+    mean_ms: float
+
+
+def benchmark_latency(
+    session: Session, repo: Repository, cases: list[EvalCase],
+    settings: Settings | None = None, repeats: int = 3,
+) -> LatencyBenchmark:
+    settings = settings or get_settings()
+    engine = build_engine(settings, embedding_provider=_provider(settings), with_graph=True)
+    samples: list[float] = []
+    for _ in range(repeats):
+        for case in cases:
+            t = time.perf_counter()
+            engine.search(session, repo.id, case.question, top_n=10)
+            samples.append((time.perf_counter() - t) * 1000)
+    samples.sort()
+
+    def pct(p: float) -> float:
+        if not samples:
+            return 0.0
+        idx = min(len(samples) - 1, int(p * len(samples)))
+        return samples[idx]
+
+    return LatencyBenchmark(
+        n=len(samples), p50_ms=pct(0.50), p95_ms=pct(0.95),
+        mean_ms=(statistics.mean(samples) if samples else 0.0),
+    )
+
+
+def persist_eval_run(
+    session: Session, name: str, dataset: str, metrics: RetrievalMetrics,
+    settings: Settings | None = None,
+) -> EvaluationRun:
+    settings = settings or get_settings()
+    run = EvaluationRun(
+        name=name, dataset=dataset,
+        recall_at_1=metrics.recall_at.get(1, 0.0),
+        recall_at_3=metrics.recall_at.get(3, 0.0),
+        recall_at_5=metrics.recall_at.get(5, 0.0),
+        recall_at_10=metrics.recall_at.get(10, 0.0),
+        mrr=metrics.mrr,
+        avg_retrieved_tokens=metrics.avg_candidate_tokens,
+        avg_context_tokens=metrics.avg_context_tokens,
+        avg_retrieval_latency_ms=metrics.avg_retrieval_latency_ms,
+        config={
+            "rrf_k": settings.rrf_k,
+            "weights": {
+                "symbol": settings.weight_symbol, "lexical": settings.weight_lexical,
+                "semantic": settings.weight_semantic, "graph": settings.weight_graph,
+            },
+            "max_context_tokens": settings.max_context_tokens,
+        },
+    )
+    session.add(run)
+    session.flush()
+    return run
+
+
+def _provider(settings: Settings):
+    from coderag.embeddings.registry import get_embedding_provider
+
+    return get_embedding_provider(settings)

coderag/git/repo.py

@@ -0,0 +1,139 @@
+"""Git-backed repository access.
+
+For git repositories we enumerate with ``git ls-files`` so ``.gitignore`` and the
+tracked/untracked distinction are honoured for free. For plain directories we
+fall back to a filesystem walk. Diffs (base..head) power incremental indexing.
+"""
+
+from __future__ import annotations
+
+import subprocess
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class DiffResult:
+    added: list[str] = field(default_factory=list)
+    modified: list[str] = field(default_factory=list)
+    deleted: list[str] = field(default_factory=list)
+
+    @property
+    def changed(self) -> list[str]:
+        return self.added + self.modified
+
+
+class GitRepo:
+    def __init__(self, path: str | Path) -> None:
+        self.path = Path(path).resolve()
+
+    # -- git helpers ------------------------------------------------------
+    def _git(self, *args: str) -> str:
+        result = subprocess.run(
+            ["git", "-C", str(self.path), *args],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        return result.stdout
+
+    def is_git(self) -> bool:
+        try:
+            out = self._git("rev-parse", "--is-inside-work-tree").strip()
+            return out == "true"
+        except (subprocess.CalledProcessError, FileNotFoundError):
+            return False
+
+    def current_commit(self) -> str | None:
+        try:
+            return self._git("rev-parse", "HEAD").strip()
+        except subprocess.CalledProcessError:
+            return None  # e.g. repo with no commits yet
+
+    def default_branch(self) -> str:
+        try:
+            return self._git("rev-parse", "--abbrev-ref", "HEAD").strip() or "main"
+        except subprocess.CalledProcessError:
+            return "main"
+
+    # -- enumeration ------------------------------------------------------
+    def list_files(self) -> list[str]:
+        """Repository-relative paths of candidate files (POSIX separators).
+
+        For git repos we include tracked files *and* untracked files that are
+        not gitignored (``--exclude-standard``), so a working tree with
+        uncommitted files still indexes correctly while honouring .gitignore.
+        """
+        if self.is_git():
+            tracked = self._git("ls-files", "-z").split("\0")
+            untracked = self._git(
+                "ls-files", "-z", "--others", "--exclude-standard"
+            ).split("\0")
+            seen: dict[str, None] = {}
+            for p in (*tracked, *untracked):
+                if p:
+                    seen[p] = None
+            return list(seen)
+        return self._walk()
+
+    def _walk(self) -> list[str]:
+        from coderag.indexing.ignore import DEFAULT_IGNORE_DIRS
+
+        files: list[str] = []
+        for p in sorted(self.path.rglob("*")):
+            if p.is_dir():
+                continue
+            rel = p.relative_to(self.path).as_posix()
+            if set(rel.split("/")[:-1]) & DEFAULT_IGNORE_DIRS:
+                continue
+            files.append(rel)
+        return files
+
+    # -- diff (incremental) ----------------------------------------------
+    def diff(self, base: str, head: str = "HEAD") -> DiffResult:
+        out = self._git("diff", "--name-status", "-z", base, head)
+        return _parse_name_status(out)
+
+    # -- reading ----------------------------------------------------------
+    def read_text(self, rel_path: str) -> str | None:
+        fp = self.path / rel_path
+        try:
+            data = fp.read_bytes()
+        except (OSError, FileNotFoundError):
+            return None
+        from coderag.indexing.ignore import looks_binary
+
+        if looks_binary(data):
+            return None
+        return data.decode("utf-8", errors="replace")
+
+
+def _parse_name_status(z: str) -> DiffResult:
+    """Parse ``git diff --name-status -z`` output.
+
+    Records are NUL-separated. Renames/copies (R/C) emit the status token, then
+    the old path, then the new path as separate records.
+    """
+    res = DiffResult()
+    tokens = [t for t in z.split("\0") if t != ""]
+    i = 0
+    while i < len(tokens):
+        status = tokens[i]
+        code = status[0]
+        if code in ("R", "C"):
+            old, new = tokens[i + 1], tokens[i + 2]
+            res.deleted.append(old)
+            res.added.append(new)
+            i += 3
+        else:
+            path = tokens[i + 1]
+            if code == "A":
+                res.added.append(path)
+            elif code == "M":
+                res.modified.append(path)
+            elif code == "D":
+                res.deleted.append(path)
+            else:  # T (type change), etc. -> treat as modified
+                res.modified.append(path)
+            i += 2
+    return res