"""E01 confirmatory retrieval runner over the frozen held-out task split.""" from __future__ import annotations from dataclasses import asdict import json from pathlib import Path import resource import time from typing import Any, Sequence from .components import Candidate from .fusion import reciprocal_rank_fusion, unique_files from .lm_studio_embeddings import LMStudioEmbeddingClient from .pilot import PilotError, research_code_revision, retrieval_metrics from .repository import GitSnapshot, SourceChunk, chunk_snapshot from .retrieval import BM25FuzzyRetriever, DenseRetriever, ExactRetriever, SQLiteEmbeddingCache from .specs import ( HarnessSpec, load_embeddings, load_experiments, load_harnesses, load_models, load_task_split, load_tasks, ) from .syntax_index import GoSymbol, SymbolGraph, SyntaxRetriever, parse_snapshot, symbol_hits from .telemetry import EventWriter, RunIdentity, run_directory from .tokenization import QwenTokenCounter from .vector_backends import FaissFlatRetriever class ConfirmatoryRetrievalError(RuntimeError): """Raised when E01 cannot preserve the frozen retrieval protocol.""" def memory_sample() -> dict[str, Any]: usage = resource.getrusage(resource.RUSAGE_SELF) return {"process_max_rss_platform_units": usage.ru_maxrss} def retrieve_treatment( harness: HarnessSpec, query: str, exact: ExactRetriever, lexical: BM25FuzzyRetriever, syntax: SyntaxRetriever, dense: FaissFlatRetriever, graph: SymbolGraph, limit: int, ) -> tuple[Candidate, ...]: rankings: list[Sequence[Candidate]] = [] if harness.lexical: rankings.append(lexical.retrieve(query, limit)) if harness.syntax == "tree_sitter": rankings.append(syntax.retrieve(query, limit)) if harness.dense: rankings.append(dense.retrieve(query, limit)) if not rankings: result = tuple(exact.retrieve(query, limit)) elif len(rankings) == 1: result = tuple(rankings[0]) else: result = reciprocal_rank_fusion(rankings, limit) result = unique_files(result) if harness.graph_hops: result = graph.expand(result, harness.graph_hops, limit) return unique_files(result)[:limit] def extended_metrics( candidates: Sequence[Candidate], gold_files: Sequence[str], gold_symbols: Sequence[str], symbols: Sequence[GoSymbol], tokenizer: QwenTokenCounter, context_budget: int, ) -> dict[str, Any]: ranked_paths = [candidate.path for candidate in candidates] metrics = retrieval_metrics(ranked_paths, gold_files) hits = symbol_hits(candidates, gold_symbols, symbols, cutoff=10) metrics["function_recall_at_10"] = len(hits) / len(gold_symbols) metrics["gold_symbol_hits_at_10"] = sorted(hits) _, included, tokens = tokenizer.pack_ranked(candidates, context_budget) included_paths = {candidate.path for candidate in included} metrics["all_gold_within_token_budget"] = set(gold_files).issubset(included_paths) metrics["packed_candidate_count"] = len(included) metrics["packed_tokens"] = tokens return metrics def run_confirmatory_retrieval( root: Path, repository: Path, experiment_id: str = "E01", task_filter: set[str] | None = None, harness_filter: set[str] | None = None, candidate_limit: int = 200, ) -> dict[str, Any]: code_revision = research_code_revision(root) experiments = load_experiments(root) harnesses = load_harnesses(root) models = load_models(root) embeddings = load_embeddings(root) tasks = load_tasks(root) experiment = experiments.get(experiment_id) if experiment is None or experiment.mode != "static_retrieval" or experiment_id == "E00": raise ConfirmatoryRetrievalError("runner requires a confirmatory static_retrieval experiment") split = load_task_split(root / "tasks" / "splits" / f"{experiment.task_split}.txt") selected_tasks = [tasks[item] for item in split if task_filter is None or item in task_filter] selected_harnesses = [ harnesses[item] for item in experiment.harness_ids if harness_filter is None or item in harness_filter ] if not selected_tasks or not selected_harnesses: raise ConfirmatoryRetrievalError("task or harness filters selected no E01 cells") model_spec = models[experiment.model_ids[0]] embedding_spec = embeddings[experiment.embedding_id] embedding_client = LMStudioEmbeddingClient(embedding_spec, timeout_seconds=120.0) embedding_runtime = embedding_client.resolve() resident_models = embedding_client.loaded_model_keys() if tuple(resident_models) != (embedding_spec.model_key,): raise ConfirmatoryRetrievalError( "E01 requires exclusive embedding-model residency; " f"observed {resident_models}" ) tokenizer = QwenTokenCounter() snapshot = GitSnapshot(repository) origin = subprocess_git(repository, ["remote", "get-url", "origin"]) if {task.repository_url for task in selected_tasks} != {origin}: raise ConfirmatoryRetrievalError("repository origin does not match frozen tasks") cache_path = root / "indexes" / "embeddings" / f"{embedding_spec.config_hash}.sqlite3" rows: list[dict[str, Any]] = [] with SQLiteEmbeddingCache(cache_path, embedding_spec) as embedding_cache: for task in selected_tasks: started = time.monotonic() chunks: tuple[SourceChunk, ...] = chunk_snapshot( snapshot, task.base_commit, embedding_spec.chunk_lines, embedding_spec.chunk_overlap_lines, embedding_spec.chunk_char_limit, ) symbols = parse_snapshot(snapshot, task.base_commit) exact = ExactRetriever(chunks) lexical = BM25FuzzyRetriever(chunks) syntax = SyntaxRetriever(symbols) dense_base, dense_stats = DenseRetriever.build( chunks, embedding_spec, embedding_client, embedding_cache ) dense = FaissFlatRetriever(dense_base) graph = SymbolGraph(symbols) shared_index_seconds = time.monotonic() - started available_paths = {chunk.path for chunk in chunks} missing = set(task.gold_files) - available_paths if missing: raise ConfirmatoryRetrievalError( f"{task.task_id} gold files absent at base commit: {sorted(missing)}" ) for harness in selected_harnesses: identity = RunIdentity( experiment_id=experiment.experiment_id, task_id=task.task_id, harness_id=harness.harness_id, harness_hash=harness.config_hash, model_id=model_spec.model_id, model_key=model_spec.expected_inference_key, model_config_hash=model_spec.config_hash, context_budget=experiment.context_budgets[0], seed=experiment.seeds[0], repetition=0, repository_sha=task.base_commit, code_revision=code_revision, ) directory = run_directory(root / "results", identity) if directory.exists(): final_path = directory / "final_metrics.json" if not final_path.exists(): raise ConfirmatoryRetrievalError(f"incomplete pre-existing run: {directory}") final = json.loads(final_path.read_text(encoding="utf-8")) rows.append({"run_id": identity.run_id, "task_id": task.task_id, "harness_id": harness.harness_id, **final}) continue with EventWriter( root / "results", identity, asdict(harness), { "agent_model_not_loaded": asdict(model_spec), "embedding_model": asdict(embedding_spec), "embedding_runtime": embedding_runtime, "tokenizer_path": str(tokenizer.path), "tokenizer_sha256": tokenizer.sha256, }, ) as writer: writer.emit( "run_started", { "confirmatory": True, "task_config_hash": task.config_hash, "candidate_limit": candidate_limit, "source_file_count": len(available_paths), "source_chunk_count": len(chunks), "symbol_count": len(symbols), "shared_index_seconds": shared_index_seconds, }, ) writer.emit("resource_sample", memory_sample()) query_started = time.monotonic() candidates = retrieve_treatment( harness, task.statement, exact, lexical, syntax, dense, graph, candidate_limit, ) query_seconds = time.monotonic() - query_started metrics = extended_metrics( candidates, task.gold_files, task.gold_symbols, symbols, tokenizer, experiment.context_budgets[0], ) metrics.update( { "candidate_count": len(candidates), "query_seconds": query_seconds, "shared_index_seconds": shared_index_seconds, "dense_total_chunks": dense_stats.total_chunks, "dense_cached_chunks": dense_stats.cached_chunks, "dense_embedded_chunks": dense_stats.embedded_chunks, "dense_vector_load_seconds": dense_stats.build_seconds, "faiss_build_seconds": dense.stats.build_seconds, } ) gold_files = set(task.gold_files) for rank, candidate in enumerate(candidates, 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 gold_files, }, ) ranking = [ { "rank": rank, "path": candidate.path, "line_start": candidate.line_start, "line_end": candidate.line_end, "score": candidate.score, "source": candidate.source, "symbol": candidate.symbol, "is_gold_file": candidate.path in gold_files, } for rank, candidate in enumerate(candidates, start=1) ] writer.write_artifact("ranking.json", json.dumps(ranking, indent=2) + "\n") writer.write_artifact("final_metrics.json", json.dumps(metrics, indent=2) + "\n") writer.emit("resource_sample", memory_sample()) writer.emit("run_finished", {"status": "completed", "metrics": metrics}) rows.append( {"run_id": identity.run_id, "task_id": task.task_id, "harness_id": harness.harness_id, **metrics} ) summary = { "schema_version": 1, "experiment_id": experiment.experiment_id, "confirmatory": True, "repository": str(repository.resolve()), "repository_origin": origin, "code_revision": code_revision, "embedding_config_hash": embedding_spec.config_hash, "tokenizer_sha256": tokenizer.sha256, "resident_models": resident_models, "task_count": len(selected_tasks), "harness_count": len(selected_harnesses), "run_count": len(rows), "runs": rows, } report = root / "results" / "reports" / f"{experiment.experiment_id}_{code_revision[:12]}_{int(time.time())}.json" report.parent.mkdir(parents=True, exist_ok=True) report.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8") summary["report_path"] = str(report) return summary def subprocess_git(repository: Path, arguments: list[str]) -> str: import subprocess result = subprocess.run( ["git", *arguments], cwd=repository, check=False, capture_output=True, text=True, timeout=30 ) if result.returncode: raise PilotError(result.stderr.strip() or "git command failed") return result.stdout.strip()