"""E05 dense-index backend systems experiment.""" from __future__ import annotations from dataclasses import asdict from hashlib import sha256 import json from pathlib import Path import statistics import time from typing import Any, Sequence import faiss import psutil from .components import Candidate from .confirmatory_retrieval import extended_metrics, subprocess_git from .fusion import reciprocal_rank_fusion, unique_files from .lm_studio_embeddings import LMStudioEmbeddingClient from .pilot import research_code_revision from .repository import GitSnapshot, chunk_snapshot from .retrieval import BM25FuzzyRetriever, DenseRetriever, SQLiteEmbeddingCache from .specs import ( BackendSpec, HarnessSpec, load_backends, load_embeddings, load_experiments, load_harnesses, load_models, load_task_split, load_tasks, ) from .syntax_index import SyntaxRetriever, parse_snapshot from .telemetry import EventWriter, RunIdentity, run_directory from .tokenization import QwenTokenCounter from .vector_backends import FaissFlatRetriever, FaissHNSWRetriever, SQLiteVecRetriever class BackendExperimentError(RuntimeError): """Raised when E05 cannot execute its frozen backend protocol.""" def percentile(values: Sequence[float], fraction: float) -> float: ordered = sorted(values) if not ordered: raise ValueError("percentile requires observations") position = (len(ordered) - 1) * fraction lower = int(position) upper = min(lower + 1, len(ordered) - 1) weight = position - lower return ordered[lower] * (1.0 - weight) + ordered[upper] * weight def build_backend( backend: BackendSpec, dense: DenseRetriever, index_path: Path, ) -> tuple[Any, int, int]: process = psutil.Process() before = process.memory_info().rss if backend.backend_id == "B001": instance = FaissFlatRetriever(dense) faiss.write_index(instance.index, str(index_path)) elif backend.backend_id == "B002": instance = FaissHNSWRetriever( dense, neighbors=int(backend.neighbors or 32), ef_construction=int(backend.ef_construction or 80), ef_search=int(backend.ef_search or 64), ) faiss.write_index(instance.index, str(index_path)) elif backend.backend_id == "B003": instance = SQLiteVecRetriever(dense, index_path) else: raise BackendExperimentError(f"unsupported backend {backend.backend_id}") after = process.memory_info().rss return instance, max(after - before, 0), index_path.stat().st_size def treatment_ranking( harness: HarnessSpec, dense_ranking: Sequence[Candidate], lexical_ranking: Sequence[Candidate], syntax_ranking: Sequence[Candidate], limit: int, ) -> tuple[Candidate, ...]: if harness.harness_id == "H003": return unique_files(tuple(dense_ranking))[:limit] if harness.harness_id == "H005": return reciprocal_rank_fusion([lexical_ranking, dense_ranking], limit) if harness.harness_id == "H007": return reciprocal_rank_fusion([lexical_ranking, syntax_ranking, dense_ranking], limit) raise BackendExperimentError(f"E05 does not implement {harness.harness_id}") def run_backend_experiment( root: Path, repository: Path, experiment_id: str = "E05", task_filter: set[str] | None = None, backend_filter: set[str] | None = None, harness_filter: set[str] | None = None, candidate_limit: int = 200, ) -> dict[str, Any]: revision = research_code_revision(root) experiments = load_experiments(root) experiment = experiments.get(experiment_id) if experiment is None or experiment.mode != "index_backend": raise BackendExperimentError("runner requires the frozen E05 index_backend experiment") harness_catalog = load_harnesses(root) backend_catalog = load_backends(root) model = load_models(root)[experiment.model_ids[0]] embedding = load_embeddings(root)[experiment.embedding_id] task_catalog = load_tasks(root) split = load_task_split(root / "tasks" / "splits" / f"{experiment.task_split}.txt") tasks = [task_catalog[item] for item in split if task_filter is None or item in task_filter] harnesses = [ harness_catalog[item] for item in experiment.harness_ids if harness_filter is None or item in harness_filter ] backends = [ backend_catalog[item] for item in experiment.backend_ids if backend_filter is None or item in backend_filter ] if not tasks or not harnesses or not backends: raise BackendExperimentError("filters selected no E05 cells") client = LMStudioEmbeddingClient(embedding, timeout_seconds=120.0) runtime = client.resolve() resident = client.loaded_model_keys() if tuple(resident) != (embedding.model_key,): raise BackendExperimentError(f"E05 requires exclusive embedding residency; observed {resident}") tokenizer = QwenTokenCounter() snapshot = GitSnapshot(repository) origin = subprocess_git(repository, ["remote", "get-url", "origin"]) rows: list[dict[str, Any]] = [] cache_path = root / "indexes" / "embeddings" / f"{embedding.config_hash}.sqlite3" with SQLiteEmbeddingCache(cache_path, embedding) as cache: for task in tasks: chunks = chunk_snapshot( snapshot, task.base_commit, embedding.chunk_lines, embedding.chunk_overlap_lines, embedding.chunk_char_limit, ) symbols = parse_snapshot(snapshot, task.base_commit) dense_base, dense_stats = DenseRetriever.build(chunks, embedding, client, cache) lexical_ranking = BM25FuzzyRetriever(chunks).retrieve(task.statement, candidate_limit) syntax_ranking = SyntaxRetriever(symbols).retrieve(task.statement, candidate_limit) index_dir = root / "indexes" / "e05" / task.base_commit index_dir.mkdir(parents=True, exist_ok=True) instances: dict[str, tuple[Any, int, int]] = {} for backend in backends: suffix = ".sqlite3" if backend.backend_id == "B003" else ".faiss" instances[backend.backend_id] = build_backend( backend, dense_base, index_dir / f"{backend.backend_id}{suffix}", ) flat_paths = [ candidate.path for candidate in unique_files( tuple(instances["B001"][0].retrieve(task.statement, candidate_limit)) )[:10] ] if "B001" in instances else [] for backend in backends: instance, ram_delta, disk_bytes = instances[backend.backend_id] for seed in experiment.seeds: timings: list[float] = [] dense_ranking: Sequence[Candidate] = () for _ in range(backend.query_repetitions): started = time.perf_counter() dense_ranking = instance.retrieve(task.statement, candidate_limit) timings.append((time.perf_counter() - started) * 1000.0) for harness in harnesses: treatment_id = f"{harness.harness_id}_{backend.backend_id}" treatment_hash = sha256( f"{harness.config_hash}\0{backend.config_hash}".encode("utf-8") ).hexdigest() identity = RunIdentity( experiment_id=experiment.experiment_id, task_id=task.task_id, harness_id=treatment_id, harness_hash=treatment_hash, model_id=model.model_id, model_key=model.expected_inference_key, model_config_hash=model.config_hash, context_budget=experiment.context_budgets[0], seed=seed, repetition=0, repository_sha=task.base_commit, code_revision=revision, ) directory = run_directory(root / "results", identity) if directory.exists(): final_path = directory / "final_metrics.json" if not final_path.exists(): raise BackendExperimentError(f"incomplete E05 run: {directory}") final = json.loads(final_path.read_text(encoding="utf-8")) rows.append({"run_id": identity.run_id, **final}) continue ranking = treatment_ranking( harness, dense_ranking, lexical_ranking, syntax_ranking, candidate_limit, ) metrics = extended_metrics( ranking, task.gold_files, task.gold_symbols, symbols, tokenizer, experiment.context_budgets[0], ) backend_top = [candidate.path for candidate in unique_files(tuple(dense_ranking))[:10]] metrics.update( { "experiment_id": experiment.experiment_id, "task_id": task.task_id, "harness_id": harness.harness_id, "backend_id": backend.backend_id, "seed": seed, "backend_recall_at_10_vs_flat": ( len(set(backend_top) & set(flat_paths)) / 10.0 if flat_paths else None ), "index_build_seconds": instance.stats.build_seconds, "index_ram_bytes_delta": ram_delta, "index_disk_bytes": disk_bytes, "query_repetitions": backend.query_repetitions, "query_mean_ms": statistics.fmean(timings), "query_p50_ms": percentile(timings, 0.50), "query_p95_ms": percentile(timings, 0.95), "dense_cached_chunks": dense_stats.cached_chunks, "dense_embedded_chunks": dense_stats.embedded_chunks, } ) with EventWriter( root / "results", identity, {"harness": asdict(harness), "backend": asdict(backend)}, { "agent_model_not_loaded": asdict(model), "embedding_model": asdict(embedding), "embedding_runtime": runtime, }, ) as writer: writer.emit("run_started", {"confirmatory": True, "candidate_limit": candidate_limit}) for rank, candidate in enumerate(ranking, start=1): writer.emit( "retrieval_candidate", { "rank": rank, "path": candidate.path, "line_start": candidate.line_start, "line_end": candidate.line_end, "source": candidate.source, "score": candidate.score, "symbol": candidate.symbol, "is_gold_file": candidate.path in set(task.gold_files), }, ) writer.write_artifact( "final_metrics.json", json.dumps(metrics, indent=2) + "\n" ) writer.emit("run_finished", {"status": "completed", "metrics": metrics}) rows.append({"run_id": identity.run_id, **metrics}) for instance, _, _ in instances.values(): if isinstance(instance, SQLiteVecRetriever): instance.close() summary = { "schema_version": 1, "experiment_id": experiment.experiment_id, "confirmatory": True, "repository_origin": origin, "code_revision": revision, "task_count": len(tasks), "harness_count": len(harnesses), "backend_count": len(backends), "seed_count": len(experiment.seeds), "run_count": len(rows), "rows": rows, } report = root / "results" / "reports" / f"E05_{revision[:12]}_{int(time.time())}.json" report.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") summary["report_path"] = str(report) return summary