"""Executable static-retrieval pilot with immutable per-run artifacts.""" from __future__ import annotations from dataclasses import asdict import json import math from pathlib import Path import resource import subprocess import time from typing import Any, Sequence from .lm_studio_embeddings import LMStudioEmbeddingClient from .repository import GitSnapshot, SourceChunk, chunk_snapshot from .retrieval import ( BM25FuzzyRetriever, DenseRetriever, ExactRetriever, SQLiteEmbeddingCache, unique_file_ranking, ) from .specs import ( ExperimentSpec, HarnessSpec, TaskSpec, load_embeddings, load_experiments, load_harnesses, load_models, load_task_split, load_tasks, ) from .telemetry import EventWriter, RunIdentity class PilotError(RuntimeError): """Raised when the development pilot cannot produce an auditable run.""" def _git_text(repository: Path, arguments: list[str]) -> str: result = subprocess.run( ["git", *arguments], cwd=repository, check=False, capture_output=True, text=True, timeout=30, ) if result.returncode != 0: raise PilotError(result.stderr.strip() or f"git {' '.join(arguments)} failed") return result.stdout.strip() def research_code_revision(root: Path) -> str: revision = _git_text(root, ["rev-parse", "HEAD"]) if _git_text(root, ["status", "--porcelain"]): raise PilotError("Research worktree is dirty; commit the exact implementation before a run") return revision def retrieval_metrics(ranked_paths: Sequence[str], gold_files: Sequence[str]) -> dict[str, Any]: gold = set(gold_files) if not gold: raise PilotError("retrieval task has no gold files") metrics: dict[str, Any] = {} for cutoff in (1, 5, 10): retrieved = set(ranked_paths[:cutoff]) metrics[f"file_recall_at_{cutoff}"] = len(gold & retrieved) / len(gold) first_gold_rank = next( (index for index, path in enumerate(ranked_paths, start=1) if path in gold), None, ) metrics["first_gold_rank"] = first_gold_rank metrics["mrr"] = 0.0 if first_gold_rank is None else 1.0 / first_gold_rank dcg = sum( 1.0 / math.log2(rank + 1) for rank, path in enumerate(ranked_paths[:10], start=1) if path in gold ) ideal_hits = min(len(gold), 10) ideal_dcg = sum(1.0 / math.log2(rank + 1) for rank in range(1, ideal_hits + 1)) metrics["ndcg_at_10"] = dcg / ideal_dcg metrics["all_gold_in_top_10"] = gold.issubset(set(ranked_paths[:10])) return metrics def _memory_sample() -> dict[str, Any]: usage = resource.getrusage(resource.RUSAGE_SELF) return {"process_max_rss_platform_units": usage.ru_maxrss} def _select_retriever( harness: HarnessSpec, chunks: Sequence[SourceChunk], embedding_spec: Any, embedding_client: LMStudioEmbeddingClient, cache: SQLiteEmbeddingCache, ) -> tuple[Any, dict[str, Any]]: started = time.monotonic() if harness.harness_id == "H000": return ExactRetriever(chunks), { "index_build_seconds": time.monotonic() - started, "index_kind": "literal_term_scan", } if harness.harness_id == "H001": retriever = BM25FuzzyRetriever(chunks) return retriever, { "index_build_seconds": time.monotonic() - started, "index_kind": "in_memory_bm25_fuzzy", } if harness.harness_id == "H003": retriever, stats = DenseRetriever.build( chunks, embedding_spec, embedding_client, cache, ) return retriever, {"index_kind": "python_flat_cosine", **asdict(stats)} raise PilotError(f"E00 runner does not implement {harness.harness_id}") def run_static_retrieval_pilot( root: Path, repository: Path, experiment_id: str = "E00", 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) try: experiment: ExperimentSpec = experiments[experiment_id] except KeyError as exc: raise PilotError(f"Unknown experiment {experiment_id}") from exc if experiment.mode != "static_retrieval": raise PilotError("pilot runner only supports static_retrieval experiments") 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 PilotError("task or harness filters selected no pilot cells") if any(task.validation_status not in {"retrieval_ready", "end_to_end_ready"} for task in selected_tasks): raise PilotError("pilot split contains a task that is not retrieval-ready") model_spec = models[experiment.model_ids[0]] embedding_spec = embeddings[experiment.embedding_id] embedding_client = LMStudioEmbeddingClient(embedding_spec, timeout_seconds=60.0) embedding_runtime = embedding_client.resolve() resident_models = embedding_client.loaded_model_keys() if set(resident_models) != {embedding_spec.model_key}: raise PilotError( "Memory-safe E00 requires only the embedding model to be resident; " f"observed {resident_models}" ) snapshot = GitSnapshot(repository) origin = _git_text(repository, ["remote", "get-url", "origin"]) expected_origins = {task.repository_url for task in selected_tasks} if expected_origins != {origin}: raise PilotError(f"Repository remote mismatch: expected {expected_origins}, observed {origin!r}") cache_path = root / "indexes" / "embeddings" / f"{embedding_spec.config_hash}.sqlite3" summary_rows: list[dict[str, Any]] = [] chunk_cache: dict[str, tuple[SourceChunk, ...]] = {} with SQLiteEmbeddingCache(cache_path, embedding_spec) as embedding_cache: for task in selected_tasks: snapshot.verify_commit(task.base_commit) snapshot.verify_commit(task.gold_commit) if task.base_commit not in chunk_cache: chunk_cache[task.base_commit] = chunk_snapshot( snapshot, task.base_commit, embedding_spec.chunk_lines, embedding_spec.chunk_overlap_lines, embedding_spec.chunk_char_limit, ) chunks = chunk_cache[task.base_commit] snapshot_paths = {chunk.path for chunk in chunks} missing_gold = set(task.gold_files) - snapshot_paths if missing_gold: raise PilotError(f"{task.task_id} gold files absent at base commit: {sorted(missing_gold)}") 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, ) with EventWriter( root / "results", identity, asdict(harness), { "agent_model": asdict(model_spec), "embedding_model": asdict(embedding_spec), "embedding_runtime": embedding_runtime, }, ) as writer: writer.emit( "run_started", { "task_config_hash": task.config_hash, "task_validation_status": task.validation_status, "repository_origin": origin, "base_commit": task.base_commit, "gold_commit": task.gold_commit, "source_file_count": len(snapshot_paths), "source_chunk_count": len(chunks), "candidate_limit": candidate_limit, }, ) writer.emit("resource_sample", _memory_sample()) retriever, index_stats = _select_retriever( harness, chunks, embedding_spec, embedding_client, embedding_cache, ) query_started = time.monotonic() candidates = tuple(retriever.retrieve(task.statement, candidate_limit)) query_seconds = time.monotonic() - query_started files = unique_file_ranking(candidates) ranked_paths = [candidate.path for candidate in files] metrics = retrieval_metrics(ranked_paths, task.gold_files) metrics.update( { "candidate_count": len(candidates), "unique_file_count": len(files), "query_seconds": query_seconds, **index_stats, } ) 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, "chunk_id": (candidate.metadata or {}).get("chunk_id"), "is_gold_file": candidate.path in set(task.gold_files), }, ) ranking_payload = [ { "rank": rank, "path": candidate.path, "line_start": candidate.line_start, "line_end": candidate.line_end, "score": candidate.score, "source": candidate.source, "is_gold_file": candidate.path in set(task.gold_files), } for rank, candidate in enumerate(files, start=1) ] writer.write_artifact("ranking.json", json.dumps(ranking_payload, 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}) summary_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, "development_only": True, "repository": str(repository.resolve()), "repository_origin": origin, "code_revision": code_revision, "embedding_config_hash": embedding_spec.config_hash, "resident_models": resident_models, "task_count": len(selected_tasks), "harness_count": len(selected_harnesses), "run_count": len(summary_rows), "runs": summary_rows, } report_path = root / "results" / "reports" / f"{experiment.experiment_id}_{code_revision[:12]}.json" report_path.parent.mkdir(parents=True, exist_ok=True) with report_path.open("x", encoding="utf-8") as handle: json.dump(summary, handle, indent=2, sort_keys=True) handle.write("\n") return summary