diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index 3f359dc..ccfe407 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -3,7 +3,7 @@ "name": "codebase-index", "displayName": "Codebase Index", "description": "Give Claude a precise local map of your codebase: find implementations, trace behavior, and predict change impact with file-line evidence.", - "version": "1.7.0", + "version": "1.8.0", "author": { "name": "codebase-index contributors" }, diff --git a/.claude/skills/codebase-index/.skill_version b/.claude/skills/codebase-index/.skill_version index bd8bf88..27f9cd3 100644 --- a/.claude/skills/codebase-index/.skill_version +++ b/.claude/skills/codebase-index/.skill_version @@ -1 +1 @@ -1.7.0 +1.8.0 diff --git a/.codex/skills/codebase-index/.skill_version b/.codex/skills/codebase-index/.skill_version index bd8bf88..27f9cd3 100644 --- a/.codex/skills/codebase-index/.skill_version +++ b/.codex/skills/codebase-index/.skill_version @@ -1 +1 @@ -1.7.0 +1.8.0 diff --git a/.opencode/skills/codebase-index/.skill_version b/.opencode/skills/codebase-index/.skill_version index bd8bf88..27f9cd3 100644 --- a/.opencode/skills/codebase-index/.skill_version +++ b/.opencode/skills/codebase-index/.skill_version @@ -1 +1 @@ -1.7.0 +1.8.0 diff --git a/CHANGELOG.md b/CHANGELOG.md index a8e224b..fe931d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,26 @@ All notable changes to this project are documented here. The format is based on ## [Unreleased] +## [1.8.0] - 2026-09-02 + +### Added + +- **Reproducible retrieval evaluation.** Added a self-repository ground-truth + query suite with Recall@K, MRR, nDCG, hit rate, precision, MAP, useful-context, + latency percentiles, and one-signal ablations. + +### Changed + +- **Hybrid retrieval quality.** Natural-language lexical queries now use safe, + down-weighted identifier/synonym expansion and soft term coverage; exact symbol + lookup preserves framing-word tolerance and bounded fuzzy matching. +- **Packaging compatibility.** Cap the build backend below the Metadata 2.5 + default until the release validation toolchain supports that metadata version. +- **Ranking defaults are evidence-driven.** Implementation/test/documentation + source priors are calibrated from the benchmark. Graph propagation is bounded + and intent-directed; graph and MMR signals remain opt-in because ablations + reduced direct retrieval quality on the reproducible corpus. + ## [1.7.0] - 2026-07-29 ### Added @@ -401,7 +421,8 @@ All notable changes to this project are documented here. The format is based on - Hooks example + `watch` mode for keeping the index fresh without blocking the edit loop (M8). - `doctor`, `stats`, `clean` diagnostics/maintenance commands. -[Unreleased]: https://github.com/denfry/codebase-index/compare/v1.7.0...HEAD +[Unreleased]: https://github.com/denfry/codebase-index/compare/v1.8.0...HEAD +[1.8.0]: https://github.com/denfry/codebase-index/compare/v1.7.0...v1.8.0 [1.7.0]: https://github.com/denfry/codebase-index/compare/v1.6.0...v1.7.0 [1.6.0]: https://github.com/denfry/codebase-index/compare/v1.5.0...v1.6.0 [1.5.0]: https://github.com/denfry/codebase-index/compare/v1.4.0...v1.5.0 diff --git a/README.md b/README.md index 9f31a88..4b43c70 100644 --- a/README.md +++ b/README.md @@ -271,7 +271,7 @@ See [MCP.md](docs/MCP.md) for client configuration. ## Project status -The latest released line is **1.7.0**. It includes: +The latest released line is **1.8.0**. It includes: - hybrid and optional vector retrieval; - Tree-sitter symbol extraction across the documented language tiers; @@ -279,6 +279,8 @@ The latest released line is **1.7.0**. It includes: - architecture communities, central nodes, and surprising cross-module links; - shortest dependency paths and node descriptions; - token-budgeted and skeletonized retrieval packets; +- benchmark-calibrated lexical expansion, fuzzy identifier matching, and source-aware ranking; +- bounded, intent-directed graph discovery with optional diversity and duplicate suppression; - CLI, Skill, plugin, and MCP delivery; - incremental updates, watch hooks, diagnostics, skill rollback, and diff-aware impact analysis. diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index 0117f29..9c4bbf0 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -2,7 +2,7 @@ ## 1. Overview -`codebase-index` is a **local-first** code intelligence layer for AI coding agents. In `1.7.0` +`codebase-index` is a **local-first** code intelligence layer for AI coding agents. In `1.8.0` it has two shipped faces: 1. **A Claude Code Skill** (`.claude/skills/codebase-index/SKILL.md`) that Claude auto-invokes for diff --git a/docs/FAQ.md b/docs/FAQ.md index 7fd8bde..1e25c0d 100644 --- a/docs/FAQ.md +++ b/docs/FAQ.md @@ -14,7 +14,7 @@ pip install codebase-index # or: pipx install codebase-index ``` To pin an exact version or grab an unreleased commit, install from a GitHub tag -instead: `pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.7.0"`. +instead: `pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.8.0"`. Then run `codebase-index init` inside your project and `codebase-index index` to build the first index. In Claude Code you can instead install the plugin @@ -162,11 +162,12 @@ Yes. Use any of these methods: ## Is it production-ready? -Yes — `codebase-index` is released as **v1.7.0**. The core indexing and search -functionality is implemented and tested. The current `1.7.0` package includes: +Yes — `codebase-index` is released as **v1.8.0**. The core indexing and search +functionality is implemented and tested. The current `1.8.0` package includes: -- Hybrid FTS/path/symbol/vector retrieval -- Import/call/reference graph expansion and `impact` +- Hybrid FTS/path/symbol/vector retrieval with benchmark-calibrated lexical expansion + and bounded fuzzy identifier matching +- Import/call/reference graph expansion, intent-directed graph discovery, and `impact` - Diff-aware blast-radius analysis for tracked working-tree changes - Optional local embeddings, with external embeddings gated behind explicit opt-in - Hooks and watch mode for freshness diff --git a/docs/INSTALLATION.md b/docs/INSTALLATION.md index 8f3ba48..695cf68 100644 --- a/docs/INSTALLATION.md +++ b/docs/INSTALLATION.md @@ -67,7 +67,7 @@ pipx install codebase-index # isolated environment uv tool install codebase-index # uv-managed tool # Pin to a GitHub tag for an exact or unreleased version -pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.7.0" +pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.8.0" # From source (editable mode) git clone https://github.com/denfry/codebase-index.git @@ -90,7 +90,7 @@ pip install -e ".[embeddings-local,watch,dev]" ### uvx / Homebrew status -As of `1.7.0`, **PyPI is shipped** — `pip install codebase-index` and +As of `1.8.0`, **PyPI is shipped** — `pip install codebase-index` and `pipx install codebase-index` are the verified paths. `uvx codebase-index init`, Homebrew tap installation, signed checksums, and SBOMs remain distribution targets for a more complete release story. @@ -130,7 +130,7 @@ Expected output: === codebase-index Doctor === [OK] Python 3.12 (requires 3.11+) -[OK] codebase-index package installed (v1.7.0) +[OK] codebase-index package installed (v1.8.0) [OK] tree-sitter is available [INFO] Cache directory not yet created: ... [INFO] Skill not installed in .claude/skills/ diff --git a/docs/RETRIEVAL.md b/docs/RETRIEVAL.md index 5b925bb..d4acf37 100644 --- a/docs/RETRIEVAL.md +++ b/docs/RETRIEVAL.md @@ -2,7 +2,8 @@ The retrieval engine turns a natural-language or symbolic query into a **compact, ranked, token-budgeted** set of file/line ranges for Claude to read. It is hybrid: multiple independent -retrievers run, their results are fused, reranked, expanded via the graph, then trimmed. +retrievers run, their results are fused and reranked, then trimmed. Graph expansion and MMR are +available as bounded opt-in signals; the shipped default prioritizes direct evidence. ``` query @@ -21,13 +22,16 @@ query [3] rank fusion ── Reciprocal Rank Fusion (RRF) across retriever result lists │ ▼ -[4] rerank ── feature-based score (symbol-kind, path proximity, recency, centrality) +[4] rerank ── feature-based score (symbol-kind, path, source role, centrality) │ ▼ -[5] graph expansion ── pull in imports/callers/callees per intent (bounded) +[5] optional graph expansion ── pull in imports/callers/callees per intent (bounded) │ ▼ -[6] token budgeting ── greedy fill under --token-budget; snippets trimmed + secret-redacted +[6] diversity / duplicate filtering (MMR opt-in, SimHash duplicate guard) + │ + ▼ +[7] token budgeting ── greedy fill under --token-budget; snippets trimmed + secret-redacted │ ▼ ranked results + recommended_reads + fallback_suggestions @@ -59,10 +63,12 @@ source)` list so fusion is source-agnostic. - **Path** — exact and glob path matches (`src/auth/*.py`, `auth.py`). Highest precision; surfaced first when the query clearly names a path. -- **Symbol** — query against `symbols` (name exact, prefix, then fuzzy/trigram). Carries `kind` - (function/class/method/...) and signature. Primary for `locate_impl` / `find_refs`. +- **Symbol** — query against `symbols` (exact, identifier parts, bounded fuzzy matching). Carries + `kind` (function/class/method/...) and signature. Primary for `locate_impl` / `find_refs`. - **FTS** — FTS5 `bm25()` over the `fts_chunks` virtual table (chunk text + symbol names + - summaries indexed). Tokenizer is code-aware (splits camelCase/snake_case). Primary lexical signal. + summaries indexed). Query-time camelCase/snake_case splitting, small down-weighted synonym + expansion, and soft coverage scoring make natural-language questions robust without weakening + exact terms. - **Vector** *(opt-in)* — cosine similarity over chunk embeddings via `sqlite-vec`. Only runs if `embeddings.enabled = true`. Adds semantic recall for paraphrased queries. Absent → pipeline degrades gracefully to FTS+symbol. @@ -96,19 +102,30 @@ A lightweight, explainable feature score (no external model required) layered on The reranker also produces the human-readable **`reason`** string per result (e.g. *"exact symbol match · called by 4 sites · in src/auth/"*). -## 5. Graph expansion (`graph/expand.py`) +## 5. Graph expansion (`graph/retrieval.py`; `graph/expand.py` for impact APIs) + +Graph expansion runs only when the tuning enables `graph_source` and the intent plan requests a +graph strategy. It is disabled by default because the reproducible self-repository ablation +reduced direct-hit MRR when related nodes displaced lexical hits. -After reranking, pull in *related* nodes per the intent's graph strategy, bounded by `--depth` -(default 1–2) and a node cap: +When enabled, it is bounded by depth and node cap: - `impact` → walk **up** edges (callers, importers) = blast radius. - `how_it_works` → walk **down** edges (callees, imported defs) = mechanism. -- `find_refs` → direct reverse edges only. -- `data_flow` → both directions along call/assignment edges. +- `find_refs` → walk **up** edges to callers/importers. +- `data_flow` → walk **both** directions along call/assignment edges. + +Expanded nodes retain edge confidence and receive distance-decayed scores so seeds stay on top. + +## 6. Diversity and duplicate control -Expanded nodes are merged into results with a discounted score so seeds stay on top. +`retrieval.diversity` provides bounded MMR selection and SimHash near-duplicate +suppression. MMR is disabled in the shipped default because the reproducible +benchmark favored relevance-only ranking; callers that need broader snippet +coverage can enable `RetrievalTuning(mmr=True)`. Duplicate suppression remains +available independently. -## 6. Token budgeting (`retrieval/budget.py`) +## 7. Token budgeting (`retrieval/budget.py`) Results are trimmed to fit `--token-budget` (default per intent, e.g. 1500 tokens): @@ -121,7 +138,7 @@ Results are trimmed to fit `--token-budget` (default per intent, e.g. 1500 token The point: Claude gets enough to decide, and a precise list of what to read next — never a dump. -## 7. Confidence & fallback +## 8. Confidence & fallback A `confidence` score (high/medium/low) is derived from: top RRF score, score gap between #1 and #2, number of agreeing retrievers, and whether a symbol matched exactly. @@ -131,13 +148,12 @@ number of agreeing retrievers, and whether a symbol matched exactly. - **low** → skill instructs Claude to **fall back** to `ripgrep`/Grep/Glob with suggested patterns emitted in `fallback_suggestions` (derived from query terms + detected symbols). -## 8. Output payload (shared by Markdown + JSON) +## 9. Output payload (shared by Markdown + JSON) ```jsonc { "query": "where is auth token refresh implemented", "intent": "locate_impl", - "index": { "exists": true, "stale": false, "built_at": "...", "head_commit": "abc1234" }, "confidence": "high", "results": [ { diff --git a/docs/RETRIEVAL_PIPELINE.md b/docs/RETRIEVAL_PIPELINE.md index 4e2901d..2a843e1 100644 --- a/docs/RETRIEVAL_PIPELINE.md +++ b/docs/RETRIEVAL_PIPELINE.md @@ -18,7 +18,7 @@ Intent detection (keyword / symbol / impact / general) │ 2. Path-based search │ │ 3. SQLite FTS5 lexical search │ │ 4. Vector search (optional embeddings) │ -│ 5. Graph expansion (from seed results) │ +│ 5. Graph expansion (explicit opt-in) │ └─────────────────────────────────────────┘ ↓ Reciprocal Rank Fusion (RRF) @@ -53,15 +53,16 @@ Ranked retrieval packet with confidence score ## 3. SQLite FTS5 Lexical Search -**Trigger:** General keyword queries. +**Trigger:** General keyword and natural-language queries. **Process:** -- Build an FTS5 query from the user's text -- Tokenize: split `snake_case`, expand `camelCase` at query time -- Search the `fts_chunks` virtual table -- Return matching chunks with BM25-style scores +- Parse identifiers into camelCase/PascalCase/snake_case subtokens. +- Add a small, explicit synonym/inflection vocabulary at lower weight. +- Use OR groups for soft matching, then require bounded term coverage and rank + by original-term coverage plus BM25. +- Quote every FTS term so query punctuation cannot inject MATCH operators. -**Score:** Based on FTS5 rank — higher for more term matches and rarer terms. +**Score:** Coverage is the primary signal; BM25 is a bounded tie-break. ## 4. Vector Search (Optional) @@ -80,52 +81,74 @@ Ranked retrieval packet with confidence score ## 5. Graph Expansion -**Trigger:** After initial results are found. +**Trigger:** `RetrievalTuning(graph_source=True)` and an intent plan with a +graph strategy. It is disabled in the shipped default because the reproducible +self-repository ablation reduced direct-hit MRR. **Process:** -- For each seed result, traverse the dependency/call graph -- Find related files: callers, callees, imports, inheritors -- Add related files with a decay factor (distance from seed) +- Seed from lexical/symbol candidates already found in SQLite. +- Traverse only indexed, resolved edges with bounded depth and node count. +- Follow `up` (callers/importers), `down` (callees/imports), or `both` + according to the intent plan. +- Apply distance decay and preserve edge confidence in the candidate reason. -**Score:** Decreases with graph distance — direct connections score higher. +Graph expansion is an opt-in context-discovery signal, not a replacement for +direct lexical or symbol evidence. -## Reciprocal Rank Fusion (RRF) +## 6. Diversity and duplicate control -Combines results from multiple retrievers: +SimHash suppresses near-duplicate snippets independently of MMR. Bounded +Maximal Marginal Relevance is available through `RetrievalTuning(mmr=True)`; +the shipped default keeps relevance-only ordering because the benchmark favored +direct hits. + +## 7. Reciprocal Rank Fusion (RRF) + +Combines ranked lists from the enabled retrievers: ``` -RRF_score(d) = Σ (1 / (k + rank_r(d))) +RRF_score(d) = Σ w_r · k / (k + rank_r(d)) ``` +The implementation multiplies textbook RRF by `k` so fusion and bounded rerank +bonuses share a comparable scale; ordering is unchanged. + Where: - `k` is a constant (default 60) - `rank_r(d)` is the rank of document `d` in retriever `r` -- Sum is over all retrievers that returned `d` +- `w_r` is the intent/tuning weight for retriever `r` -This ensures documents that appear in multiple retrievers rank higher. +The implementation merges co-located chunks into one per-file bucket before +fusion, preventing a large file from dominating the result list. -## Reranking +## 8. Reranking -After fusion, apply additional boosts: +After fusion, apply bounded explainable boosts and penalties: -| Factor | Boost | Rationale | -|---|---|---| -| Exact symbol match | +0.3 | User named a specific symbol | -| File type relevance | +0.1 | `.ts` for TypeScript queries, etc. | -| Recency | +0.05 | Recently modified files may be more relevant | -| File size | -0.05 per 10KB | Prefer focused files over large ones | +| Factor | Effect | Rationale | +|---|---:|---| +| Exact symbol match | +0.20 | User named a specific symbol | +| Symbol definition kind | +0.05 | Prefer actionable definitions | +| Path term match | +0.05 | User supplied a location clue | +| Degree / reference evidence | up to +0.08 | Stable structural tiebreaker | +| Implementation source prior | +0.08 | Prefer source over prose/tests | +| Documentation source prior | -0.05 | Avoid docs displacing implementation | +| Generated/vendor/build | -0.12 | Suppress low-value derived code | +| Test path on non-test query | -0.06 | Keep tests as supporting evidence | -## Confidence Score +## 9. Confidence -The final confidence score (0.0 to 1.0) determines how Claude should proceed: +Confidence is categorical (`high`, `medium`, `low`) and is derived from +exact-symbol evidence, multi-retriever agreement, score separation, and result +count. Exact symbol matches are `high`, including a single-result response. -| Confidence | Meaning | Action | -|---|---|---| -| 0.8 - 1.0 | High | Read recommended ranges and answer directly | -| 0.5 - 0.8 | Medium | Read ranges; optionally confirm with one Grep | -| 0.0 - 0.5 | Low | Use fallback suggestions (ripgrep, Glob) | +| Confidence | Action | +|---|---| +| `high` | Read recommended ranges and answer directly | +| `medium` | Read ranges; optionally confirm with one Grep | +| `low` | Use fallback suggestions (ripgrep, Glob) | -## Token Budget Enforcement +## 10. Token Budget Enforcement The output is capped at a configurable token budget: @@ -134,9 +157,9 @@ The output is capped at a configurable token budget: 3. Remaining results are listed without snippets 4. The `recommended_reads` field contains only the most critical line ranges -Default budget: 2000 tokens (configurable in `.codeindex.json`). +Default budget: 1500 tokens (configurable in `.codeindex.json`). -## Fallback Suggestions +## 11. Fallback Suggestions When confidence is low, the pipeline generates fallback strategies: diff --git a/docs/SEO.md b/docs/SEO.md index 18b832b..bd867d3 100644 --- a/docs/SEO.md +++ b/docs/SEO.md @@ -177,7 +177,7 @@ Features: - Secret redaction - Respects .gitignore -Install: pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.7.0" +Install: pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.8.0" GitHub: https://github.com/denfry/codebase-index ``` @@ -192,5 +192,5 @@ codebase-index builds a local hybrid index so Claude finds the right files witho - No network by default - Token-efficient output -pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.7.0" +pip install "codebase-index @ git+https://github.com/denfry/codebase-index.git@v1.8.0" ``` diff --git a/docs/installer.md b/docs/installer.md index 689fda3..10fddd1 100644 --- a/docs/installer.md +++ b/docs/installer.md @@ -108,7 +108,7 @@ pwsh ./install.ps1 -Target claude -InstallDir "D:\skills\codebase-index" **Pinning по ветке/тегу** (воспроизводимость и безопасность): ```sh -sh install.sh --branch v1.7.0 +sh install.sh --branch v1.8.0 ``` --- @@ -151,8 +151,7 @@ sh install.sh --branch v1.7.0 ```json { "skill_name": "codebase-index", - "version": "1.7.0", - "installed_at": "2026-05-29T12:00:00Z", + "version": "1.8.0", "target": "claude", "os": "linux", "source_repo": "https://github.com/denfry/codebase-index", diff --git a/pyproject.toml b/pyproject.toml index 35ec711..c967628 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["hatchling"] +requires = ["hatchling<1.32"] build-backend = "hatchling.build" [project] diff --git a/requirements.lock b/requirements.lock index 679ed63..4a600d6 100644 --- a/requirements.lock +++ b/requirements.lock @@ -1,3 +1,3 @@ -codebase-index @ https://github.com/denfry/codebase-index/archive/refs/tags/v1.7.0.tar.gz +codebase-index @ https://github.com/denfry/codebase-index/archive/refs/tags/v1.8.0.tar.gz tree-sitter==0.25.2 tree-sitter-language-pack==1.8.1 diff --git a/src/codebase_index/__init__.py b/src/codebase_index/__init__.py index 62be96d..f4bccfd 100644 --- a/src/codebase_index/__init__.py +++ b/src/codebase_index/__init__.py @@ -4,4 +4,4 @@ See docs/ARCHITECTURE.md for the module map. """ -__version__ = "1.7.0" +__version__ = "1.8.0" diff --git a/src/codebase_index/graph/expand.py b/src/codebase_index/graph/expand.py index 02ee989..ba75a0f 100644 --- a/src/codebase_index/graph/expand.py +++ b/src/codebase_index/graph/expand.py @@ -12,6 +12,7 @@ from __future__ import annotations +import math import sqlite3 from collections import deque from typing import Optional @@ -20,6 +21,44 @@ from ..storage import repo +_CONFIDENCE_RANK = { + "extracted": 0, + "inferred": 1, + "ambiguous": 2, +} + + +def _validated_decay(decay: float) -> float: + """Return a finite distance-decay factor in the inclusive [0, 1] range.""" + try: + value = float(decay) + except (TypeError, ValueError) as exc: + raise ValueError("decay must be a finite number between 0 and 1") from exc + if not math.isfinite(value) or not 0.0 <= value <= 1.0: + raise ValueError("decay must be a finite number between 0 and 1") + return value + + +def _confidence_rank(confidence: Optional[str]) -> int: + """Rank edge confidence conservatively (lower is better).""" + return _CONFIDENCE_RANK.get(confidence or "", len(_CONFIDENCE_RANK)) + + +def _set_optional_impact_strength(node: ImpactNode, strength: float) -> None: + """Populate an optional model field without changing the shared schema. + + Older ImpactNode models intentionally have no score field. If a consumer + supplies a compatible model with an ``impact_strength`` field, expose the + calculated value; otherwise distance remains the public ranking signal. + """ + fields = getattr(type(node), "model_fields", None) + if fields is None: + fields = getattr(type(node), "__fields__", {}) + if "impact_strength" in fields: + setattr(node, "impact_strength", strength) + + + def _freshness(conn: sqlite3.Connection) -> IndexFreshness: return IndexFreshness( exists=True, @@ -78,33 +117,98 @@ def _node_meta(conn, kind, node_id) -> Optional[ImpactNode]: line_start=row["line_start"], distance=0) +def _neighbor_sort_key(neighbor: tuple[str, int, Optional[str], Optional[str]]) -> tuple: + """Make edge traversal independent of SQLite's unspecified row order.""" + kind, node_id, edge_type, confidence = neighbor + return ( + _confidence_rank(confidence), + kind, + node_id, + edge_type or "", + confidence or "", + ) + + +def _impact_sort_key(node: ImpactNode) -> tuple: + """Rank direct impact before transitive impact, then resolve ties safely.""" + return ( + node.distance, + _confidence_rank(node.via_confidence), + node.path, + node.kind, + node.name or "", + node.line_start if node.line_start is not None else -1, + node.via_edge or "", + ) + + def walk_impact( - conn: sqlite3.Connection, target: str, *, depth: int, direction: str + conn: sqlite3.Connection, + target: str, + *, + depth: int, + direction: str, + decay: float = 1.0, ) -> list[ImpactNode]: + """Return bounded impact nodes ranked by hop distance. + + ``decay`` controls the optional strength score as ``decay ** (distance - 1)``: + direct nodes score 1.0 and transitive nodes score progressively less. The + current ImpactNode schema has no score field, so its observable output uses + stable distance ordering (direct before transitive) and confidence as a tie + breaker. A compatible model that declares ``impact_strength`` receives it. + """ + decay_value = _validated_decay(decay) seeds = _seed_nodes(conn, target) if not seeds: return [] - visited: set[tuple[str, int]] = set(seeds) + + seed_keys = set(seeds) queue: deque[tuple[str, int, int]] = deque((k, i, 0) for k, i in seeds) - out: list[ImpactNode] = [] + states: dict[tuple[str, int], tuple[int, int]] = { + key: (0, 0) for key in seed_keys + } + nodes: dict[tuple[str, int], ImpactNode] = {} while queue: kind, node_id, dist = queue.popleft() + state = states.get((kind, node_id)) + if state is None or state[0] != dist: + continue if dist >= depth: continue - for nk, nid, etype, conf in _neighbors(conn, kind, node_id, direction): - if (nk, nid) in visited: + for nk, nid, etype, conf in sorted( + _neighbors(conn, kind, node_id, direction), key=_neighbor_sort_key + ): + key = (nk, nid) + if key in seed_keys: continue - visited.add((nk, nid)) + next_dist = dist + 1 + candidate_state = (next_dist, _confidence_rank(conf)) + current_state = states.get(key) + if current_state is not None: + # Shorter paths always win. At equal distance, only a more + # confident edge may replace the existing audit trail. + if candidate_state >= current_state: + continue + meta = _node_meta(conn, nk, nid) if meta is None: continue - meta.distance = dist + 1 + meta.distance = next_dist meta.via_edge = etype meta.via_confidence = conf - out.append(meta) - queue.append((nk, nid, dist + 1)) - return out + _set_optional_impact_strength( + meta, decay_value ** max(next_dist - 1, 0) + ) + states[key] = candidate_state + nodes[key] = meta + # Equal-distance confidence improvements do not need another walk; + # shorter paths can change the reachable frontier and are queued. + if current_state is None or next_dist < current_state[0]: + queue.append((nk, nid, next_dist)) + + return sorted(nodes.values(), key=_impact_sort_key) def _target_paths(conn: sqlite3.Connection, target: str) -> list[str]: @@ -119,11 +223,17 @@ def _target_paths(conn: sqlite3.Connection, target: str) -> list[str]: return [suffix[0]["path"]] return [] - def impact_lookup( - conn: sqlite3.Connection, target: str, *, depth: int, direction: str + conn: sqlite3.Connection, + target: str, + *, + depth: int, + direction: str, + decay: float = 1.0, ) -> ImpactResponse: - nodes = walk_impact(conn, target, depth=depth, direction=direction) + nodes = walk_impact( + conn, target, depth=depth, direction=direction, decay=decay + ) best: dict[str, int] = {} for n in nodes: if n.path not in best or n.distance < best[n.path]: diff --git a/src/codebase_index/graph/retrieval.py b/src/codebase_index/graph/retrieval.py new file mode 100644 index 0000000..ebef077 --- /dev/null +++ b/src/codebase_index/graph/retrieval.py @@ -0,0 +1,456 @@ +"""Bounded graph-based candidate retrieval. + +The retriever deliberately works from the lexical/symbol candidates already found by +other retrievers. It only follows indexed, resolved SQLite edges from those seeds; +it never enumerates the repository's complete node set. +""" + +from __future__ import annotations + +import math +import sqlite3 +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import Iterable, Optional + +from ..retrieval.types import Candidate + +Node = tuple[str, int] + +# Confidence is both a ranking signal and an audit trail. A malformed/unknown value +# is treated conservatively rather than being presented as an exact relationship. +_CONFIDENCE_WEIGHT = { + "extracted": 1.0, + "inferred": 0.75, + "ambiguous": 0.35, +} + + +@dataclass(frozen=True) +class _Edge: + source: Node + target: Node + edge_type: str + confidence: str + line: Optional[int] + edge_id: int + direction: str + + +def _safe_float(value: object, default: float) -> float: + try: + if not isinstance(value, (str, bytes, bytearray, int, float)): + return default + number = float(value) + except (TypeError, ValueError): + return default + return number if math.isfinite(number) else default + + +def _safe_int(value: object, default: int = 0) -> int: + try: + if not isinstance(value, (str, bytes, bytearray, int, float)): + return default + return int(value) + except (TypeError, ValueError): + return default + +def _row_value(row: object, key: str, index: int, default: object = None) -> object: + """Read sqlite Row and tuple rows alike (useful for small test connections).""" + try: + if isinstance(row, sqlite3.Row): + return row[key] + return row[index] # type: ignore[index] + except (IndexError, KeyError, TypeError): + return default + + +def _execute(conn: sqlite3.Connection, sql: str, params: tuple[object, ...]) -> list[object]: + """Execute a read, making absent/partial tables a harmless empty graph.""" + try: + return list(conn.execute(sql, params).fetchall()) + except (AttributeError, sqlite3.Error, TypeError, ValueError): + return [] + + +def _edge_execute( + conn: sqlite3.Connection, + primary_sql: str, + fallback_sql: str, + params: tuple[object, ...], +) -> list[object]: + """Read edge rows, tolerating pre-confidence/partial edge tables.""" + try: + return list(conn.execute(primary_sql, params).fetchall()) + except (AttributeError, sqlite3.Error, TypeError, ValueError): + return _execute(conn, fallback_sql, params) + + + + +def _seed_nodes(conn: sqlite3.Connection, candidate: Candidate) -> list[Node]: + path = getattr(candidate, "path", None) + if not isinstance(path, str) or not path: + return [] + files = _execute(conn, "SELECT id FROM files WHERE path = ? LIMIT 1", (path,)) + if not files: + return [] + file_id = _safe_int(_row_value(files[0], "id", 0), -1) + if file_id < 0: + return [] + + symbol_name = getattr(candidate, "symbol", None) + if isinstance(symbol_name, str) and symbol_name: + line = _safe_int(getattr(candidate, "line_start", 0), 0) + rows = _execute( + conn, + """ + SELECT s.id + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE f.path = ? AND s.name = ? + ORDER BY CASE WHEN s.line_start = ? THEN 0 ELSE 1 END, + s.line_start, s.id + LIMIT 1 + """, + (path, symbol_name, line), + ) + if not rows: + rows = _execute( + conn, + """ + SELECT s.id + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE f.path = ? AND s.name = ? + ORDER BY s.id + LIMIT 1 + """, + (path, symbol_name), + ) + if rows: + symbol_id = _safe_int(_row_value(rows[0], "id", 0), -1) + if symbol_id >= 0: + return [("symbol", symbol_id)] + return [("file", file_id)] + + +def _edge_rows( + conn: sqlite3.Connection, + node: Node, + *, + limit: int, + direction: str, +) -> list[_Edge]: + kind, node_id = node + # Separate indexed lookups retain SQLite's idx_edges_src/idx_edges_dst plans. + incoming: list[object] = [] + outgoing: list[object] = [] + if direction in ("up", "both"): + incoming = _edge_execute( + conn, + """ + SELECT id, edge_type, src_kind, src_id, confidence, line + FROM edges + WHERE resolved = 1 AND dst_kind = ? AND dst_id = ? + ORDER BY src_kind, src_id, edge_type, id + LIMIT ? + """, + """ + SELECT id, edge_type, src_kind, src_id, NULL AS confidence, line + FROM edges + WHERE resolved = 1 AND dst_kind = ? AND dst_id = ? + ORDER BY src_kind, src_id, edge_type, id + LIMIT ? + """, + (kind, node_id, limit), + ) + if direction in ("down", "both"): + outgoing = _edge_execute( + conn, + """ + SELECT id, edge_type, dst_kind, dst_id, confidence, line + FROM edges + WHERE resolved = 1 AND src_kind = ? AND src_id = ? AND dst_id IS NOT NULL + ORDER BY dst_kind, dst_id, edge_type, id + LIMIT ? + """, + """ + SELECT id, edge_type, dst_kind, dst_id, NULL AS confidence, line + FROM edges + WHERE resolved = 1 AND src_kind = ? AND src_id = ? AND dst_id IS NOT NULL + ORDER BY dst_kind, dst_id, edge_type, id + LIMIT ? + """, + (kind, node_id, limit), + ) + edges: list[_Edge] = [] + for row in incoming: + nk = _row_value(row, "src_kind", 2) + ni = _safe_int(_row_value(row, "src_id", 3), -1) + if nk not in ("file", "symbol") or ni < 0: + continue + edges.append( + _Edge( + source=(str(nk), ni), + target=node, + edge_type=str(_row_value(row, "edge_type", 1) or "unknown"), + confidence=str(_row_value(row, "confidence", 4) or "unknown").lower(), + line=(lambda x: _safe_int(x) if x is not None else None)(_row_value(row, "line", 5)), + edge_id=_safe_int(_row_value(row, "id", 0), -1), + direction="incoming", + ) + ) + for row in outgoing: + nk = _row_value(row, "dst_kind", 2) + ni = _safe_int(_row_value(row, "dst_id", 3), -1) + if nk not in ("file", "symbol") or ni < 0: + continue + edges.append( + _Edge( + source=node, + target=(str(nk), ni), + edge_type=str(_row_value(row, "edge_type", 1) or "unknown"), + confidence=str(_row_value(row, "confidence", 4) or "unknown").lower(), + line=(lambda x: _safe_int(x) if x is not None else None)(_row_value(row, "line", 5)), + edge_id=_safe_int(_row_value(row, "id", 0), -1), + direction="outgoing", + ) + ) + # A malformed database can expose duplicate rows through both indexes. Keep all + # provenance but make traversal order stable. + edges.sort(key=lambda e: (e.target, e.edge_type, e.confidence, e.edge_id, e.direction)) + return edges + + +def _node_metadata(conn: sqlite3.Connection, node: Node) -> Optional[dict[str, object]]: + kind, node_id = node + if kind == "file": + rows = _execute( + conn, "SELECT path, is_generated, summary FROM files WHERE id = ? LIMIT 1", (node_id,) + ) + if not rows: + rows = _execute( + conn, + "SELECT path, 0 AS is_generated, NULL AS summary " + "FROM files WHERE id = ? LIMIT 1", + (node_id,), + ) + if not rows: + return None + row = rows[0] + path = _row_value(row, "path", 0) + if not isinstance(path, str) or not path: + return None + return { + "path": path, + "line_start": 1, + "line_end": 1, + "kind": "file", + "symbol": None, + "content": _row_value(row, "summary", 2), + "is_generated": bool(_row_value(row, "is_generated", 1, 0)), + } + + rows = _execute( + conn, + """ + SELECT s.name, s.kind, s.line_start, s.line_end, s.signature, + f.path, f.is_generated + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE s.id = ? LIMIT 1 + """, + (node_id,), + ) + if not rows: + rows = _execute( + conn, + """ + SELECT s.name, s.kind, s.line_start, s.line_end, NULL AS signature, + f.path, 0 AS is_generated + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE s.id = ? LIMIT 1 + """, + (node_id,), + ) + if not rows: + rows = _execute( + conn, + """ + SELECT s.name, 'symbol' AS kind, 1 AS line_start, 1 AS line_end, + NULL AS signature, f.path, 0 AS is_generated + FROM symbols AS s JOIN files AS f ON f.id = s.file_id + WHERE s.id = ? LIMIT 1 + """, + (node_id,), + ) + if not rows: + return None + row = rows[0] + path = _row_value(row, "path", 5) + name = _row_value(row, "name", 0) + if not isinstance(path, str) or not path or not isinstance(name, str): + return None + return { + "path": path, + "line_start": max(1, _safe_int(_row_value(row, "line_start", 2), 1)), + "line_end": max(1, _safe_int(_row_value(row, "line_end", 3), 1)), + "kind": _row_value(row, "kind", 1), + "symbol": name, + "content": _row_value(row, "signature", 4), + "is_generated": bool(_row_value(row, "is_generated", 6, 0)), + } + + +def _provenance(edges: Iterable[_Edge]) -> str: + parts: list[str] = [] + for edge in sorted(edges, key=lambda e: (e.edge_type, e.confidence, e.line or 0, e.edge_id)): + confidence = edge.confidence if edge.confidence in _CONFIDENCE_WEIGHT else "unknown" + line = f", line={edge.line}" if edge.line is not None else "" + direction = "<-" if edge.direction == "incoming" else "->" + parts.append(f"{direction}{edge.edge_type} (confidence={confidence}{line})") + return "; ".join(dict.fromkeys(parts)) or "unknown edge provenance" + + +def graph_candidates( + conn: sqlite3.Connection, + seeds: Iterable[Candidate], + *, + depth: int = 2, + node_cap: int = 50, + damping: float = 0.85, + iterations: int = 12, + direction: str = "both", +) -> list[Candidate]: + """Return bounded, deterministic graph neighbors of lexical/symbol *seeds*. + + ``direction`` controls traversal from each seed: ``up`` follows incoming + edges (callers/importers), ``down`` follows outgoing edges (callees/imports), + and ``both`` preserves the general related-code behavior. + Personalized PageRank is evaluated only on the bounded depth-limited + subgraph. Seed candidates themselves are excluded. + """ + if conn is None: + return [] + if direction not in {"up", "down", "both"}: + raise ValueError("direction must be 'up', 'down', or 'both'") + depth = max(0, _safe_int(depth)) + node_cap = max(0, _safe_int(node_cap)) + iterations = max(0, _safe_int(iterations)) + damping = min(1.0, max(0.0, _safe_float(damping, 0.85))) + if depth == 0 or node_cap == 0 or iterations == 0: + return [] + + seed_scores: dict[Node, float] = defaultdict(float) + seed_keys: set[tuple[str, int, int]] = set() + for candidate in seeds or (): + if not isinstance(candidate, Candidate): + continue + try: + seed_keys.add(candidate.key()) + except (AttributeError, TypeError): + pass + nodes = _seed_nodes(conn, candidate) + score = max(0.0, _safe_float(getattr(candidate, "score", 0.0), 0.0)) + for node in nodes: + seed_scores[node] = max(seed_scores[node], score) + if not seed_scores: + return [] + + # Candidate lists are normally already small. Sorting here makes results + # independent of input order and bounds fan-out when a caller supplies many hits. + ordered_seeds = sorted(seed_scores, key=lambda n: n) + seed_total = sum(seed_scores.values()) + if seed_total <= 0.0: + teleport = {node: 1.0 / len(ordered_seeds) for node in ordered_seeds} + else: + teleport = {node: seed_scores[node] / seed_total for node in ordered_seeds} + + distances: dict[Node, int] = {node: 0 for node in ordered_seeds} + adjacency: dict[Node, dict[Node, float]] = defaultdict(dict) + via: dict[Node, list[_Edge]] = defaultdict(list) + queue: deque[Node] = deque(ordered_seeds) + # Seeds are retained; node_cap limits graph nodes returned, while this bound + # prevents a depth walk from accumulating an unbounded frontier. + max_discovered = len(ordered_seeds) + node_cap + while queue and len(distances) <= max_discovered: + node = queue.popleft() + distance = distances[node] + if distance >= depth: + continue + for edge in _edge_rows( + conn, node, limit=max(1, node_cap), direction=direction + ): + neighbor = edge.source if edge.direction == "incoming" else edge.target + if neighbor not in distances: + if len(distances) >= max_discovered: + continue + distances[neighbor] = distance + 1 + queue.append(neighbor) + via[neighbor].append(edge) + elif distances[neighbor] == distance + 1: + via[neighbor].append(edge) + weight = _CONFIDENCE_WEIGHT.get(edge.confidence, 0.5) + adjacency[node][neighbor] = adjacency[node].get(neighbor, 0.0) + weight + if direction == "both": + adjacency[neighbor][node] = adjacency[neighbor].get(node, 0.0) + weight + + if not adjacency: + return [] + + rank = dict(teleport) + for node in distances: + rank.setdefault(node, 0.0) + for _ in range(iterations): + next_rank = {node: (1.0 - damping) * teleport.get(node, 0.0) for node in rank} + dangling = 0.0 + for node, score in rank.items(): + neighbors = adjacency.get(node, {}) + total = sum(neighbors.values()) + if total <= 0.0: + dangling += score + continue + for neighbor, weight in neighbors.items(): + next_rank[neighbor] = next_rank.get(neighbor, 0.0) + damping * score * weight / total + if dangling: + for node, probability in teleport.items(): + next_rank[node] = next_rank.get(node, 0.0) + damping * dangling * probability + rank = next_rank + + out: list[tuple[float, Candidate]] = [] + seen: set[tuple[str, int, int]] = set(seed_keys) + for node, score in rank.items(): + if node in seed_scores or node not in via: + continue + metadata = _node_metadata(conn, node) + if metadata is None: + continue + provenance = _provenance(via[node]) + base_content = metadata["content"] if isinstance(metadata["content"], str) else "" + content = f"{base_content}\n[graph provenance: {provenance}]" if base_content else ( + f"[graph provenance: {provenance}]" + ) + candidate = Candidate( + path=str(metadata["path"]), + line_start=_safe_int(metadata["line_start"], 1), + line_end=_safe_int(metadata["line_end"], 1), + source="graph", + score=float(score), + kind=metadata["kind"] if isinstance(metadata["kind"], str) else None, + symbol=metadata["symbol"] if isinstance(metadata["symbol"], str) else None, + content=content, + is_generated=bool(metadata["is_generated"]), + reason=( + f"graph propagation distance={distances[node]}; " + f"provenance={provenance}" + ), + ) + if candidate.key() in seen: + continue + seen.add(candidate.key()) + out.append((score, candidate)) + + out.sort(key=lambda item: (-item[0], item[1].path, item[1].line_start, item[1].line_end)) + return [candidate for _, candidate in out[:node_cap]] + + +# Descriptive alias for callers that prefer the operation-oriented name. +retrieve_graph_candidates = graph_candidates diff --git a/src/codebase_index/mcp/server.py b/src/codebase_index/mcp/server.py index c842311..6540297 100644 --- a/src/codebase_index/mcp/server.py +++ b/src/codebase_index/mcp/server.py @@ -30,7 +30,7 @@ from ..config import Config try: - from mcp.server.fastmcp import FastMCP + from mcp.server.fastmcp import FastMCP # type: ignore[attr-defined] except ImportError as exc: # pragma: no cover raise ImportError( "MCP server needs the optional extra: pip install codebase-index[mcp]" diff --git a/src/codebase_index/retrieval/diversity.py b/src/codebase_index/retrieval/diversity.py new file mode 100644 index 0000000..dbbb2a5 --- /dev/null +++ b/src/codebase_index/retrieval/diversity.py @@ -0,0 +1,289 @@ +"""Deterministic diversity helpers for retrieval candidates. + +The module deliberately has no index, model, or repository dependencies. Code +snippets are reduced to lexical tokens, making formatting and comments +irrelevant to near-duplicate detection. MMR uses the same compact lexical +features when no caller-provided similarity function is supplied. +""" + +from __future__ import annotations + +import hashlib +import math +from collections.abc import Callable, Iterable, Sequence + +from .types import Candidate + +_MASK64 = (1 << 64) - 1 + + +def normalize_code_tokens(content: str | None) -> tuple[str, ...]: + """Return normalized code tokens, omitting whitespace and comments.""" + if not content: + return () + + tokens: list[str] = [] + length = len(content) + index = 0 + operators = ( + "===", "!==", ">>>=", "**=", "...", "=>", "->", "::", "==", "!=", "<=", ">=", + "&&", "||", "++", "--", "+=", "-=", "*=", "/=", "%=", "<<", ">>", "**", "??", + ) + + while index < length: + char = content[index] + if char.isspace(): + index += 1 + continue + if content.startswith("//", index) or char == "#": + newline = content.find("\n", index + (2 if content.startswith("//", index) else 1)) + index = length if newline < 0 else newline + 1 + continue + if content.startswith("/*", index): + end = content.find("*/", index + 2) + index = length if end < 0 else end + 2 + continue + + if char in "'\"`": + quote = char + triple = content.startswith(char * 3, index) + width = 3 if triple else 1 + cursor = index + width + while cursor < length: + if content[cursor] == "\\": + cursor += 2 + continue + if content.startswith(quote * width, cursor): + cursor += width + break + cursor += 1 + tokens.append("str:" + content[index + width : max(index + width, cursor - width)]) + index = cursor + continue + + if char.isalpha() or char == "_" or ord(char) >= 128: + cursor = index + 1 + while cursor < length and ( + content[cursor].isalnum() or content[cursor] == "_" or ord(content[cursor]) >= 128 + ): + cursor += 1 + tokens.append(content[index:cursor].lower()) + index = cursor + continue + + if char.isdigit(): + cursor = index + 1 + while cursor < length and (content[cursor].isalnum() or content[cursor] in "._"): + cursor += 1 + tokens.append(content[index:cursor].lower()) + index = cursor + continue + + operator = next((candidate for candidate in operators if content.startswith(candidate, index)), None) + if operator is not None: + tokens.append(operator) + index += len(operator) + else: + tokens.append(char) + index += 1 + + return tuple(tokens) + + +def token_fingerprint(tokens: Iterable[str] | str | None) -> int: + """Compute a stable unsigned 64-bit SimHash for tokens or source content.""" + if tokens is None: + return 0 + if isinstance(tokens, str): + tokens = normalize_code_tokens(tokens) + values = tuple(tokens) + if not values: + return 0 + + weights = [0] * 64 + for token in values: + digest = hashlib.blake2b(str(token).encode("utf-8"), digest_size=8).digest() + hashed = int.from_bytes(digest, "big", signed=False) + for bit in range(64): + weights[bit] += 1 if (hashed >> bit) & 1 else -1 + + fingerprint = 0 + for bit, weight in enumerate(weights): + if weight >= 0: + fingerprint |= 1 << bit + return fingerprint & _MASK64 + + +simhash = token_fingerprint + + +def simhash_distance(left: int, right: int) -> int: + """Return the Hamming distance between two 64-bit fingerprints.""" + return ((int(left) ^ int(right)) & _MASK64).bit_count() + + +def _features(candidate: Candidate) -> tuple[tuple[str, ...], frozenset[str], int | None]: + tokens = normalize_code_tokens(candidate.content) + return tokens, frozenset(tokens), token_fingerprint(tokens) if tokens else None + + +def _score_value(candidate: Candidate) -> float: + score = float(candidate.score) + return score if math.isfinite(score) else float("-inf") + + +def _same_or_better(left: Candidate, right: Candidate) -> bool: + """Whether left is the representative to retain; ties favor input order.""" + return _score_value(left) > _score_value(right) + + +def deduplicate(candidates: Sequence[Candidate], hamming_distance: int = 3) -> list[Candidate]: + """Suppress near-identical snippets, retaining the highest-scoring hit.""" + if not candidates: + return [] + threshold = max(0, min(64, int(hamming_distance))) + representatives: list[Candidate] = [] + fingerprints: list[int | None] = [] + + for candidate in candidates: + _, _, fingerprint = _features(candidate) + if fingerprint is None: + representatives.append(candidate) + fingerprints.append(None) + continue + + matches = [ + index + for index, existing in enumerate(fingerprints) + if existing is not None and simhash_distance(fingerprint, existing) <= threshold + ] + if not matches: + representatives.append(candidate) + fingerprints.append(fingerprint) + continue + + first = matches[0] + winner = candidate + for index in matches: + incumbent = representatives[index] + if _same_or_better(incumbent, winner): + winner = incumbent + for index in reversed(matches): + representatives.pop(index) + fingerprints.pop(index) + representatives.insert(first, winner) + fingerprints.insert(first, token_fingerprint(normalize_code_tokens(winner.content))) + + return representatives + + +def _normalized_relevance(candidates: Sequence[Candidate]) -> list[float]: + scores = [_score_value(candidate) for candidate in candidates] + finite = [score for score in scores if math.isfinite(score)] + if not finite: + return [0.0] * len(candidates) + low, high = min(finite), max(finite) + if high <= low: + return [1.0 if math.isfinite(score) else 0.0 for score in scores] + return [ + max(0.0, min(1.0, (score - low) / (high - low))) if math.isfinite(score) else 0.0 + for score in scores + ] + + +def _similarity_from_features( + left: Candidate, + right: Candidate, + left_features: tuple[tuple[str, ...], frozenset[str], int | None], + right_features: tuple[tuple[str, ...], frozenset[str], int | None], +) -> float: + _, left_set, left_hash = left_features + _, right_set, right_hash = right_features + if left_set or right_set: + union = left_set | right_set + lexical = len(left_set & right_set) / len(union) if union else 0.0 + sim = 1.0 - simhash_distance(left_hash or 0, right_hash or 0) / 64.0 + value = 0.75 * lexical + 0.25 * sim + else: + value = 0.0 + + if left.path == right.path: + value = max(value, 0.35) + if left.symbol and right.symbol and left.symbol == right.symbol: + value = max(value, 0.45) + return max(0.0, min(1.0, value)) + + +def candidate_similarity(left: Candidate, right: Candidate) -> float: + """Fallback similarity from lexical Jaccard/SimHash plus structure.""" + return _similarity_from_features(left, right, _features(left), _features(right)) + + +def mmr_select( + candidates: Sequence[Candidate], + limit: int, + lambda_: float, + *, + similarity: Callable[[Candidate, Candidate], float] | None = None, +) -> list[Candidate]: + """Select up to ``limit`` candidates with maximal marginal relevance.""" + if not candidates or limit <= 0: + return [] + count = min(int(limit), len(candidates)) + weight = float(lambda_) + weight = 0.5 if not math.isfinite(weight) else max(0.0, min(1.0, weight)) + relevance = _normalized_relevance(candidates) + feature_cache = [_features(candidate) for candidate in candidates] + compare = similarity + remaining = set(range(len(candidates))) + selected: list[int] = [] + + while remaining and len(selected) < count: + best_index: int | None = None + best_key: tuple[float, float, int] | None = None + for index in sorted(remaining): + if compare is None: + redundancy = max( + ( + _similarity_from_features( + candidates[index], + candidates[chosen], + feature_cache[index], + feature_cache[chosen], + ) + for chosen in selected + ), + default=0.0, + ) + else: + redundancy = max( + ( + max( + 0.0, + min(1.0, float(compare(candidates[index], candidates[chosen]))), + ) + for chosen in selected + ), + default=0.0, + ) + utility = weight * relevance[index] - (1.0 - weight) * redundancy + key = (utility, relevance[index], -index) + if best_key is None or key > best_key: + best_key = key + best_index = index + assert best_index is not None + remaining.remove(best_index) + selected.append(best_index) + + return [candidates[index] for index in selected] + + +__all__ = [ + "candidate_similarity", + "deduplicate", + "mmr_select", + "normalize_code_tokens", + "simhash", + "simhash_distance", + "token_fingerprint", +] diff --git a/src/codebase_index/retrieval/fuzzy.py b/src/codebase_index/retrieval/fuzzy.py new file mode 100644 index 0000000..95246c9 --- /dev/null +++ b/src/codebase_index/retrieval/fuzzy.py @@ -0,0 +1,278 @@ +"""Deterministic fuzzy matching for symbol identifiers. + +The module deliberately works on one identifier pair (or a caller-provided +candidate sequence). It does not inspect the repository or perform any I/O; +searchers can therefore decide how many rows to fetch before reranking. +""" + +from __future__ import annotations + +import math +import re +import unicodedata +from collections.abc import Mapping, Sequence +from typing import Any, TypeVar + + +_IDENTIFIER_PART_RE = re.compile(r"[^\W_]+", re.UNICODE) +_CAMEL_PART_RE = re.compile( + r"[A-Z]+(?=[A-Z][a-z]|[0-9]|$)|[A-Z]?[a-z]+|[0-9]+" +) +_MAX_EDIT_INPUT = 256 + +T = TypeVar("T") + + +def _fold(value: str) -> str: + """Return a case-folded, compatibility-normalized string.""" + return unicodedata.normalize("NFKC", value).casefold() + + +def _tokens(value: str) -> tuple[str, ...]: + """Split an identifier into case-insensitive word and camel-case parts.""" + normalized = unicodedata.normalize("NFKC", value) + parts: list[str] = [] + for segment in _IDENTIFIER_PART_RE.findall(normalized): + camel_parts = _CAMEL_PART_RE.findall(segment) + parts.extend(_fold(part) for part in (camel_parts or [segment]) if part) + return tuple(parts) + + +def _compact(value: str) -> str: + """Remove identifier separators while retaining Unicode letters/digits.""" + folded = _fold(value) + return "".join(char for char in folded if char.isalnum()) + + +def _acronym(tokens: tuple[str, ...]) -> str: + return "".join(token[0] for token in tokens if token) + + +def _bounded_levenshtein_similarity(left: str, right: str, max_distance: int) -> float: + """Compute normalized Levenshtein similarity with a bounded edit window.""" + if left == right: + return 1.0 + if not left or not right: + return 0.0 + longest = max(len(left), len(right)) + if longest > _MAX_EDIT_INPUT: + # Containment and token features still work for long names, while this + # guard prevents an accidental quadratic allocation for pathological input. + return 0.0 + if abs(len(left) - len(right)) > max_distance: + return 0.0 + + # Keep the shorter value on the columns to reduce memory use. + if len(left) < len(right): + left, right = right, left + width = len(right) + previous = list(range(width + 1)) + infinity = max_distance + longest + 1 + for row, left_char in enumerate(left, 1): + current = [infinity] * (width + 1) + current[0] = row + lo = max(1, row - max_distance) + hi = min(width, row + max_distance) + row_min = current[0] + for column in range(lo, hi + 1): + substitution = previous[column - 1] + (left_char != right[column - 1]) + insertion = current[column - 1] + 1 + deletion = previous[column] + 1 + value = min(substitution, insertion, deletion) + current[column] = value + row_min = min(row_min, value) + if row_min > max_distance: + return 0.0 + previous = current + + distance = previous[width] + if distance > max_distance: + return 0.0 + return max(0.0, 1.0 - distance / longest) + + +def _ordered_subtoken_coverage(query: str, name_tokens: tuple[str, ...]) -> tuple[float, float]: + """Measure name subtokens found in a concatenated query. + + This is what makes ``userid`` match ``getUserById``: ``user`` and ``id`` + occur in order even though the candidate has an intervening ``by`` token. + The second return value is the fraction of candidate subtokens covered. + """ + if not query or not name_tokens: + return 0.0, 0.0 + cursor = 0 + matched_chars = 0 + matched_tokens = 0 + for token in name_tokens: + if len(token) < 2: + continue + position = query.find(token, cursor) + if position < 0: + continue + cursor = position + len(token) + matched_chars += len(token) + matched_tokens += 1 + return matched_chars / len(query), matched_tokens / len(name_tokens) + + +def _token_overlap(query_tokens: tuple[str, ...], name_tokens: tuple[str, ...]) -> float: + if not query_tokens or not name_tokens: + return 0.0 + covered = 0 + for query_token in query_tokens: + match_size = 0 + for name_token in name_tokens: + if query_token == name_token or query_token in name_token: + match_size = max(match_size, len(query_token)) + elif name_token in query_token: + match_size = max(match_size, len(name_token)) + covered += match_size + query_size = sum(len(token) for token in query_tokens) + return covered / query_size if query_size else 0.0 + + +def identifier_similarity( + query: str, + name: str, + *, + max_edit_distance: int = 64, +) -> float: + """Return deterministic identifier similarity in the inclusive range [0, 1]. + + The score combines case-folded/separator-insensitive equality, token and + concatenation containment, prefixes, acronyms, and bounded edit similarity. + A score is not an exactness claim; callers must keep their separate + ``exact_symbol``/database exact flag unchanged. + """ + if not isinstance(query, str) or not isinstance(name, str): + return 0.0 + if max_edit_distance < 0: + raise ValueError("max_edit_distance must be non-negative") + + query_tokens = _tokens(query) + name_tokens = _tokens(name) + query_compact = _compact(query) + name_compact = _compact(name) + if not query_compact or not name_compact: + return 0.0 + if query_compact == name_compact: + return 1.0 + # One-character identifiers are too ambiguous for non-exact matching. + if min(len(query_compact), len(name_compact)) < 2: + return 0.0 + + score = 0.0 + name_acronym = _acronym(name_tokens) + query_is_acronym = len(query) >= 2 and query.strip().isupper() and len(query_compact) >= 2 + if query_compact == name_acronym and len(query_compact) >= 2: + score = max(score, 0.94) + elif len(query_compact) >= 3 and name_acronym.startswith(query_compact): + score = max(score, 0.77 + 0.15 * len(query_compact) / len(name_acronym)) + + if query_compact in name_compact: + ratio = len(query_compact) / len(name_compact) + score = max(score, (0.72 if name_compact.startswith(query_compact) else 0.64) + 0.22 * ratio) + elif name_compact in query_compact: + ratio = len(name_compact) / len(query_compact) + score = max(score, 0.60 + 0.22 * ratio) + + overlap = _token_overlap(query_tokens, name_tokens) + if overlap: + score = max(score, 0.32 + 0.34 * overlap) + + query_fit, candidate_token_ratio = _ordered_subtoken_coverage(query_compact, name_tokens) + if query_fit >= 0.8: + score = max(score, 0.76 + 0.15 * query_fit * candidate_token_ratio) + elif query_fit: + score = max(score, 0.25 + 0.55 * query_fit * candidate_token_ratio) + + edit = _bounded_levenshtein_similarity(query_compact, name_compact, max_edit_distance) + # Edit distance alone can make a short query look deceptively close to a + # much longer identifier merely because they share a common token. Keep + # useful typo matches, but dampen length-mismatched pairs; containment and + # token features above remain responsible for those matches. + length_ratio = min(len(query_compact), len(name_compact)) / max( + len(query_compact), len(name_compact) + ) + if length_ratio < 0.6: + edit *= length_ratio / 0.6 + # An all-uppercase query is generally an acronym; near-edit matches to an + # unrelated acronym should not pass a fuzzy threshold by edit distance alone. + if query_is_acronym: + edit *= 0.55 + score = max(score, edit) + return min(1.0, max(0.0, score)) + + +def _candidate_value(candidate: Any, key: str, default: Any = None) -> Any: + if isinstance(candidate, Mapping): + return candidate.get(key, default) + try: + return candidate[key] + except (KeyError, IndexError, TypeError): + return getattr(candidate, key, default) + + +def _candidate_exact(candidate: Any) -> bool: + """Read either the public Candidate flag or a repository row's exact flag.""" + value = _candidate_value(candidate, "exact_symbol", None) + if value is None: + value = _candidate_value(candidate, "is_exact", False) + return bool(value) + + +def _candidate_name(candidate: Any) -> str: + value = _candidate_value(candidate, "symbol", None) + if not value: + value = _candidate_value(candidate, "name", "") + return value if isinstance(value, str) else "" + + +def rank_fuzzy_symbols( + query: str, + candidates: Sequence[T], + *, + threshold: float = 0.55, + limit: int | None = None, + max_edit_distance: int = 64, +) -> list[T]: + """Filter and rank a provided candidate list by identifier similarity. + + Candidates may be ``Candidate`` objects, mappings/SQLite rows with a + ``symbol`` or ``name`` field, or small compatible records. Returned values + are the original objects (not copies), and their exactness flags are never + changed. Ties use score, pre-existing score, identifier, path, location, + and original position to make ordering reproducible across runs. + """ + if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0: + raise ValueError("threshold must be finite and between 0 and 1") + if limit is not None and limit < 0: + raise ValueError("limit must be non-negative or None") + if limit == 0: + return [] + + ranked: list[tuple[tuple[Any, ...], T]] = [] + for position, candidate in enumerate(candidates): + name = _candidate_name(candidate) + score = identifier_similarity(query, name, max_edit_distance=max_edit_distance) + if score < threshold: + continue + exact = _candidate_exact(candidate) + base_score = _candidate_value(candidate, "score", 0.0) + try: + base_score = float(base_score) + except (TypeError, ValueError): + base_score = 0.0 + path = str(_candidate_value(candidate, "path", "")).casefold() + line_start = _candidate_value(candidate, "line_start", 0) + line_end = _candidate_value(candidate, "line_end", 0) + identifier = name.casefold() + key = (-int(exact), -score, -base_score, identifier, path, line_start, line_end, position) + ranked.append((key, candidate)) + + ranked.sort(key=lambda item: item[0]) + result = [candidate for _, candidate in ranked] + return result if limit is None else result[:limit] + + +__all__ = ["identifier_similarity", "rank_fuzzy_symbols"] diff --git a/src/codebase_index/retrieval/intent.py b/src/codebase_index/retrieval/intent.py index 19e58e2..cefe865 100644 --- a/src/codebase_index/retrieval/intent.py +++ b/src/codebase_index/retrieval/intent.py @@ -1,7 +1,7 @@ """Cheap rule-first intent classifier (regex/keyword heuristics). Each intent maps to retriever weights over {"path","symbol","fts"}, a default -token budget, and a graph strategy (consumed later by M5). +token budget, and an opt-in graph strategy consumed by the retrieval pipeline. """ from __future__ import annotations diff --git a/src/codebase_index/retrieval/lexical.py b/src/codebase_index/retrieval/lexical.py new file mode 100644 index 0000000..ca05d5a --- /dev/null +++ b/src/codebase_index/retrieval/lexical.py @@ -0,0 +1,409 @@ +"""Safe, code-aware lexical query parsing. + +The lexical retriever needs two things that a plain ``str.split`` cannot provide: +identifier-aware variants (``getUserById`` -> ``get``, ``user``, ``by``, ``id``) +and a deliberately tiny vocabulary of programming synonyms. This module keeps +those concerns independent of SQLite/FTS so callers can apply their own ranking +and ablation policy. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import math +import re +from types import MappingProxyType +from typing import Literal, Mapping + + +# A token may contain the separators used by identifiers. Other punctuation is +# a boundary, which also means arbitrary FTS operators never reach the builder. +_TOKEN_RE = re.compile(r"[^\W_]+(?:[_-][^\W_]+)*", re.UNICODE) + +# Keep this list conservative: words such as ``get``, ``set`` and ``run`` can +# be real identifiers and must not be discarded from a code search. +LEXICAL_STOPWORDS = frozenset( + { + "the", + "a", + "an", + "is", + "are", + "was", + "were", + "be", + "been", + "being", + "how", + "does", + "do", + "did", + "what", + "where", + "which", + "who", + "whom", + "when", + "why", + "to", + "of", + "in", + "on", + "for", + "and", + "or", + "with", + "from", + "it", + "this", + "that", + "these", + "those", + "into", + "during", + "if", + "via", + "across", + "between", + "about", + "their", + "its", + # Common query framing, not useful lexical evidence. + "work", + "works", + "working", + "function", + "functions", + "method", + "methods", + "class", + "classes", + "interface", + "interfaces", + "enum", + "enums", + "type", + "types", + } +) + +# Values are tuples rather than sets to make expansion order deterministic. +# The map is intentionally explicit and small; it is not a general thesaurus. +SYNONYM_MAP: Mapping[str, tuple[str, ...]] = MappingProxyType( + { + "auth": ("authentication",), + "authentication": ("auth",), + "config": ("configuration",), + "configuration": ("config",), + "delete": ("remove",), + "remove": ("delete",), + "create": ("add", "insert"), + "add": ("create", "insert"), + "insert": ("create", "add"), + "repo": ("repository",), + "repository": ("repo",), + "retry": ("retries",), + "retries": ("retry",), + "caller": ("callers",), + "callers": ("caller",), + "test": ("tests",), + "tests": ("test",), + # Small, explicit inflection bridges for natural-language questions. + "redacted": ("redact",), + "redaction": ("redact",), + "secrets": ("secret",), + "files": ("file",), + "chosen": ("choose",), + "produced": ("produce",), + "raises": ("raise",), + "parsing": ("parse",), + "loaded": ("load",), + "merged": ("merge", "merg"), + "merging": ("merge", "merg"), + "ranking": ("rank",), + } +) +TermKind = Literal["original", "subtoken", "synonym"] + + +@dataclass(frozen=True, slots=True) +class WeightedTerm: + """A lexical term variant and its relative ranking weight.""" + + term: str + weight: float + kind: TermKind + origin: str + + +@dataclass(frozen=True, slots=True) +class LexicalQuery: + """Parsed query with original terms kept distinct from expansions. + + ``original_terms`` are salient, normalized query tokens. ``expanded_terms`` + never contains an original term and is ordered by source term, then + subtokens before synonyms. Ranking code can use :attr:`weighted_terms` to + preserve the stronger score of originals without parsing query strings. + """ + + original_terms: tuple[str, ...] + expanded_terms: tuple[WeightedTerm, ...] + + @property + def weighted_terms(self) -> tuple[WeightedTerm, ...]: + originals = tuple( + WeightedTerm(term=term, weight=1.0, kind="original", origin=term) + for term in self.original_terms + ) + return originals + self.expanded_terms + + @property + def terms(self) -> tuple[str, ...]: + """All terms in deterministic ranking order, originals first per source.""" + return tuple(item.term for item in self.weighted_terms) + + @property + def weights(self) -> Mapping[str, float]: + """Highest applicable weight per term, useful for score accumulation.""" + out: dict[str, float] = {} + for item in self.weighted_terms: + out[item.term] = max(out.get(item.term, 0.0), item.weight) + return MappingProxyType(out) + + def to_fts(self, *, include_expansions: bool = True) -> str: + return build_fts_query(self, include_expansions=include_expansions) + + +def _is_identifier_char(char: str) -> bool: + # ``isalnum`` handles non-ASCII letters and decimal digits without relying + # on locale or third-party Unicode tables. + return char.isalnum() + + +def _camel_boundary(previous: str, current: str, following: str | None) -> bool: + """Whether ``current`` starts a new camel/Pascal component.""" + if not current.isupper(): + return False + if previous.islower() or previous.isdigit(): + return True + # ``HTTPServer`` -> ``HTTP``, ``Server``: the S starts a word because the + # acronym run is ending before a lowercase letter. + return previous.isupper() and following is not None and following.islower() + + +def split_identifier(identifier: str) -> tuple[str, ...]: + """Split common identifier forms into lowercase components. + + Underscores and hyphens are separators. Case transitions cover camelCase, + PascalCase and acronym runs, while non-ASCII alphanumeric characters are + retained. Concatenations such as ``userid`` are intentionally not guessed + apart: without a dictionary, doing so creates noisy false positives. + """ + + if not identifier: + return () + + parts: list[str] = [] + current: list[str] = [] + + def flush() -> None: + if current: + parts.append("".join(current).casefold()) + current.clear() + + for index, char in enumerate(identifier): + if char in "_-" or not _is_identifier_char(char): + flush() + continue + if current: + previous = current[-1] + following = identifier[index + 1] if index + 1 < len(identifier) else None + if _camel_boundary(previous, char, following): + flush() + current.append(char) + flush() + return tuple(part for part in parts if part) + + +def salient_terms(query: str) -> tuple[str, ...]: + """Return normalized, deduplicated query terms worth lexical matching.""" + + out: list[str] = [] + seen: set[str] = set() + for token in _TOKEN_RE.findall(query): + term = token.casefold() + if len(term) < 2 or term in LEXICAL_STOPWORDS or term in seen: + continue + seen.add(term) + out.append(term) + return tuple(out) + + +def _bounded_weight(value: float) -> float: + try: + value = float(value) + except (TypeError, ValueError): + return 0.0 + if not math.isfinite(value): + return 0.0 + return max(0.0, min(value, 0.999999)) + + +def build_lexical_query( + query: str, + *, + synonym_weight: float = 0.35, + subtoken_weight: float = 0.65, + include_synonyms: bool = True, + include_subtokens: bool = True, + tuning: object | None = None, +) -> LexicalQuery: + """Parse ``query`` and produce deterministic, down-weighted expansions. + + ``tuning`` is optional to avoid coupling this pure module to a config type; + when supplied, ``query_expansion`` and ``expansion_weight`` attributes are + honored. Callers may still explicitly disable either expansion family. + """ + + if tuning is not None: + if hasattr(tuning, "query_expansion"): + include_synonyms = include_synonyms and bool(tuning.query_expansion) + if hasattr(tuning, "expansion_weight"): + synonym_weight = float(tuning.expansion_weight) + + raw_terms: dict[str, str] = {} + for token in _TOKEN_RE.findall(query): + normalized = token.casefold() + if len(normalized) >= 2 and normalized not in LEXICAL_STOPWORDS: + raw_terms.setdefault(normalized, token) + originals = salient_terms(query) + original_set = set(originals) + expanded: list[WeightedTerm] = [] + emitted: set[str] = set(original_set) + sub_weight = _bounded_weight(subtoken_weight) + syn_weight = _bounded_weight(synonym_weight) + + for original in originals: + identifier = raw_terms.get(original, original) + if include_subtokens: + for subtoken in split_identifier(identifier): + if subtoken in emitted or len(subtoken) < 2: + continue + emitted.add(subtoken) + expanded.append( + WeightedTerm( + term=subtoken, + weight=sub_weight, + kind="subtoken", + origin=original, + ) + ) + if include_synonyms: + for synonym in SYNONYM_MAP.get(original, ()): + if synonym in emitted or len(synonym) < 2: + continue + emitted.add(synonym) + expanded.append( + WeightedTerm( + term=synonym, + weight=syn_weight, + kind="synonym", + origin=original, + ) + ) + + return LexicalQuery(original_terms=originals, expanded_terms=tuple(expanded)) + + +def expansion_weights( + query: str | LexicalQuery, + *, + synonym_weight: float = 0.35, + subtoken_weight: float = 0.65, + include_synonyms: bool = True, + include_subtokens: bool = True, + tuning: object | None = None, +) -> Mapping[str, float]: + """Return the strongest relative weight for each parsed lexical term.""" + + parsed = ( + query + if isinstance(query, LexicalQuery) + else build_lexical_query( + query, + synonym_weight=synonym_weight, + subtoken_weight=subtoken_weight, + include_synonyms=include_synonyms, + include_subtokens=include_subtokens, + tuning=tuning, + ) + ) + return parsed.weights + + +def escape_fts_term(term: str) -> str: + """Quote one FTS5 term so operators and quotes are treated literally.""" + + # FTS5 phrase syntax escapes an embedded quote by doubling it. NUL is not a + # valid SQLite string character in all bindings, so replace it defensively. + return '"' + term.replace("\x00", " ").replace('"', '""') + '"' + + +def build_fts_query( + query: str | LexicalQuery, + *, + include_expansions: bool = True, + synonym_weight: float = 0.35, + subtoken_weight: float = 0.65, + include_synonyms: bool = True, + include_subtokens: bool = True, + tuning: object | None = None, +) -> str: + """Build a safe FTS5 expression with originals first in each OR group. + + FTS5 MATCH itself has no portable per-term boost. The returned expression + therefore groups each original with its variants, while + :func:`expansion_weights` exposes graduated weights for the caller's ranker. + Every emitted value is quoted, so input cannot inject MATCH operators. + """ + + parsed = ( + query + if isinstance(query, LexicalQuery) + else build_lexical_query( + query, + synonym_weight=synonym_weight, + subtoken_weight=subtoken_weight, + include_synonyms=include_synonyms, + include_subtokens=include_subtokens, + tuning=tuning, + ) + ) + if not parsed.original_terms: + return "" + + by_origin: dict[str, list[str]] = {term: [] for term in parsed.original_terms} + if include_expansions: + for item in parsed.expanded_terms: + by_origin.setdefault(item.origin, []).append(item.term) + + groups: list[str] = [] + for original in parsed.original_terms: + variants = [original, *by_origin.get(original, [])] + quoted = [escape_fts_term(term) for term in variants] + groups.append(quoted[0] if len(quoted) == 1 else "(" + " OR ".join(quoted) + ")") + return " AND ".join(groups) + + +__all__ = [ + "LEXICAL_STOPWORDS", + "SYNONYM_MAP", + "LexicalQuery", + "WeightedTerm", + "build_fts_query", + "build_lexical_query", + "escape_fts_term", + "expansion_weights", + "salient_terms", + "split_identifier", +] diff --git a/src/codebase_index/retrieval/pipeline.py b/src/codebase_index/retrieval/pipeline.py index d270e30..a6cf2e2 100644 --- a/src/codebase_index/retrieval/pipeline.py +++ b/src/codebase_index/retrieval/pipeline.py @@ -1,7 +1,7 @@ """Orchestrate the hybrid retrieval pipeline (RETRIEVAL.md §1–§7). query -> intent -> retrievers -> RRF fuse -> rerank -> budget -> payload. -Graph expansion (§5) and vector retrieval (§2 vector) are deferred to M5/M6. +Graph expansion is bounded and opt-in; vector retrieval remains optional. """ from __future__ import annotations @@ -15,17 +15,16 @@ from ..indexer.freshness import compute_freshness from . import searchers from .budget import apply_budget +from .diversity import deduplicate, mmr_select from .fusion import fuse from .intent import detect_intent from .rerank import rerank +from .tuning import DEFAULT_TUNING, RetrievalTuning from .types import Confidence +from ..graph.retrieval import graph_candidates _TERM_RE = re.compile(r"[A-Za-z0-9_]+") -_RRF_K = 60 -# Max results kept per file before extras are pushed to the tail. Bucketed fusion -# already collapses co-located hits; this caps the long tail of one big file -# dominating the page so distinct files get surfaced. -_MAX_PER_FILE = 3 +# rrf_k / max_per_file now live on RetrievalTuning so they are ablatable. _KIND_ALIASES = { "method": "method", "methods": "method", @@ -51,28 +50,55 @@ def _requested_symbol_kind(query: str) -> str | None: return next(iter(kinds)) if len(kinds) == 1 else None -def _run_retrievers(conn, query, *, mode, limit, weights, backend=None): +def _run_retrievers( + conn, query, *, mode, limit, weights, backend=None, tuning=DEFAULT_TUNING, + graph_depth: int = 2, graph_node_cap: int = 40, graph_strategy: str = "none", +): lists = {} symbol_kind = _requested_symbol_kind(query) if mode in ("hybrid", "fts"): - lists["fts"] = searchers.fts_candidates(conn, query, limit=limit) + lists["fts"] = searchers.fts_candidates(conn, query, limit=limit, tuning=tuning) if mode in ("hybrid", "symbol"): - lists["symbol"] = searchers.symbol_candidates(conn, query, limit=limit, kind=symbol_kind) + lists["symbol"] = searchers.symbol_candidates( + conn, query, limit=limit, kind=symbol_kind, tuning=tuning + ) if mode == "hybrid": - lists["path"] = searchers.path_candidates(conn, query, limit=limit) + lists["path"] = searchers.path_candidates( + conn, query, limit=limit, tuning=tuning + ) if mode in ("hybrid", "vector") and backend is not None and getattr(backend, "enabled", False): lists["vector"] = searchers.vector_candidates(conn, query, backend, limit=limit) + if mode != "hybrid": weights = {mode: 1.0} + elif tuning.graph_source and graph_strategy != "none": + seeds = [candidate for candidates in lists.values() for candidate in candidates] + related = graph_candidates( + conn, + seeds, + depth=graph_depth, + node_cap=graph_node_cap, + damping=tuning.graph_damping, + iterations=tuning.graph_iterations, + direction={"up": "up", "down": "down", "refs": "up", "both": "both"}.get( + graph_strategy, "both" + ), + ) + if related: + lists["graph"] = related + weights = {**weights, "graph": tuning.graph_weight} return lists, weights - def _confidence(ranked) -> Confidence: if not ranked: return Confidence.LOW top = ranked[0] if top.score <= 0: return Confidence.LOW + exact = getattr(top, "exact_symbol", False) + # Exact symbol matches are high confidence even when they are the sole hit. + if exact: + return Confidence.HIGH if len(ranked) == 1: return Confidence.MEDIUM # Relative gap, not absolute: scale-invariant, so it stays meaningful regardless @@ -80,11 +106,7 @@ def _confidence(ranked) -> Confidence: # surfaced the winning file at all), the signal RRF agreement is meant to capture. rel_gap = (top.score - ranked[1].score) / top.score agree = getattr(top, "agreeing_sources", 1) - exact = getattr(top, "exact_symbol", False) n = len(ranked) - # Exact symbol match always high confidence - if exact: - return Confidence.HIGH # Strong multi-source agreement with a clear score gap if agree >= 3 and rel_gap > 0.15: return Confidence.HIGH @@ -131,21 +153,43 @@ def search( token_budget: int, no_fallback: bool, backend=None, + tuning: Optional[RetrievalTuning] = None, root: Optional[Path] = None, config: Optional[Config] = None, offset: int = 0, compact: bool = True, compact_min_reduction: float = 0.25, ) -> dict: + tuning = tuning or DEFAULT_TUNING plan = detect_intent(query) if token_budget <= 0: token_budget = plan.token_budget fetch_limit = limit + offset + pool_limit = ( + max(fetch_limit * 2, 20) + if (tuning.mmr or tuning.dedup) + else fetch_limit + ) lists, weights = _run_retrievers( - conn, query, mode=mode, limit=fetch_limit, weights=plan.weights, backend=backend + conn, + query, + mode=mode, + limit=pool_limit, + weights=plan.weights, + backend=backend, + tuning=tuning, + graph_depth=tuning.graph_depth, + graph_node_cap=tuning.graph_node_cap, + graph_strategy=plan.graph_strategy, ) - fused = fuse(lists, weights=weights, k=_RRF_K) - ranked = _diversify(rerank(fused, query=query, intent=plan.intent), per_file=_MAX_PER_FILE) + fused = fuse(lists, weights=weights, k=tuning.rrf_k) + ranked = rerank(fused, query=query, intent=plan.intent, tuning=tuning) + if tuning.dedup: + ranked = deduplicate(ranked, hamming_distance=tuning.dedup_hamming) + if tuning.mmr: + ranked = mmr_select(ranked, fetch_limit, tuning.mmr_lambda) + else: + ranked = _diversify(ranked, per_file=tuning.max_per_file) ranked = ranked[:fetch_limit] confidence = _confidence(ranked) # Scale budget proportionally so later pages receive snippet coverage. diff --git a/src/codebase_index/retrieval/priors.py b/src/codebase_index/retrieval/priors.py new file mode 100644 index 0000000..3318e82 --- /dev/null +++ b/src/codebase_index/retrieval/priors.py @@ -0,0 +1,211 @@ +"""Source-role priors for retrieval ranking. + +The helpers in this module inspect only a candidate's path and query metadata. They +never touch the filesystem, and their deliberately small scores are intended to +break near-ties rather than replace lexical or symbol evidence. +""" + +from __future__ import annotations + +import re +from enum import Enum +from pathlib import PurePosixPath + +from ..discovery.classify import detect_language, is_generated, is_test_path +from .types import Intent + + +class SourceRole(str, Enum): + """Coarse role used to apply a ranking prior to a source path.""" + + IMPLEMENTATION = "implementation" + TEST = "test" + DOCUMENTATION = "documentation" + GENERATED_VENDOR_BUILD = "generated_vendor_build" + UNKNOWN = "unknown" + + +# Directory names are compared as complete path components, not substrings. This +# keeps paths such as ``contest`` and ``builder`` from receiving a false penalty. +_GENERATED_DIRS = frozenset( + { + "build", + "coverage", + "dist", + "generated", + "gen", + "out", + "target", + "vendor", + "vendors", + "site-packages", + "dist-packages", + "thirdparty", + "external", + "node_modules", + "bower_components", + "__pycache__", + ".venv", + "venv", + } +) +_DOCUMENTATION_DIRS = frozenset({"doc", "docs", "documentation"}) +_IMPLEMENTATION_DIRS = frozenset( + { + "app", + "bin", + "client", + "cmd", + "include", + "lib", + "module", + "modules", + "pkg", + "server", + "service", + "services", + "src", + } +) +_DOCUMENTATION_SUFFIXES = frozenset({".adoc", ".md", ".mdx", ".rst", ".textile"}) +_DOCUMENTATION_NAMES = frozenset( + { + "authors", + "changelog", + "codeowners", + "contributing", + "copying", + "license", + "notice", + "readme", + "security", + } +) +# Source extensions not covered by discovery.classify.detect_language (for +# extensionless special cases, detect_language remains the authoritative check). +_IMPLEMENTATION_SUFFIXES = frozenset( + { + ".asm", + ".dart", + ".ex", + ".exs", + ".fs", + ".fsx", + ".groovy", + ".hs", + ".jl", + ".m", + ".mm", + ".nim", + ".pl", + ".proto", + ".scala", + ".sh", + ".swift", + ".v", + ".vue", + ".zig", + } +) +_MINIFIED_NAME_RE = re.compile(r"(?:^|[._-])min(?:ified)?(?:[._-]|$)", re.IGNORECASE) +_TEST_QUERY_RE = re.compile( + r"\b(?:test|tests|testing|e2e|spec|specs|fixture|fixtures|mock|mocks|pytest|jest)\b", + re.IGNORECASE, +) + +_ROLE_PRIORS = { + SourceRole.IMPLEMENTATION: 0.08, + SourceRole.TEST: -0.06, + SourceRole.DOCUMENTATION: -0.05, + SourceRole.GENERATED_VENDOR_BUILD: -0.12, + SourceRole.UNKNOWN: 0.0, +} +_TEST_QUERY_PRIOR = 0.05 + + +def _path_parts(path: str) -> tuple[str, ...]: + """Normalize either separator style without touching the filesystem.""" + + normalized = str(path).replace("\\", "/") + return tuple(part for part in PurePosixPath(normalized).parts if part not in {"", "."}) + + +def classify_source_role(path: str) -> SourceRole: + """Classify *path* into a small set of ranking roles. + + Generated/vendor/build markers take precedence over test markers: a vendored + or generated copy is generally less useful even when it happens to live under + a test-looking directory. Test detection delegates to the repository's + word-boundary-aware :func:`is_test_path` helper. + """ + + parts = _path_parts(path) + lowered = tuple(part.lower() for part in parts) + name = lowered[-1] if lowered else "" + stem = name.rsplit(".", 1)[0] if "." in name else name + suffix = PurePosixPath(name).suffix.lower() + + if ( + is_generated(path) + or name.endswith(".map") + or _MINIFIED_NAME_RE.search(name) + or any(part in _GENERATED_DIRS for part in lowered[:-1]) + ): + return SourceRole.GENERATED_VENDOR_BUILD + if is_test_path(path): + return SourceRole.TEST + if ( + any(part in _DOCUMENTATION_DIRS for part in lowered[:-1]) + or suffix in _DOCUMENTATION_SUFFIXES + or stem in _DOCUMENTATION_NAMES + ): + return SourceRole.DOCUMENTATION + if ( + detect_language(path) is not None + or suffix in _IMPLEMENTATION_SUFFIXES + or any(part in _IMPLEMENTATION_DIRS for part in lowered[:-1]) + ): + return SourceRole.IMPLEMENTATION + return SourceRole.UNKNOWN + + +def is_test_intent(query: str = "", intent: Intent | str | None = None) -> bool: + """Return whether ranking should favor test sources for this request. + + ``intent`` accepts the retrieval ``Intent`` enum, its string value, or another + enum-like object. Query terms are still checked because the normal intent plan + intentionally has no dedicated ``TEST`` enum. + """ + + if _TEST_QUERY_RE.search(query): + return True + if intent is None: + return False + value = getattr(intent, "value", intent) + return bool(isinstance(value, str) and _TEST_QUERY_RE.search(value)) + + +def source_role_prior( + path: str, + *, + query: str = "", + intent: Intent | str | None = None, +) -> float: + """Return a deterministic, bounded additive score for a source path. + + The test-role penalty is replaced with a small positive bonus for an explicit + test-oriented query/intent. Other role scores remain unchanged, and generated + or vendored paths stay penalized even when they contain tests. + """ + + role = classify_source_role(path) + if role is SourceRole.TEST and is_test_intent(query, intent): + return _TEST_QUERY_PRIOR + return _ROLE_PRIORS[role] + + +# Concise alias for call sites that refer to the signal as a source prior. +def source_prior(path: str, *, query: str = "", intent: Intent | str | None = None) -> float: + """Alias for :func:`source_role_prior`.""" + + return source_role_prior(path, query=query, intent=intent) diff --git a/src/codebase_index/retrieval/rerank.py b/src/codebase_index/retrieval/rerank.py index 32a6ac0..4b976d8 100644 --- a/src/codebase_index/retrieval/rerank.py +++ b/src/codebase_index/retrieval/rerank.py @@ -11,6 +11,8 @@ import re from ..discovery.classify import is_test_path +from .priors import source_role_prior +from .tuning import DEFAULT_TUNING, RetrievalTuning from .types import Candidate, Intent _TERM_RE = re.compile(r"[A-Za-z0-9_]+") @@ -25,7 +27,13 @@ _DEGREE_CAP = 0.08 -def rerank(candidates: list[Candidate], *, query: str, intent: Intent) -> list[Candidate]: +def rerank( + candidates: list[Candidate], + *, + query: str, + intent: Intent, + tuning: RetrievalTuning = DEFAULT_TUNING, +) -> list[Candidate]: terms = {t.lower() for t in _TERM_RE.findall(query)} for c in candidates: bonus = 0.0 @@ -57,6 +65,12 @@ def rerank(candidates: list[Candidate], *, query: str, intent: Intent) -> list[C if intent is Intent.ARCHITECTURE and (c.in_degree + c.out_degree): bonus += min(_DEGREE_CAP, math.log1p(c.in_degree + c.out_degree) * (_DEGREE_SCALE / 2)) + if tuning.source_priors and ("/" in c.path or "\\" in c.path): + prior = source_role_prior(c.path, query=query, intent=intent) + if prior: + bonus += prior + reasons.append(f"source prior {prior:+.2f}") + wants_tests = "test" in terms or "tests" in terms if c.is_generated or (is_test_path(c.path) and not wants_tests): bonus -= 0.15 diff --git a/src/codebase_index/retrieval/searchers.py b/src/codebase_index/retrieval/searchers.py index d82aca5..a53874a 100644 --- a/src/codebase_index/retrieval/searchers.py +++ b/src/codebase_index/retrieval/searchers.py @@ -1,11 +1,12 @@ """Three retrievers, each emitting a uniform list[Candidate]. -Vector retrieval (RETRIEVAL.md §2) is M6 and intentionally absent here; the -pipeline degrades to path+symbol+fts. +Vector retrieval is optional and supplied by the configured embedding backend; +without it the pipeline degrades cleanly to path, symbol, and FTS retrieval. """ from __future__ import annotations +import math import re import sqlite3 from pathlib import Path @@ -22,27 +23,146 @@ SymbolResponse, ) from ..storage import repo +from .fuzzy import identifier_similarity +from .lexical import ( + LexicalQuery, + build_fts_query, + build_lexical_query, + escape_fts_term, + salient_terms as lexical_salient_terms, +) +from .tuning import DEFAULT_TUNING from .types import Candidate as M4Candidate _WORD_RE = re.compile(r"[A-Za-z0-9_]+") _CAMEL_RE = re.compile(r"[A-Z]+(?![a-z])|[A-Z]?[a-z0-9]+") -def fts_candidates(conn: sqlite3.Connection, query: str, *, limit: int) -> list[M4Candidate]: - match = build_match_query(query) +def _soft_match_query(parsed: LexicalQuery) -> str: + """OR lexical groups so Python can rank by group coverage. + + FTS5 cannot express per-term boosts portably. Each original term is grouped + with its down-weighted variants, then groups are OR-ed. Candidate scoring + below restores coverage and original-term precedence. + """ + by_origin: dict[str, list[str]] = {term: [] for term in parsed.original_terms} + for item in parsed.expanded_terms: + by_origin.setdefault(item.origin, []).append(item.term) + groups: list[str] = [] + for original in parsed.original_terms: + variants = [original, *by_origin.get(original, [])] + quoted = [escape_fts_term(term) for term in variants] + groups.append(quoted[0] if len(quoted) == 1 else "(" + " OR ".join(quoted) + ")") + return " OR ".join(groups) + + +def _fts_coverage(content: str | None, parsed: LexicalQuery) -> tuple[int, float]: + text = (content or "").casefold() + matched = 0 + weighted = 0.0 + weights = parsed.weights + for original in parsed.original_terms: + variants = [ + original, + *(item.term for item in parsed.expanded_terms if item.origin == original), + ] + # Longer code terms are safe and much cheaper as substring probes. + # Two-character terms need token boundaries to avoid `id` matching + # every `grid`/`identifier` occurrence. + found = next( + ( + term + for term in variants + if ( + term.casefold() in text + if len(term) > 2 + else bool( + re.search( + rf"(? list[M4Candidate]: + parsed = build_lexical_query( + query, + include_synonyms=bool(tuning.query_expansion), + include_subtokens=True, + tuning=tuning, + ) + match = build_match_query(query, tuning=tuning) if not match: return [] + # Soft OR queries need a wider first-stage pool; ranking by coverage happens + # locally after FTS. The multiplier is bounded so long natural-language + # questions do not turn into an unbounded scan. + fetch_limit = max(limit, limit * min(5, max(2, len(parsed.original_terms)))) + rows = repo.fts_search(conn, match, limit=fetch_limit) + if not rows: + return [] + + scored_rows: list[tuple[float, int, float, sqlite3.Row]] = [] + for row in rows: + if tuning.soft_lexical and parsed.original_terms: + matched, weighted = _fts_coverage(row["content"], parsed) + required = max(1, math.ceil(len(parsed.original_terms) * tuning.min_term_coverage)) + if matched < required: + continue + coverage = matched / len(parsed.original_terms) + score = coverage * 2.0 + weighted / len(parsed.original_terms) * 0.35 + else: + score = 0.0 + matched = 0 + weighted = 0.0 + # BM25 is an additional tie-break, never the dominant signal in soft mode. + bm25 = max(0.0, -float(row["bm25"])) + score += min(bm25, 4.0) * (0.08 if tuning.soft_lexical else 1.0) + scored_rows.append((score, matched, bm25, row)) + + # If coverage threshold was too strict for an unusual query, retain FTS + # recall rather than returning an empty list. + if tuning.soft_lexical and not scored_rows: + scored_rows = [ + ( + max(0.0, -float(row["bm25"])), + 0, + max(0.0, -float(row["bm25"])), + row, + ) + for row in rows + ] + scored_rows.sort( + key=lambda item: (-item[0], -item[1], -item[2], item[3]["path"], item[3]["line_start"]) + ) out: list[M4Candidate] = [] - for row in repo.fts_search(conn, match, limit=limit): + for score, matched, _, row in scored_rows[:limit]: + reason = ( + f"lexical coverage {matched}/{len(parsed.original_terms)}" + if tuning.soft_lexical and parsed.original_terms + else "fts bm25" + ) out.append( M4Candidate( path=row["path"], line_start=row["line_start"], line_end=row["line_end"], source="fts", - score=-float(row["bm25"]), + score=score, content=row["content"], token_est=int(row["token_est"]), + reason=reason, ) ) return out @@ -51,10 +171,56 @@ def fts_candidates(conn: sqlite3.Connection, query: str, *, limit: int) -> list[ # Natural-language filler that is never a useful symbol query term. Kept deliberately small: # anything that could plausibly be an identifier (get/set/run/...) is NOT a stopword. _SYMBOL_STOPWORDS = { - "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "how", "does", - "do", "did", "what", "where", "which", "who", "whom", "when", "why", "to", "of", "in", - "on", "for", "and", "or", "with", "from", "it", "this", "that", "these", "those", - "into", "during", "if", "via", "across", "between", "about", "their", "its", + "the", + "a", + "an", + "is", + "are", + "was", + "were", + "be", + "been", + "being", + "how", + "does", + "do", + "did", + "what", + "where", + "which", + "who", + "whom", + "when", + "why", + "to", + "of", + "in", + "on", + "for", + "and", + "or", + "with", + "from", + "it", + "this", + "that", + "these", + "those", + "into", + "during", + "if", + "via", + "across", + "between", + "about", + "their", + "its", + "find", + "locate", + "show", + "me", + "implemented", + "implementation", } @@ -75,49 +241,116 @@ def _name_subtokens(name: str) -> set[str]: return {s.lower() for s in _subtokens(name)} -def symbol_candidates( - conn: sqlite3.Connection, query: str, *, limit: int, kind: str | None = None -) -> list[M4Candidate]: - """Symbol retriever that scores by how many query terms a symbol's name covers. +def _fuzzy_symbol_rows( + conn: sqlite3.Connection, query: str, *, limit: int, kind: str | None +) -> list[sqlite3.Row]: + """Fetch a small lexical neighborhood for fuzzy symbol scoring. - The old behaviour searched only the single longest term, so "religion manager" matched - the bare `Religion` class (exact) and never reached `ReligionManager`. Now every salient - term is searched and candidates are ranked by camelCase/underscore-split *coverage* of the - query, so the multi-word concept lands on the multi-word symbol. + The symbol table has a name index but no trigram index. Query terms, their + four-character prefixes, and acronym initials provide a bounded candidate + pool without scanning every symbol; identifier_similarity does the precise + comparison locally. """ - terms = _salient_terms(query) + raw_terms = [t for t in _WORD_RE.findall(query) if len(t) >= 2] + needles: set[str] = set() + for term in raw_terms: + needles.add(term) + if len(term) >= 4: + needles.add(term[:4]) + if len(term) >= 2 and term.isupper(): + needles.add(term[:1]) + rows_by_key: dict[tuple, sqlite3.Row] = {} + per_query = max(limit, min(limit * 2, 40)) + for needle in sorted(needles, key=lambda item: (len(item), item.casefold())): + for row in repo.symbol_search(conn, needle, limit=per_query, kind=kind): + rows_by_key.setdefault((row["path"], row["line_start"], row["name"]), row) + return list(rows_by_key.values()) + + +def symbol_candidates( + conn: sqlite3.Connection, + query: str, + *, + limit: int, + kind: str | None = None, + tuning=DEFAULT_TUNING, +) -> list[M4Candidate]: + """Return symbols ranked by name coverage, exactness, and optional fuzziness.""" + terms = ( + [term for term in lexical_salient_terms(query) if term not in _SYMBOL_STOPWORDS] + if tuning.query_expansion + else _salient_terms(query) + ) + expanded_terms = [] + if tuning.query_expansion: + parsed = build_lexical_query(query, tuning=tuning) + expanded_terms = [ + item.term + for item in parsed.expanded_terms + if item.kind == "synonym" and item.term not in terms + ] + symbol_terms = list(dict.fromkeys((*terms, *expanded_terms))) if not terms: return [] + fuzzy_enabled = tuning.fuzzy_symbols and ( + len(terms) <= 3 or any(char.isupper() for char in query) + ) term_set = set(terms) joined = "".join(terms) rows_by_key: dict[tuple, sqlite3.Row] = {} - for term in terms: - for row in repo.symbol_search(conn, term, limit=limit, kind=kind): - key = (row["path"], row["line_start"], row["name"]) - rows_by_key.setdefault(key, row) + if fuzzy_enabled: + for row in _fuzzy_symbol_rows(conn, query, limit=limit, kind=kind): + rows_by_key.setdefault((row["path"], row["line_start"], row["name"]), row) + for needle in symbol_terms: + for row in repo.symbol_search(conn, needle, limit=limit, kind=kind): + rows_by_key.setdefault((row["path"], row["line_start"], row["name"]), row) scored: list[tuple] = [] for row in rows_by_key.values(): subs = _name_subtokens(row["name"]) - name_l = (row["name"] or "").lower() - covered = sum(1 for t in terms if t in subs or t in name_l) + name_l = (row["name"] or "").casefold() + covered = sum(1 for t in terms if t in subs or t.casefold() in name_l) + expanded_covered = sum(1 for t in expanded_terms if t in subs or t.casefold() in name_l) tightness = len(subs & term_set) / len(subs) if subs else 0.0 - # Exact-match precedence is for *precise* lookups only. With one salient term it's a - # real identifier query; with many it must match the whole camelCase-joined name - # (e.g. "religion manager" -> ReligionManager). A single shared term ("token" hitting - # a generated `Token` type) must NOT count as exact. - exact = (len(terms) == 1 and bool(row["is_exact"])) or (bool(joined) and name_l == joined) - # Ranking: most query terms covered, then exact-name match, then a tighter name - # (fewer junk subtokens), then more-referenced (in_degree), then a shorter name. - sort_key = (covered, int(exact), tightness, int(row["in_degree"]), -len(name_l)) - score = covered + tightness + (2.0 if exact else 0.0) + exact = (len(terms) == 1 and bool(row["is_exact"])) or ( + bool(joined) and name_l == joined.casefold() + ) + fuzzy = 0.0 + if fuzzy_enabled: + fuzzy = max( + identifier_similarity(term, row["name"]) for term in (*terms, joined) if term + ) + if ( + not exact + and covered == 0 + and expanded_covered == 0 + and fuzzy < tuning.fuzzy_threshold + ): + continue + sort_key = ( + int(exact), + covered, + expanded_covered, + tightness, + int(row["in_degree"]), + -len(name_l), + name_l, + row["path"], + int(row["line_start"]), + ) + score = ( + covered + + 0.35 * expanded_covered + + tightness + + (2.0 if exact else 0.0) + + fuzzy * (0.65 if fuzzy_enabled else 0.0) + ) scored.append((sort_key, score, exact, row)) - scored.sort(key=lambda x: x[0], reverse=True) - + scored.sort(key=lambda item: item[0], reverse=True) out: list[M4Candidate] = [] - for sort_key, score, exact, row in scored[:limit]: + for _, score, exact, row in scored[:limit]: out.append( M4Candidate( path=row["path"], @@ -136,8 +369,6 @@ def symbol_candidates( ) ) - # Damped centrality fallback: symbols whose name is not globally unique never - # get a resolved in_degree, so back-fill a name-reference count for the zero ones. zero_deg = [c.symbol for c in out if not c.in_degree and c.symbol] if zero_deg: counts = repo.name_ref_counts(conn, zero_deg) @@ -147,9 +378,29 @@ def symbol_candidates( return out -def path_candidates(conn: sqlite3.Connection, query: str, *, limit: int) -> list[M4Candidate]: +def path_candidates( + conn: sqlite3.Connection, + query: str, + *, + limit: int, + tuning=DEFAULT_TUNING, +) -> list[M4Candidate]: + if tuning.query_expansion: + parsed = build_lexical_query( + query, + include_synonyms=True, + include_subtokens=False, + tuning=tuning, + ) + variants = [ + *parsed.original_terms, + *(item.term for item in parsed.expanded_terms if item.kind == "synonym"), + ] + path_query = " ".join(dict.fromkeys(variants)) + else: + path_query = query out: list[M4Candidate] = [] - for rank, row in enumerate(repo.path_search(conn, query, limit=limit)): + for rank, row in enumerate(repo.path_search(conn, path_query, limit=limit)): out.append( M4Candidate( path=row["path"], @@ -170,31 +421,40 @@ def _subtokens(term: str) -> list[str]: return [p for p in parts if len(p) >= 2] -def build_match_query(query: str) -> str: - """Build the FTS5 MATCH expression for `query`. +def build_match_query(query: str, *, tuning=None) -> str: + """Build a safe FTS5 MATCH expression. - Each whitespace term expands to an OR group over the term and its - camelCase/snake_case subtokens; groups are AND-ed. Natural-language filler - ("how does X work") is dropped first: otherwise FTS would AND-in stopwords - that code chunks never contain, collapsing recall to zero on the very intents - (HOW_IT_WORKS / DEBUG_ERROR) that weight FTS highest. If *every* term is a - stopword we fall back to the full set rather than emit an empty match. + Calls without ``tuning`` preserve the 1.7 public helper exactly (including + its deliberate retention of identifier-like words such as ``work``). + Tuned retrieval uses the lexical parser: soft mode ORs term groups so + coverage can be ranked in Python; hard mode keeps groups AND-ed. """ - groups: list[str] = [] - salient: list[str] = [] - for term in _WORD_RE.findall(query): - variants = {term, *_subtokens(term)} - variants = {v for v in variants if len(v) >= 2} - if not variants: - continue - ored = " OR ".join(f'"{v}"' for v in sorted(variants, key=str.lower)) - # FTS5 rejects implicit AND (space) when a group contains parenthesised OR - # expressions; explicit AND is required between all groups. - group = f"({ored})" if len(variants) > 1 else ored - groups.append(group) - if term.lower() not in _SYMBOL_STOPWORDS: - salient.append(group) - return " AND ".join(salient or groups) + if tuning is None: + groups: list[str] = [] + salient: list[str] = [] + for term in _WORD_RE.findall(query): + variants = {term, *_subtokens(term)} + variants = {v for v in variants if len(v) >= 2} + if not variants: + continue + ored = " OR ".join(f'"{v}"' for v in sorted(variants, key=str.lower)) + group = f"({ored})" if len(variants) > 1 else ored + groups.append(group) + if term.lower() not in _SYMBOL_STOPWORDS: + salient.append(group) + return " AND ".join(salient or groups) + + parsed = build_lexical_query( + query, + include_synonyms=bool(tuning.query_expansion), + include_subtokens=True, + tuning=tuning, + ) + if not parsed.original_terms: + return "" + if tuning.soft_lexical: + return _soft_match_query(parsed) + return build_fts_query(parsed, include_expansions=bool(tuning.query_expansion)) def _freshness( diff --git a/src/codebase_index/retrieval/tuning.py b/src/codebase_index/retrieval/tuning.py new file mode 100644 index 0000000..ba52676 --- /dev/null +++ b/src/codebase_index/retrieval/tuning.py @@ -0,0 +1,106 @@ +"""Retrieval tuning knobs — the ablation contract. + +Every ranking signal added after 1.7.0 sits behind a flag here so it can be +switched off independently and measured (`tests/eval/run_eval.py --ablate`). +A signal that cannot demonstrate a win in the eval harness does not ship. + +`RetrievalTuning.baseline()` reproduces the pre-1.8.0 pipeline byte-for-byte and +is the honest "before" column in every benchmark table. +""" + +from __future__ import annotations + +from dataclasses import dataclass, fields, replace + + +@dataclass(frozen=True) +class RetrievalTuning: + """Immutable retrieval configuration. + + Invariants: + * Every field is either a bool feature flag or a bounded numeric parameter. + * Defaults are the shipped configuration, chosen by benchmark, not taste. + * `baseline()` must keep matching the 1.7.0 behaviour; the ablation table is + meaningless if the baseline drifts. + """ + + # --- candidate generation ------------------------------------------------ + fuzzy_symbols: bool = True + """Acronym / concatenation / edit-distance identifier matching in the symbol + retriever, so `userid` and `get user` both reach `getUserById`.""" + fuzzy_threshold: float = 0.55 + """Minimum identifier similarity for fuzzy symbol candidates.""" + + query_expansion: bool = True + """Down-weighted code-synonym expansion (auth->authentication, ...). Original + terms always keep a strictly higher weight so precision is preserved.""" + + graph_source: bool = False + """Personalized-PageRank candidate source seeded from lexical/symbol hits. + Disabled by default: the self-repository ablation showed lower MRR when + architectural neighbors displaced direct lexical hits. + """ + + soft_lexical: bool = True + """Coverage-scored lexical matching instead of a hard AND over every term. + + Measured cause of the baseline's concept-query collapse: a question like + "how does incremental update avoid reparsing everything" AND-ed five terms, + two of which ("avoid", "everything") appear in no code chunk, so FTS returned + nothing from the target file. Soft matching requires a *fraction* of terms and + ranks by how many matched.""" + + min_term_coverage: float = 0.5 + """Fraction of salient query terms a chunk must contain under soft matching.""" + # --- selection ----------------------------------------------------------- + mmr: bool = False + """Maximal Marginal Relevance re-selection of the ranked list. + Disabled by default: the benchmark's direct file-level metrics favored + relevance-only selection; enable it when snippet diversity is the priority. + """ + + mmr_lambda: float = 0.7 + """MMR relevance/diversity trade-off. 1.0 == pure relevance (no diversity).""" + + dedup: bool = True + """SimHash near-duplicate suppression across returned chunks.""" + + dedup_hamming: int = 3 + """Max Hamming distance between 64-bit SimHashes still considered duplicate.""" + + source_priors: bool = True + """Prefer implementation files over their tests and over prose docs when both + match a code question. Measured: tests/test_fusion.py outranked fusion.py.""" + + # --- fixed parameters ---------------------------------------------------- + rrf_k: int = 60 + max_per_file: int = 3 + graph_damping: float = 0.85 + graph_iterations: int = 12 + graph_weight: float = 0.1 + expansion_weight: float = 0.35 + graph_depth: int = 2 + graph_node_cap: int = 40 + + @classmethod + def baseline(cls) -> RetrievalTuning: + return cls( + fuzzy_symbols=False, + query_expansion=False, + graph_source=False, + soft_lexical=False, + mmr=False, + dedup=False, + source_priors=False, + ) + + def without(self, flag: str) -> RetrievalTuning: + """Return a copy with one boolean signal disabled (single-signal ablation).""" + field_names = {f.name for f in fields(self)} + if flag not in field_names: + raise KeyError(f"unknown tuning flag: {flag!r}") + return replace(self, **{flag: False}) + + +DEFAULT_TUNING = RetrievalTuning() +"""Module-level default so call sites never construct ad-hoc tunings.""" diff --git a/src/codebase_index/storage/repo.py b/src/codebase_index/storage/repo.py index 6bef9eb..3dbd716 100644 --- a/src/codebase_index/storage/repo.py +++ b/src/codebase_index/storage/repo.py @@ -306,7 +306,7 @@ def fts_search( JOIN chunks c ON c.id = fts_chunks.rowid JOIN files f ON f.id = c.file_id WHERE fts_chunks MATCH ? - ORDER BY bm25(fts_chunks) + ORDER BY bm25(fts_chunks), c.id LIMIT ? """, (match_query, limit), @@ -328,7 +328,7 @@ def path_search( ({score_expr}) AS hits FROM files WHERE {' OR '.join(['path LIKE ?'] * len(tokens))} - ORDER BY hits DESC, length(path) ASC + ORDER BY hits DESC, length(path) ASC, path ASC, id ASC LIMIT ? """, (*like_args, *like_args, limit), @@ -363,7 +363,11 @@ def symbol_search( WHERE {name_clause} {kind_clause} ORDER BY is_exact DESC, (s.name LIKE :prefix COLLATE NOCASE) DESC, - s.in_degree DESC + s.in_degree DESC, + s.name COLLATE NOCASE ASC, + f.path ASC, + s.line_start ASC, + s.id ASC LIMIT :limit """, { diff --git a/tests/eval/__init__.py b/tests/eval/__init__.py new file mode 100644 index 0000000..3342c18 --- /dev/null +++ b/tests/eval/__init__.py @@ -0,0 +1 @@ +"""Retrieval evaluation harness (IR metrics, ablation sweeps, latency).""" diff --git a/tests/eval/harness.py b/tests/eval/harness.py new file mode 100644 index 0000000..279a086 --- /dev/null +++ b/tests/eval/harness.py @@ -0,0 +1,257 @@ +"""Retrieval evaluation harness: query set x tuning -> IR metrics + latency. + +Design constraints: + * Reproducible. The index is built once per corpus into a temp dir and reused + across every tuning variant in a sweep, so ablation deltas measure ranking, + not indexing noise. + * Honest. Ground truth is validated against the source tree before any query + runs (`validate_queries`); a stale expectation fails the run loudly instead + of silently deflating scores. + * In-process. `search()` is called directly, so latency excludes interpreter + start-up and measures the thing we can actually optimise. +""" + +from __future__ import annotations + +import statistics +import time +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from pathlib import Path + +import yaml + +from codebase_index.config import Config +from codebase_index.indexer.pipeline import build_index +from codebase_index.retrieval.pipeline import search +from codebase_index.retrieval.tuning import RetrievalTuning +from codebase_index.storage.db import Database + +from . import metrics + +QUERY_DIR = Path(__file__).parent / "queries" +DEFAULT_BUDGET = 1500 +DEFAULT_LIMIT = 10 + + +@dataclass(frozen=True) +class EvalQuery: + query: str + category: str + expected_files: tuple[str, ...] + expected_symbols: tuple[str, ...] = () + + +@dataclass +class QueryOutcome: + query: EvalQuery + ranked_files: list[str] + returned: list[tuple[str, int]] + latency_ms: float + + +@dataclass +class EvalReport: + """Aggregate metrics for one (query set, tuning) pair.""" + + label: str + n_queries: int + recall_at_5: float + recall_at_10: float + mrr: float + ndcg_at_10: float + hit_rate_at_3: float + precision_at_5: float + map_score: float + useful_context: float + p50_ms: float + p95_ms: float + p99_ms: float + mean_ms: float + per_category: dict[str, float] = field(default_factory=dict) + + def as_row(self) -> dict[str, float | str | int]: + return { + "label": self.label, + "n": self.n_queries, + "recall@5": self.recall_at_5, + "recall@10": self.recall_at_10, + "MRR": self.mrr, + "nDCG@10": self.ndcg_at_10, + "hit@3": self.hit_rate_at_3, + "P@5": self.precision_at_5, + "MAP": self.map_score, + "useful@budget": self.useful_context, + "p50_ms": self.p50_ms, + "p95_ms": self.p95_ms, + "p99_ms": self.p99_ms, + } + + +def load_queries(name_or_path: str | Path) -> list[EvalQuery]: + """Load a query set by bare name (`self_repo`) or explicit path.""" + path = Path(name_or_path) + if not path.is_file(): + path = QUERY_DIR / f"{name_or_path}.yml" + raw = yaml.safe_load(path.read_text(encoding="utf-8")) or [] + out: list[EvalQuery] = [] + for entry in raw: + out.append( + EvalQuery( + query=entry["query"], + category=entry.get("category", "uncategorized"), + expected_files=tuple(entry.get("expected_files", ())), + expected_symbols=tuple(entry.get("expected_symbols", ())), + ) + ) + return out + + +def validate_queries(queries: Iterable[EvalQuery], root: Path) -> list[str]: + """Return a list of ground-truth problems (missing files). Empty == valid.""" + problems: list[str] = [] + for q in queries: + if not q.expected_files: + problems.append(f"{q.query!r}: no expected_files (unscoreable)") + for rel in q.expected_files: + if not (root / rel).is_file(): + problems.append(f"{q.query!r}: expected file missing from tree: {rel}") + return problems + + +# The ground-truth YAML quotes every benchmark query verbatim, so leaving it in +# the corpus makes it the top lexical hit for almost every query — a measurement +# artifact that depresses scores and hides real ranking behaviour. Benchmark +# scaffolding is excluded from the corpus it grades. +CORPUS_EXCLUDES = ("tests/eval/**", "tests/benchmark_*", "tests/fixtures/expected_answers.yml") + + +def build_corpus_index(root: Path, db_path: Path) -> Database: + """Build a fresh index for `root` at `db_path` and return the open handle.""" + cfg = Config() + cfg.root = str(root) + cfg.embeddings.enabled = False + cfg.extra_ignore = [*cfg.extra_ignore, *CORPUS_EXCLUDES] + db = Database(db_path).open() + build_index(cfg, db, root=root) + return db + + +def _normalise(path: str) -> str: + return path.replace("\\", "/") + + +def run_query( + conn, + q: EvalQuery, + *, + tuning: RetrievalTuning, + limit: int, + token_budget: int, + mode: str = "hybrid", +) -> QueryOutcome: + start = time.perf_counter() + payload = search( + conn, + q.query, + mode=mode, + limit=limit, + token_budget=token_budget, + no_fallback=True, + tuning=tuning, + ) + latency_ms = (time.perf_counter() - start) * 1000.0 + + ranked_files: list[str] = [] + returned: list[tuple[str, int]] = [] + for r in payload.get("results", []): + p = _normalise(r["path"]) + # File-level ranking: the agent's unit of decision is "which file do I + # open", so multiple hits inside one file collapse to its best rank. + if p not in ranked_files: + ranked_files.append(p) + returned.append((p, int(r.get("token_est") or 0))) + return QueryOutcome(query=q, ranked_files=ranked_files, returned=returned, + latency_ms=latency_ms) + + +def evaluate( + conn, + queries: Sequence[EvalQuery], + *, + tuning: RetrievalTuning, + label: str, + limit: int = DEFAULT_LIMIT, + token_budget: int = DEFAULT_BUDGET, + repeats: int = 1, +) -> EvalReport: + """Run every query and aggregate metrics. `repeats` only affects latency.""" + outcomes: list[QueryOutcome] = [] + latencies: list[float] = [] + for q in queries: + outcome = run_query(conn, q, tuning=tuning, limit=limit, token_budget=token_budget) + outcomes.append(outcome) + latencies.append(outcome.latency_ms) + for _ in range(max(0, repeats - 1)): + latencies.append( + run_query(conn, q, tuning=tuning, limit=limit, + token_budget=token_budget).latency_ms + ) + + def mean(fn) -> float: + vals = [fn(o) for o in outcomes] + return statistics.fmean(vals) if vals else 0.0 + + per_category: dict[str, list[float]] = {} + for o in outcomes: + per_category.setdefault(o.query.category, []).append( + metrics.reciprocal_rank(o.ranked_files, o.query.expected_files) + ) + + return EvalReport( + label=label, + n_queries=len(outcomes), + recall_at_5=mean(lambda o: metrics.recall_at_k(o.ranked_files, o.query.expected_files, 5)), + recall_at_10=mean(lambda o: metrics.recall_at_k(o.ranked_files, o.query.expected_files, 10)), + mrr=mean(lambda o: metrics.reciprocal_rank(o.ranked_files, o.query.expected_files)), + ndcg_at_10=mean(lambda o: metrics.ndcg_at_k(o.ranked_files, o.query.expected_files, 10)), + hit_rate_at_3=mean(lambda o: metrics.hit_rate_at_k(o.ranked_files, o.query.expected_files, 3)), + precision_at_5=mean(lambda o: metrics.precision_at_k(o.ranked_files, o.query.expected_files, 5)), + map_score=mean(lambda o: metrics.average_precision(o.ranked_files, o.query.expected_files)), + useful_context=mean( + lambda o: metrics.useful_context_at_budget( + o.returned, o.query.expected_files, token_budget + ) + ), + p50_ms=metrics.percentile(latencies, 50), + p95_ms=metrics.percentile(latencies, 95), + p99_ms=metrics.percentile(latencies, 99), + mean_ms=statistics.fmean(latencies) if latencies else 0.0, + per_category={ + cat: statistics.fmean(vals) for cat, vals in sorted(per_category.items()) + }, + ) + + +def format_table(reports: Sequence[EvalReport], *, baseline: EvalReport | None = None) -> str: + """Render reports as a Markdown table, with deltas against `baseline`.""" + cols = ["label", "recall@5", "recall@10", "MRR", "nDCG@10", "hit@3", "P@5", + "MAP", "useful@budget", "p50_ms", "p95_ms"] + lines = ["| " + " | ".join(cols) + " |", + "|" + "|".join("---" for _ in cols) + "|"] + for rep in reports: + row = rep.as_row() + cells = [] + for c in cols: + v = row[c] + if isinstance(v, float): + cell = f"{v:.3f}" if c.endswith("_ms") is False else f"{v:.1f}" + if baseline is not None and rep is not baseline: + delta = v - float(baseline.as_row()[c]) + if abs(delta) >= 0.0005: + cell += f" ({delta:+.3f})" if not c.endswith("_ms") else f" ({delta:+.1f})" + cells.append(cell) + else: + cells.append(str(v)) + lines.append("| " + " | ".join(cells) + " |") + return "\n".join(lines) diff --git a/tests/eval/metrics.py b/tests/eval/metrics.py new file mode 100644 index 0000000..5ab1cef --- /dev/null +++ b/tests/eval/metrics.py @@ -0,0 +1,129 @@ +"""Standard IR metrics for the retrieval eval harness. + +All functions take a *ranked* list of retrieved item ids (best first, already +deduplicated) and a set of relevant ids. Binary relevance throughout: a file +either does or does not contain the answer. Graded relevance was considered and +rejected — assigning 0..3 grades by hand to hundreds of (query, file) pairs is +not reproducible, and binary gains keep nDCG comparable across corpora. + +Every function is pure and total: an empty ranking or empty ground truth yields +0.0 rather than raising, so a query that returns nothing scores as a miss +instead of crashing the sweep. +""" + +from __future__ import annotations + +import math +from collections.abc import Iterable, Sequence + + +def _prefix(ranked: Sequence[str], k: int) -> Sequence[str]: + return ranked[: max(0, k)] + + +def recall_at_k(ranked: Sequence[str], relevant: Iterable[str], k: int) -> float: + """Fraction of the relevant set retrieved within the top k.""" + rel = set(relevant) + if not rel: + return 0.0 + return len(set(_prefix(ranked, k)) & rel) / len(rel) + + +def precision_at_k(ranked: Sequence[str], relevant: Iterable[str], k: int) -> float: + """Fraction of the top k that is relevant.""" + rel = set(relevant) + top = _prefix(ranked, k) + if not top: + return 0.0 + return sum(1 for item in top if item in rel) / len(top) + + +def hit_rate_at_k(ranked: Sequence[str], relevant: Iterable[str], k: int) -> float: + """1.0 if any relevant item appears in the top k. The 'did it work at all' metric.""" + rel = set(relevant) + return 1.0 if any(item in rel for item in _prefix(ranked, k)) else 0.0 + + +def reciprocal_rank(ranked: Sequence[str], relevant: Iterable[str]) -> float: + """1/rank of the first relevant hit; 0.0 when the ranking misses entirely. + + Averaged over a query set this is MRR — the metric that matters most for an + agent, which reads results top-down and stops early. + """ + rel = set(relevant) + for i, item in enumerate(ranked, start=1): + if item in rel: + return 1.0 / i + return 0.0 + + +def average_precision(ranked: Sequence[str], relevant: Iterable[str]) -> float: + """Mean of precision@i taken at every rank i holding a relevant item.""" + rel = set(relevant) + if not rel: + return 0.0 + hits = 0 + total = 0.0 + for i, item in enumerate(ranked, start=1): + if item in rel: + hits += 1 + total += hits / i + return total / len(rel) if hits else 0.0 + + +def ndcg_at_k(ranked: Sequence[str], relevant: Iterable[str], k: int) -> float: + """Binary-gain nDCG@k with the standard log2(rank+1) discount. + + The ideal ranking places min(|relevant|, k) hits at the very top, so a query + with a single relevant file can still reach 1.0 — nDCG is not implicitly + penalised for small ground-truth sets. + """ + rel = set(relevant) + if not rel: + return 0.0 + dcg = sum( + 1.0 / math.log2(i + 1) + for i, item in enumerate(_prefix(ranked, k), start=1) + if item in rel + ) + ideal_hits = min(len(rel), max(0, k)) + idcg = sum(1.0 / math.log2(i + 1) for i in range(1, ideal_hits + 1)) + return dcg / idcg if idcg else 0.0 + + +def useful_context_at_budget( + returned: Sequence[tuple[str, int]], + relevant: Iterable[str], + budget_tokens: int, +) -> float: + """Agent-centric metric: relevant items recovered per unit of context spent. + + `returned` is the ranked list of (item_id, token_cost) actually placed in the + agent's context. We walk it in rank order, spending tokens until the budget is + exhausted, and report the fraction of the relevant set that made it in. + + This is what the whole system optimises: not "is the answer somewhere in the + ranking" but "is the answer in the context window the agent can afford". + """ + rel = set(relevant) + if not rel: + return 0.0 + spent = 0 + found: set[str] = set() + for item, cost in returned: + cost = max(0, cost) + if spent + cost > budget_tokens: + continue + spent += cost + if item in rel: + found.add(item) + return len(found) / len(rel) + + +def percentile(values: Sequence[float], pct: float) -> float: + """Nearest-rank percentile (no interpolation) — stable for small samples.""" + if not values: + return 0.0 + ordered = sorted(values) + idx = math.ceil(pct / 100.0 * len(ordered)) - 1 + return ordered[min(max(idx, 0), len(ordered) - 1)] diff --git a/tests/eval/queries/self_repo.yml b/tests/eval/queries/self_repo.yml new file mode 100644 index 0000000..9abdba7 --- /dev/null +++ b/tests/eval/queries/self_repo.yml @@ -0,0 +1,204 @@ +# Ground-truth query set over this repository (self-hosted eval corpus). +# +# Why self-hosted: it is a real ~4k-LOC multi-package Python codebase that ships +# with the project, so the benchmark is reproducible by anyone who clones it — no +# external checkout, no network. A second, larger, Java corpus is evaluated via +# `--repo` (see tests/eval/README.md) to guard against overfitting to Python. +# +# Ground truth is OBJECTIVE: `expected_files` are the files that actually define +# the answer, verified against the source tree by `harness.validate_queries()`, +# which fails the run if a path or symbol no longer exists. The index cannot +# grade its own homework — these are written from the source, not from output. +# +# `category` groups queries for per-category reporting; it is not fed to the +# retriever. Categories follow the agent question taxonomy: navigation, +# exact_symbol, concept, bug, architecture, impact, test. + +- query: "where is query intent classified" + category: navigation + expected_files: [src/codebase_index/retrieval/intent.py] + expected_symbols: [detect_intent] + +- query: "how does reciprocal rank fusion combine retrievers" + category: concept + expected_files: [src/codebase_index/retrieval/fusion.py] + expected_symbols: [fuse] + +- query: "where is the token budget applied to results" + category: navigation + expected_files: [src/codebase_index/retrieval/budget.py] + expected_symbols: [apply_budget] + +- query: "find walk_impact" + category: exact_symbol + expected_files: [src/codebase_index/graph/expand.py] + expected_symbols: [walk_impact] + +- query: "build_match_query" + category: exact_symbol + expected_files: [src/codebase_index/retrieval/searchers.py] + expected_symbols: [build_match_query] + +- query: "detect_intent" + category: exact_symbol + expected_files: [src/codebase_index/retrieval/intent.py] + expected_symbols: [detect_intent] + +- query: "compute_freshness" + category: exact_symbol + expected_files: [src/codebase_index/indexer/freshness.py] + expected_symbols: [compute_freshness] + +- query: "resolve_backend" + category: exact_symbol + expected_files: [src/codebase_index/embeddings/backend.py] + expected_symbols: [resolve_backend] + +- query: "how does incremental update avoid reparsing everything" + category: concept + expected_files: [src/codebase_index/indexer/pipeline.py] + expected_symbols: [update_index] + +- query: "how are files chosen for indexing" + category: concept + expected_files: + - src/codebase_index/discovery/walker.py + - src/codebase_index/discovery/ignore.py + - src/codebase_index/discovery/classify.py + +- query: "where do we split camelCase identifiers" + category: concept + expected_files: + - src/codebase_index/retrieval/searchers.py + - src/codebase_index/retrieval/lexical.py + expected_symbols: [split_identifier] +- query: "tree-sitter parsing of source files into symbols" + category: navigation + expected_files: + - src/codebase_index/parsers/treesitter.py + - src/codebase_index/parsers/languages.py + expected_symbols: [parse_file] + +- query: "sqlite pragmas and schema version guard" + category: navigation + expected_files: [src/codebase_index/storage/db.py] + expected_symbols: [Database] + +- query: "how does the MCP stdio server expose tools" + category: architecture + expected_files: [src/codebase_index/mcp/server.py] + +- query: "where are secrets redacted before output" + category: navigation + expected_files: [src/codebase_index/output/redact.py] + expected_symbols: [redact_snippet] + +- query: "how is the dependency graph built from AST edges" + category: concept + expected_files: [src/codebase_index/graph/builder.py] + expected_symbols: [build_graph] + +- query: "detecting god nodes and modules in the codebase map" + category: architecture + expected_files: [src/codebase_index/graph/analysis.py] + expected_symbols: [god_nodes] + +- query: "shortest dependency path between two symbols" + category: navigation + expected_files: [src/codebase_index/graph/navigate.py] + expected_symbols: [path_payload] + +- query: "what breaks if I change the Candidate dataclass" + category: impact + expected_files: + - src/codebase_index/retrieval/types.py + - src/codebase_index/retrieval/searchers.py + - src/codebase_index/retrieval/fusion.py + - src/codebase_index/retrieval/rerank.py + - src/codebase_index/retrieval/budget.py + +- query: "what depends on the Database class" + category: impact + expected_files: + - src/codebase_index/storage/db.py + - src/codebase_index/cli.py + - src/codebase_index/service.py + +- query: "how does configuration get loaded and merged" + category: concept + expected_files: [src/codebase_index/config.py] + expected_symbols: [find_root, config_hash] + +- query: "where is the project root discovered" + category: navigation + expected_files: [src/codebase_index/config.py] + expected_symbols: [find_root] + +- query: "which tests cover the retrieval budget" + category: test + expected_files: [tests/test_budget.py] + +- query: "tests for tree-sitter symbol extraction" + category: test + expected_files: [tests/test_treesitter.py] + +- query: "unsupported language raises during parsing" + category: bug + expected_files: [src/codebase_index/parsers/treesitter.py] + +- query: "why would search return zero results for a stopword-only query" + category: bug + expected_files: [src/codebase_index/retrieval/searchers.py] + expected_symbols: [build_match_query] + +- query: "index reports stale after files changed" + category: bug + expected_files: [src/codebase_index/indexer/freshness.py] + expected_symbols: [compute_freshness] + +- query: "how does the CLI wire subcommands to the service layer" + category: architecture + expected_files: + - src/codebase_index/cli.py + - src/codebase_index/service.py + +- query: "config and security diagnostics" + category: navigation + expected_files: [src/codebase_index/doctor.py] + expected_symbols: [run_doctor] + +- query: "how are chunks produced when no grammar exists" + category: concept + expected_files: [src/codebase_index/parsers/line_chunker.py] + expected_symbols: [chunk_text] + +- query: "skeletonize a snippet to save tokens" + category: concept + expected_files: [src/codebase_index/retrieval/skeleton.py] + expected_symbols: [make_compactor] + +- query: "semantic vector search over chunks" + category: navigation + expected_files: + - src/codebase_index/retrieval/searchers.py + - src/codebase_index/embeddings/backend.py + expected_symbols: [vector_candidates] + +- query: "scaffolding the skill directory on init" + category: navigation + expected_files: [src/codebase_index/scaffold.py] + +- query: "watching the filesystem for changes" + category: navigation + expected_files: [src/codebase_index/watch/watcher.py] + +- query: "how does ranking demote generated and test files" + category: concept + expected_files: + - src/codebase_index/retrieval/rerank.py + - src/codebase_index/discovery/classify.py + expected_symbols: [rerank, is_test_path] + +- query: "exporting the graph to HTML for a human" + category: navigation + expected_files: [src/codebase_index/graph/export.py] diff --git a/tests/eval/run_eval.py b/tests/eval/run_eval.py new file mode 100644 index 0000000..d2a5da8 --- /dev/null +++ b/tests/eval/run_eval.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +"""Run the retrieval benchmark and ablation sweep. + + python tests/eval/run_eval.py # baseline vs shipped default + python tests/eval/run_eval.py --ablate # + one-signal-off sweep + python tests/eval/run_eval.py --repo --queries + +The index is built once and shared by every variant, so reported deltas isolate +ranking changes from indexing variance. +""" + +from __future__ import annotations + +import argparse +import json +import sys +import tempfile +from dataclasses import fields +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +sys.path.insert(0, str(REPO_ROOT / "tests")) +sys.path.insert(0, str(REPO_ROOT / "src")) + +from eval import harness # noqa: E402 +from codebase_index.retrieval.tuning import RetrievalTuning # noqa: E402 + +# Signals swept by --ablate. Numeric parameters are excluded: turning off a bool +# answers "does this signal earn its complexity", which is the decision we make. +ABLATABLE = ("soft_lexical", "query_expansion", "fuzzy_symbols", "graph_source", + "mmr", "dedup", "source_priors") + + +def main(argv: list[str] | None = None) -> int: + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--repo", default=str(REPO_ROOT), + help="corpus to index and query (default: this repository)") + ap.add_argument("--queries", default="self_repo", + help="query set name under tests/eval/queries or a path") + ap.add_argument("--ablate", action="store_true", + help="also run a one-signal-off sweep") + ap.add_argument("--limit", type=int, default=harness.DEFAULT_LIMIT) + ap.add_argument("--token-budget", type=int, default=harness.DEFAULT_BUDGET) + ap.add_argument("--repeats", type=int, default=3, + help="latency repeats per query (quality is unaffected)") + ap.add_argument("--json", dest="as_json", action="store_true") + args = ap.parse_args(argv) + + root = Path(args.repo).resolve() + queries = harness.load_queries(args.queries) + problems = harness.validate_queries(queries, root) + if problems: + print("Ground-truth validation FAILED:", file=sys.stderr) + for p in problems: + print(f" - {p}", file=sys.stderr) + return 2 + + print(f"corpus: {root}") + print(f"queries: {len(queries)} from {args.queries}") + print("building index (once, shared by all variants)...", flush=True) + + with tempfile.TemporaryDirectory() as tmp: + db = harness.build_corpus_index(root, Path(tmp) / "eval.sqlite") + try: + variants: list[tuple[str, RetrievalTuning]] = [ + ("baseline (1.7.0)", RetrievalTuning.baseline()), + ("default (all signals)", RetrievalTuning()), + ] + if args.ablate: + default = RetrievalTuning() + known = {f.name for f in fields(default)} + for flag in ABLATABLE: + if flag in known: + variants.append((f" -{flag}", default.without(flag))) + + reports = [] + for label, tuning in variants: + print(f"running: {label}", flush=True) + reports.append( + harness.evaluate( + db.conn, queries, tuning=tuning, label=label, + limit=args.limit, token_budget=args.token_budget, + repeats=args.repeats, + ) + ) + finally: + db.close() + + if args.as_json: + print(json.dumps([r.as_row() for r in reports], indent=2)) + return 0 + + print() + print(harness.format_table(reports, baseline=reports[0])) + print() + print("Per-category MRR (default config):") + for cat, val in reports[1].per_category.items(): + print(f" {cat:14} {val:.3f}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/golden/explain_auth.json b/tests/golden/explain_auth.json index 80c361f..e680331 100644 --- a/tests/golden/explain_auth.json +++ b/tests/golden/explain_auth.json @@ -1,11 +1,6 @@ { - "confidence": "low", - "fallback_suggestions": { - "ripgrep": [ - "rg -n \"authentication\"", - "rg -n \"how.*does.*authentication\"" - ] - }, + "confidence": "high", + "fallback_suggestions": {}, "index": { "built_at": "", "exists": true, @@ -16,6 +11,44 @@ "intent": "how_it_works", "mode": "hybrid", "query": "how does authentication work", - "recommended_reads": [], - "results": [] + "recommended_reads": [ + { + "line_end": 5, + "line_start": 1, + "path": "src/api/service.py" + }, + { + "line_end": 1, + "line_start": 1, + "path": "src/auth/token.py" + } + ], + "results": [ + { + "elided_lines": 0, + "line_end": 5, + "line_start": 1, + "path": "src/api/service.py", + "rank": 1, + "reason": "source prior +0.08", + "score": 1.08, + "skeletonized": false, + "snippet": "\"\"\"Service layer (fixture) - exercises cross-file edges for impact tests.\"\"\"\n\nfrom auth.token import refresh_access_token\nfrom models.user import User\n", + "symbols": [], + "token_est": 38 + }, + { + "elided_lines": 0, + "line_end": 1, + "line_start": 1, + "path": "src/auth/token.py", + "rank": 2, + "reason": "source prior +0.08", + "score": 0.38, + "skeletonized": false, + "snippet": null, + "symbols": [], + "token_est": 0 + } + ] } diff --git a/tests/golden/impact_user_model.json b/tests/golden/impact_user_model.json index cb8d033..a5107f6 100644 --- a/tests/golden/impact_user_model.json +++ b/tests/golden/impact_user_model.json @@ -17,15 +17,6 @@ "stale": false }, "nodes": [ - { - "distance": 1, - "kind": "file", - "line_start": null, - "name": null, - "path": "src/api/service.py", - "via_confidence": "inferred", - "via_edge": "import" - }, { "distance": 1, "kind": "symbol", @@ -34,6 +25,15 @@ "path": "src/api/service.py", "via_confidence": "extracted", "via_edge": "extends" + }, + { + "distance": 1, + "kind": "file", + "line_start": null, + "name": null, + "path": "src/api/service.py", + "via_confidence": "inferred", + "via_edge": "import" } ], "target": "src/models/user.py" diff --git a/tests/golden/mcp_explain_code.json b/tests/golden/mcp_explain_code.json index 74eebe9..46f9bbf 100644 --- a/tests/golden/mcp_explain_code.json +++ b/tests/golden/mcp_explain_code.json @@ -1,11 +1,6 @@ { - "confidence": "low", - "fallback_suggestions": { - "ripgrep": [ - "rg -n \"authentication\"", - "rg -n \"how.*does.*authentication\"" - ] - }, + "confidence": "high", + "fallback_suggestions": {}, "index": { "built_at": "", "exists": true, @@ -16,8 +11,46 @@ "intent": "how_it_works", "mode": "hybrid", "query": "how does authentication work", - "recommended_reads": [], - "results": [], + "recommended_reads": [ + { + "line_end": 5, + "line_start": 1, + "path": "src/api/service.py" + }, + { + "line_end": 1, + "line_start": 1, + "path": "src/auth/token.py" + } + ], + "results": [ + { + "elided_lines": 0, + "line_end": 5, + "line_start": 1, + "path": "src/api/service.py", + "rank": 1, + "reason": "source prior +0.08", + "score": 1.08, + "skeletonized": false, + "snippet": "\"\"\"Service layer (fixture) - exercises cross-file edges for impact tests.\"\"\"\n\nfrom auth.token import refresh_access_token\nfrom models.user import User\n", + "symbols": [], + "token_est": 38 + }, + { + "elided_lines": 0, + "line_end": 1, + "line_start": 1, + "path": "src/auth/token.py", + "rank": 2, + "reason": "source prior +0.08", + "score": 0.38, + "skeletonized": false, + "snippet": null, + "symbols": [], + "token_est": 0 + } + ], "schema_version": 1, "tool": "explain_code" } diff --git a/tests/golden/mcp_impact_of.json b/tests/golden/mcp_impact_of.json index 1629921..1ce4815 100644 --- a/tests/golden/mcp_impact_of.json +++ b/tests/golden/mcp_impact_of.json @@ -17,15 +17,6 @@ "stale": false }, "nodes": [ - { - "distance": 1, - "kind": "file", - "line_start": null, - "name": null, - "path": "src/api/service.py", - "via_confidence": "inferred", - "via_edge": "import" - }, { "distance": 1, "kind": "symbol", @@ -34,6 +25,15 @@ "path": "src/api/service.py", "via_confidence": "extracted", "via_edge": "extends" + }, + { + "distance": 1, + "kind": "file", + "line_start": null, + "name": null, + "path": "src/api/service.py", + "via_confidence": "inferred", + "via_edge": "import" } ], "schema_version": 1, diff --git a/tests/golden/mcp_search_code.json b/tests/golden/mcp_search_code.json index b2e1a8a..247be25 100644 --- a/tests/golden/mcp_search_code.json +++ b/tests/golden/mcp_search_code.json @@ -25,8 +25,8 @@ "line_start": 4, "path": "src/auth/token.py", "rank": 1, - "reason": "in src/auth/ · 2 callers", - "score": 2.233, + "reason": "in src/auth/ · 2 callers · source prior +0.08", + "score": 2.313, "skeletonized": false, "snippet": "def refresh_access_token(refresh_token: str) -> str:", "symbols": [ @@ -40,8 +40,8 @@ "line_start": 7, "path": "src/api/service.py", "rank": 2, - "reason": "fts", - "score": 0.9375, + "reason": "source prior +0.08", + "score": 1.0175, "skeletonized": false, "snippet": "class AdminUser(User):\n \"\"\"Subclass of User; imported-from edge target for impact tests.\"\"\"\n\n def renew(self, refresh_token: str) -> str:\n return refresh_access_token(refresh_token)", "symbols": [], diff --git a/tests/golden/search_token.json b/tests/golden/search_token.json index 295730f..dd69cdd 100644 --- a/tests/golden/search_token.json +++ b/tests/golden/search_token.json @@ -25,8 +25,8 @@ "line_start": 4, "path": "src/auth/token.py", "rank": 1, - "reason": "in src/auth/ · 2 callers", - "score": 2.233, + "reason": "in src/auth/ · 2 callers · source prior +0.08", + "score": 2.313, "skeletonized": false, "snippet": "def refresh_access_token(refresh_token: str) -> str:", "symbols": [ @@ -40,8 +40,8 @@ "line_start": 7, "path": "src/api/service.py", "rank": 2, - "reason": "fts", - "score": 0.9375, + "reason": "source prior +0.08", + "score": 1.0175, "skeletonized": false, "snippet": "class AdminUser(User):\n \"\"\"Subclass of User; imported-from edge target for impact tests.\"\"\"\n\n def renew(self, refresh_token: str) -> str:\n return refresh_access_token(refresh_token)", "symbols": [], diff --git a/tests/test_diversity.py b/tests/test_diversity.py new file mode 100644 index 0000000..6956d87 --- /dev/null +++ b/tests/test_diversity.py @@ -0,0 +1,49 @@ +from codebase_index.retrieval.diversity import deduplicate, mmr_select +from codebase_index.retrieval.types import Candidate + + +def _candidate(path: str, score: float, content: str, *, symbol: str | None = None) -> Candidate: + return Candidate( + path=path, + line_start=1, + line_end=4, + source="fts", + score=score, + content=content, + symbol=symbol, + ) + + +def test_deduplicate_collapses_formatting_and_comments_and_keeps_best(): + weaker = _candidate("weak.py", 0.4, "def add(a, b):\n return a + b") + stronger = _candidate( + "strong.py", 0.9, "# comment\ndef add(a,b):\n\treturn a + b # trailing comment" + ) + + result = deduplicate([weaker, stronger]) + + assert result == [stronger] + assert weaker.score == 0.4 # caller-owned candidates are not mutated + + +def test_deduplicate_keeps_structurally_distinct_snippets(): + left = _candidate("left.py", 0.8, "def add(a, b): return a + b") + right = _candidate("right.py", 0.7, "def multiply(a, b): return a * b") + + assert deduplicate([left, right]) == [left, right] + + +def test_mmr_balances_relevance_against_redundancy(): + first = _candidate("first.py", 1.0, "def add(a, b): return a + b", symbol="add") + near = _candidate("near.py", 0.95, "def add(a,b): return a + b", symbol="add") + distinct = _candidate("distinct.py", 0.55, "def parse(value): return decode(value)", symbol="parse") + + assert mmr_select([first, near, distinct], 2, lambda_=1.0) == [first, near] + assert mmr_select([first, near, distinct], 2, lambda_=0.0) == [first, distinct] + + +def test_selection_handles_empty_and_singleton_inputs_without_mutation(): + candidate = _candidate("one.py", 1.0, "return 1") + assert deduplicate([]) == [] + assert mmr_select([], 3, lambda_=0.5) == [] + assert mmr_select([candidate], 3, lambda_=0.5) == [candidate] diff --git a/tests/test_fuzzy.py b/tests/test_fuzzy.py new file mode 100644 index 0000000..9539f40 --- /dev/null +++ b/tests/test_fuzzy.py @@ -0,0 +1,42 @@ +from codebase_index.retrieval.fuzzy import identifier_similarity, rank_fuzzy_symbols +from codebase_index.retrieval.types import Candidate + + +def test_identifier_similarity_handles_case_tokens_and_concatenation(): + assert identifier_similarity("Get_User", "getUser") == 1.0 + assert identifier_similarity("userid", "getUserById") > 0.8 + assert identifier_similarity("getUser", "getUserById") > 0.8 + + +def test_identifier_similarity_handles_acronyms_and_edit_distance(): + assert identifier_similarity("GUBI", "getUserById") > 0.85 + assert identifier_similarity("getUsrById", "getUserById") > 0.65 + assert identifier_similarity("GUBI", "getUserByName") < 0.55 + + +def test_identifier_similarity_short_and_empty_inputs_are_safe(): + assert identifier_similarity("", "anything") == 0.0 + assert identifier_similarity("a", "anything") == 0.0 + assert identifier_similarity("", "") == 0.0 + assert identifier_similarity("ID", "id") == 1.0 + + +def test_rank_fuzzy_symbols_preserves_exactness_and_has_deterministic_ties(): + candidates = [ + Candidate("z.py", 1, 2, "symbol", 0.1, symbol="getUserById"), + Candidate("a.py", 4, 5, "symbol", 0.1, symbol="getUserById"), + Candidate("exact.py", 1, 2, "symbol", 0.0, symbol="userid", exact_symbol=True), + ] + ranked = rank_fuzzy_symbols("userid", candidates, threshold=0.55) + assert ranked[0] is candidates[2] + assert [c.path for c in ranked[1:]] == ["a.py", "z.py"] + assert all(c.exact_symbol is (c is candidates[2]) for c in ranked) + + +def test_rank_fuzzy_symbols_accepts_mapping_rows_and_threshold(): + rows = [ + {"name": "getUserById", "path": "src/user.py", "line_start": 1}, + {"name": "getUserByName", "path": "src/name.py", "line_start": 1}, + ] + ranked = rank_fuzzy_symbols("GUBI", rows, threshold=0.55) + assert [row["name"] for row in ranked] == ["getUserById"] diff --git a/tests/test_graph_retrieval.py b/tests/test_graph_retrieval.py new file mode 100644 index 0000000..b781e69 --- /dev/null +++ b/tests/test_graph_retrieval.py @@ -0,0 +1,100 @@ +import sqlite3 + +from codebase_index.graph.retrieval import graph_candidates +from codebase_index.retrieval.types import Candidate + + +def _graph(*, confidence=True): + conn = sqlite3.connect(":memory:") + conn.row_factory = sqlite3.Row + confidence_col = ", confidence TEXT" if confidence else "" + conn.executescript( + f""" + CREATE TABLE files ( + id INTEGER PRIMARY KEY, path TEXT UNIQUE NOT NULL, + is_generated INTEGER DEFAULT 0, summary TEXT + ); + CREATE TABLE symbols ( + id INTEGER PRIMARY KEY, file_id INTEGER, name TEXT, kind TEXT, + line_start INTEGER, line_end INTEGER, signature TEXT + ); + CREATE TABLE edges ( + id INTEGER PRIMARY KEY, edge_type TEXT, src_kind TEXT, src_id INTEGER, + dst_kind TEXT, dst_id INTEGER, dst_name TEXT, file_id INTEGER, + line INTEGER, resolved INTEGER{confidence_col} + ); + CREATE INDEX idx_edges_src ON edges(src_kind, src_id); + CREATE INDEX idx_edges_dst ON edges(dst_kind, dst_id); + """ + ) + conn.executemany( + "INSERT INTO files(id, path, summary) VALUES (?, ?, ?)", + [(1, "src/a.py", "A"), (2, "src/b.py", "B"), (3, "src/c.py", "C")], + ) + conn.executemany( + "INSERT INTO symbols VALUES (?, ?, ?, ?, ?, ?, ?)", + [(10, 1, "alpha", "function", 2, 4, "def alpha()"), + (20, 2, "beta", "function", 3, 5, "def beta()"), + (30, 3, "gamma", "function", 7, 9, "def gamma()")], + ) + if confidence: + conn.executemany( + "INSERT INTO edges VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + [(1, "call", "symbol", 10, "symbol", 20, "beta", 1, 3, 1, "extracted"), + (2, "call", "symbol", 20, "symbol", 30, "gamma", 2, 4, 1, "inferred")], + ) + else: + conn.executemany( + "INSERT INTO edges VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + [(1, "call", "symbol", 10, "symbol", 20, "beta", 1, 3, 1), + (2, "call", "symbol", 20, "symbol", 30, "gamma", 2, 4, 1)], + ) + return conn + + +def test_graph_candidates_are_bounded_and_keep_provenance(): + conn = _graph() + seed = Candidate("src/a.py", 2, 4, "symbol", 1.0, symbol="alpha") + result = graph_candidates(conn, [seed], depth=2, node_cap=1, iterations=8) + assert len(result) == 1 + assert result[0].symbol == "beta" + assert "confidence=extracted" in result[0].reason + assert "graph provenance" in (result[0].content or "") + assert result[0].key() != seed.key() + + +def test_graph_candidates_are_deterministic_and_support_old_edges(): + seed = Candidate("src/a.py", 2, 4, "symbol", 1.0, symbol="alpha") + first = graph_candidates(_graph(confidence=False), [seed], depth=2, node_cap=5, iterations=5) + second = graph_candidates(_graph(confidence=False), [seed], depth=2, node_cap=5, iterations=5) + assert [(c.path, c.line_start, c.score) for c in first] == [ + (c.path, c.line_start, c.score) for c in second + ] + assert first and "confidence=unknown" in first[0].reason + + + +def test_graph_candidates_honor_direction(): + down = graph_candidates( + _graph(), + [Candidate("src/a.py", 2, 4, "symbol", 1.0, symbol="alpha")], + depth=2, + node_cap=5, + direction="down", + ) + assert [candidate.symbol for candidate in down] == ["beta", "gamma"] + + up = graph_candidates( + _graph(), + [Candidate("src/b.py", 3, 5, "symbol", 1.0, symbol="beta")], + depth=2, + node_cap=5, + direction="up", + ) + assert [candidate.symbol for candidate in up] == ["alpha"] + +def test_graph_candidates_tolerate_empty_or_partial_graphs(): + assert graph_candidates(sqlite3.connect(":memory:"), [], depth=2, node_cap=5) == [] + conn = sqlite3.connect(":memory:") + conn.execute("CREATE TABLE files(id INTEGER PRIMARY KEY, path TEXT)") + assert graph_candidates(conn, [Candidate("missing.py", 1, 1, "path", 1.0)]) == [] diff --git a/tests/test_impact_decay.py b/tests/test_impact_decay.py new file mode 100644 index 0000000..1fe9f5e --- /dev/null +++ b/tests/test_impact_decay.py @@ -0,0 +1,110 @@ +from __future__ import annotations + +from codebase_index.graph.builder import build_graph +from codebase_index.graph.expand import impact_lookup, walk_impact +from codebase_index.parsers.base import Symbol +from codebase_index.storage import repo +from codebase_index.storage.db import Database + + +def _file(db, path: str, sha: str) -> int: + return repo.upsert_file( + db.conn, + path=path, + lang="python", + size_bytes=1, + sha256=sha, + mtime_ns=1, + git_status=None, + parser="treesitter", + indexed_at="t", + is_generated=False, + ) + + +def _graph(db): + target_file = _file(db, "src/target.py", "target") + direct_file = _file(db, "src/direct.py", "direct") + transitive_file = _file(db, "src/transitive.py", "transitive") + repo.replace_symbols( + db.conn, + target_file, + [Symbol(name="root", kind="function", line_start=1, line_end=1, qualified="root")], + ) + direct = repo.replace_symbols( + db.conn, + direct_file, + [Symbol(name="direct", kind="function", line_start=1, line_end=1, qualified="direct")], + )[0] + transitive = repo.replace_symbols( + db.conn, + transitive_file, + [ + Symbol( + name="transitive", + kind="function", + line_start=1, + line_end=1, + qualified="transitive", + ) + ], + )[0] + repo.replace_edges( + db.conn, + direct_file, + [ + { + "edge_type": "call", + "src_kind": "symbol", + "src_id": direct, + "dst_kind": None, + "dst_id": None, + "dst_name": "root", + "line": 1, + "resolved": 0, + } + ], + ) + repo.replace_edges( + db.conn, + transitive_file, + [ + { + "edge_type": "call", + "src_kind": "symbol", + "src_id": transitive, + "dst_kind": None, + "dst_id": None, + "dst_name": "direct", + "line": 1, + "resolved": 0, + } + ], + ) + build_graph(db.conn) + return target_file + + +def test_impact_decay_keeps_direct_before_transitive(tmp_path): + db = Database(tmp_path / "index.sqlite").open() + _graph(db) + + nodes = walk_impact( + db.conn, "src/target.py", depth=2, direction="up", decay=0.5 + ) + assert [(node.name, node.distance) for node in nodes] == [ + ("direct", 1), + ("transitive", 2), + ] + + response = impact_lookup( + db.conn, "src/target.py", depth=2, direction="up", decay=0.5 + ) + assert response.files == ["src/direct.py", "src/transitive.py"] + db.close() + + +def test_impact_decay_missing_target_is_empty(tmp_path): + db = Database(tmp_path / "index.sqlite").open() + assert walk_impact(db.conn, "missing.py", depth=2, direction="up", decay=0.5) == [] + db.close() diff --git a/tests/test_lexical.py b/tests/test_lexical.py new file mode 100644 index 0000000..4e3b1b6 --- /dev/null +++ b/tests/test_lexical.py @@ -0,0 +1,66 @@ +from codebase_index.retrieval.lexical import ( + build_fts_query, + build_lexical_query, + escape_fts_term, + expansion_weights, + salient_terms, + split_identifier, +) + + +def test_split_identifier_handles_camel_pascal_snake_kebab_and_concatenation(): + assert split_identifier("getUserById") == ("get", "user", "by", "id") + assert split_identifier("PascalCase") == ("pascal", "case") + assert split_identifier("refresh_access-token") == ("refresh", "access", "token") + assert split_identifier("userid") == ("userid",) + + +def test_salient_terms_excludes_natural_language_filler(): + assert salient_terms("How does the auth config work?") == ("auth", "config") + + +def test_lexical_query_preserves_originals_and_adds_synonyms_subtokens(): + parsed = build_lexical_query("refresh_access_token auth") + + assert parsed.original_terms == ("refresh_access_token", "auth") + assert [term.term for term in parsed.expanded_terms] == [ + "refresh", + "access", + "token", + "authentication", + ] + assert all(term.weight < 1.0 for term in parsed.expanded_terms) + +def test_query_expansion_retains_camel_case_boundaries(): + parsed = build_lexical_query("getUserById") + + assert parsed.original_terms == ("getuserbyid",) + assert [term.term for term in parsed.expanded_terms] == ["get", "user", "by", "id"] + + +def test_synonym_weights_are_lower_than_original_terms(): + weights = expansion_weights("delete config") + + assert weights["delete"] == 1.0 + assert weights["remove"] < weights["delete"] + assert weights["config"] == 1.0 + assert weights["configuration"] < weights["config"] + + +def test_fts_terms_are_quoted_and_unsafe_syntax_is_literal(): + assert escape_fts_term('name" OR secret') == '"name"" OR secret"' + query = build_fts_query('delete " OR drop') + assert ' OR drop' not in query + assert '"delete"' in query + + +def test_natural_language_inflections_get_downweighted_bridges(): + parsed = build_lexical_query("secrets were redacted") + assert parsed.original_terms == ("secrets", "redacted") + assert [term.term for term in parsed.expanded_terms] == ["secret", "redact"] + assert all(term.weight < 1.0 for term in parsed.expanded_terms) + + +def test_morphological_merging_bridge_matches_merging_text(): + parsed = build_lexical_query("configuration get loaded and merged") + assert "merg" in [term.term for term in parsed.expanded_terms] diff --git a/tests/test_priors.py b/tests/test_priors.py new file mode 100644 index 0000000..574f0ec --- /dev/null +++ b/tests/test_priors.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from codebase_index.retrieval.priors import ( + SourceRole, + classify_source_role, + is_test_intent, + source_role_prior, +) + + +def test_classifies_roles_across_posix_and_windows_paths(): + assert ( + classify_source_role("src/codebase_index/retrieval/searchers.py") + is SourceRole.IMPLEMENTATION + ) + assert classify_source_role(r"C:\work\tests\test_searchers.py") is SourceRole.TEST + assert classify_source_role("docs/retrieval.md") is SourceRole.DOCUMENTATION + assert ( + classify_source_role(r"C:\work\node_modules\pkg\index.js") + is SourceRole.GENERATED_VENDOR_BUILD + ) + assert classify_source_role("misc/data.bin") is SourceRole.UNKNOWN + + +def test_prior_preserves_order_and_is_bounded(): + values = { + role: source_role_prior( + { + SourceRole.IMPLEMENTATION: "src/search.py", + SourceRole.TEST: "tests/test_search.py", + SourceRole.DOCUMENTATION: "docs/search.md", + SourceRole.GENERATED_VENDOR_BUILD: "dist/search.min.js", + SourceRole.UNKNOWN: "misc/data.bin", + }[role], + query="where is search implemented", + ) + for role in SourceRole + } + assert values[SourceRole.IMPLEMENTATION] > values[SourceRole.TEST] + assert values[SourceRole.IMPLEMENTATION] > values[SourceRole.DOCUMENTATION] + assert values[SourceRole.DOCUMENTATION] > values[SourceRole.GENERATED_VENDOR_BUILD] + assert all(-0.15 <= value <= 0.15 for value in values.values()) + assert ( + source_role_prior("src/search.py", query="where is search implemented") + == values[SourceRole.IMPLEMENTATION] + ) + + +def test_test_intent_does_not_demote_tests(): + assert is_test_intent("find the unit tests for search") + assert ( + source_role_prior("tests/test_search.py", query="find the unit tests for search") + > 0 + ) + assert ( + source_role_prior("tests/test_search.py", query="find the unit tests for search") + > source_role_prior( + "tests/test_search.py", query="where is search implemented" + ) + ) + assert ( + source_role_prior("tests/test_search.py", query="anything", intent="find_refs") + < 0 + ) + assert ( + source_role_prior("tests/test_search.py", query="anything", intent="test") + > 0 + ) + + +def test_role_classification_is_path_only_and_deterministic(): + path = r"C:\repo\src\contest\latest.py" + assert classify_source_role(path) is SourceRole.IMPLEMENTATION + assert classify_source_role(path) == classify_source_role(path) diff --git a/tests/test_search_cli.py b/tests/test_search_cli.py index 905891a..3924c65 100644 --- a/tests/test_search_cli.py +++ b/tests/test_search_cli.py @@ -35,6 +35,19 @@ def test_search_json_runs(tmp_path, monkeypatch): assert "results" in payload +def test_exact_symbol_search_reports_high_confidence(tmp_path, monkeypatch): + db_path = _build(tmp_path, monkeypatch) + monkeypatch.setenv("CBX_DB_PATH", str(db_path)) + result = runner.invoke( + app, + ["search", "refresh_access_token", "--mode", "symbol", "--json"], + ) + assert result.exit_code == 0, result.output + payload = _json.loads(result.stdout) + assert payload["confidence"] == "high" + assert payload["results"][0]["symbols"] == ["refresh_access_token"] + + def test_search_auto_indexes_when_missing(sample_repo, tmp_path): root = tmp_path / "copy" shutil.copytree(sample_repo, root) diff --git a/tests/test_symbol_ranking.py b/tests/test_symbol_ranking.py index b917628..93ec481 100644 --- a/tests/test_symbol_ranking.py +++ b/tests/test_symbol_ranking.py @@ -40,11 +40,27 @@ def test_multiterm_camelcase_coverage_outranks_single_term(tmp_path): db = _db(tmp_path) conn = db.conn rel = _insert_file(conn, path="src/Religion.java", lang="java", mtime_ns=1) - _insert_symbol(conn, rel, name="Religion", kind="class", - line_start=1, line_end=9, signature="class Religion", in_degree=3) + _insert_symbol( + conn, + rel, + name="Religion", + kind="class", + line_start=1, + line_end=9, + signature="class Religion", + in_degree=3, + ) mgr = _insert_file(conn, path="src/managers/ReligionManager.java", lang="java", mtime_ns=2) - _insert_symbol(conn, mgr, name="ReligionManager", kind="class", - line_start=1, line_end=99, signature="class ReligionManager", in_degree=5) + _insert_symbol( + conn, + mgr, + name="ReligionManager", + kind="class", + line_start=1, + line_end=99, + signature="class ReligionManager", + in_degree=5, + ) conn.commit() cands = symbol_candidates(conn, "religion manager belief and faith handling", limit=10) @@ -62,11 +78,25 @@ def test_exact_single_identifier_still_wins(tmp_path): db = _db(tmp_path) conn = db.conn rel = _insert_file(conn, path="src/Religion.java", lang="java", mtime_ns=1) - _insert_symbol(conn, rel, name="Religion", kind="class", - line_start=1, line_end=9, signature="class Religion") + _insert_symbol( + conn, + rel, + name="Religion", + kind="class", + line_start=1, + line_end=9, + signature="class Religion", + ) mgr = _insert_file(conn, path="src/managers/ReligionManager.java", lang="java", mtime_ns=2) - _insert_symbol(conn, mgr, name="ReligionManager", kind="class", - line_start=1, line_end=99, signature="class ReligionManager") + _insert_symbol( + conn, + mgr, + name="ReligionManager", + kind="class", + line_start=1, + line_end=99, + signature="class ReligionManager", + ) conn.commit() cands = symbol_candidates(conn, "Religion", limit=10) @@ -79,13 +109,80 @@ def test_underscore_names_also_covered(tmp_path): db = _db(tmp_path) conn = db.conn f = _insert_file(conn, path="auth/token.py", lang="python", mtime_ns=1) - _insert_symbol(conn, f, name="refresh_access_token", kind="function", - line_start=1, line_end=6, signature="def refresh_access_token()") - _insert_symbol(conn, f, name="refresh", kind="function", - line_start=8, line_end=9, signature="def refresh()") + _insert_symbol( + conn, + f, + name="refresh_access_token", + kind="function", + line_start=1, + line_end=6, + signature="def refresh_access_token()", + ) + _insert_symbol( + conn, + f, + name="refresh", + kind="function", + line_start=8, + line_end=9, + signature="def refresh()", + ) conn.commit() cands = symbol_candidates(conn, "how does refresh access token work", limit=10) names = [c.symbol for c in cands] db.close() assert names.index("refresh_access_token") < names.index("refresh") + + +def test_framing_words_do_not_hide_exact_symbol(tmp_path): + db = _db(tmp_path) + conn = db.conn + f = _insert_file(conn, path="src/graph/expand.py", lang="python", mtime_ns=1) + _insert_symbol( + conn, + f, + name="walk_impact", + kind="function", + line_start=1, + line_end=8, + signature="def walk_impact()", + ) + _insert_symbol( + conn, + f, + name="walk", + kind="function", + line_start=10, + line_end=12, + signature="def walk()", + ) + conn.commit() + + cands = symbol_candidates(conn, "find walk_impact", limit=10) + db.close() + + assert cands[0].symbol == "walk_impact" + assert cands[0].exact_symbol is True + + +def test_symbol_synonym_expansion_retrieves_alias(tmp_path): + db = _db(tmp_path) + conn = db.conn + f = _insert_file(conn, path="src/config.py", lang="python", mtime_ns=1) + _insert_symbol( + conn, + f, + name="Config", + kind="class", + line_start=1, + line_end=12, + signature="class Config", + ) + conn.commit() + + cands = symbol_candidates(conn, "configuration", limit=10) + db.close() + + assert cands[0].symbol == "Config" + assert cands[0].exact_symbol is False