"""E04 stale-index and plausible-distractor robustness experiment.""" from __future__ import annotations from dataclasses import asdict from hashlib import sha256 import json from pathlib import Path import time from typing import Any, Sequence from .components import Candidate from .confirmatory_retrieval import extended_metrics, memory_sample, retrieve_treatment, subprocess_git from .lm_studio_embeddings import LMStudioEmbeddingClient from .pilot import research_code_revision from .repository import GitSnapshot, SourceChunk, SourceFile, chunk_file, chunk_snapshot from .retrieval import BM25FuzzyRetriever, DenseRetriever, ExactRetriever, SQLiteEmbeddingCache, query_terms from .specs import HarnessSpec, TaskSpec, load_embeddings, load_experiments, load_harnesses, load_models, load_task_split, load_tasks from .syntax_index import GoSymbol, SymbolGraph, SyntaxRetriever, parse_go_file, parse_snapshot from .telemetry import EventWriter, RunIdentity, run_directory from .tokenization import QwenTokenCounter from .vector_backends import FaissFlatRetriever E02_QUERY_REVISION = "1a7066f6c7682793b6f04445c776a93cc4fac895" DISTRACTOR_SEVERITIES = {0: 1, 1: 5, 2: 10} DISTRACTOR_PREFIX = "__harness_distractors__" class RobustnessExperimentError(RuntimeError): """Raised when the frozen E04 robustness protocol cannot be preserved.""" def distractor_severity(seed: int) -> int: """Map the frozen E04 seed coordinate to a nested distractor dose.""" try: return DISTRACTOR_SEVERITIES[seed] except KeyError as exc: raise RobustnessExperimentError(f"E04 has no distractor dose for seed {seed}") from exc def synthetic_distractor_sources(task: TaskSpec, severity: int) -> tuple[SourceFile, ...]: """Create deterministic, valid, in-memory Go distractors with issue vocabulary.""" if severity not in set(DISTRACTOR_SEVERITIES.values()): raise RobustnessExperimentError(f"unsupported distractor severity {severity}") terms = query_terms(task.statement) vocabulary = " ".join(terms[:40]) or "repository repair" issue = " ".join(task.statement.split()) task_name = "".join(part.capitalize() for part in task.task_id.lower().split("_")) sources: list[SourceFile] = [] for index in range(1, severity + 1): path = f"{DISTRACTOR_PREFIX}/{task.task_id.lower()}/distractor_{index:03d}.go" name = f"{task_name}PlausibleResolver{index:03d}" text = ( "package harnessdistractor\n\n" f"// {name} appears related to this issue: {issue}\n" f"type {name} struct {{\n" "\tEnabled bool\n" "}\n\n" f"// Resolve handles {vocabulary}.\n" f"func (value {name}) Resolve() string {{\n" f"\treturn {json.dumps(vocabulary)}\n" "}\n" ) sources.append(SourceFile(path=path, text=text)) return tuple(sources) def robustness_shift( baseline: dict[str, Any], perturbed: dict[str, Any], baseline_paths: Sequence[str], perturbed_paths: Sequence[str], gold_files: Sequence[str], missing_rank: int, ) -> dict[str, Any]: """Compute paired degradation measures, including a predeclared censored rank.""" baseline_rank = baseline["first_gold_rank"] perturbed_rank = perturbed["first_gold_rank"] baseline_censored = baseline_rank if baseline_rank is not None else missing_rank perturbed_censored = perturbed_rank if perturbed_rank is not None else missing_rank baseline_gold_at_10 = set(baseline_paths[:10]) & set(gold_files) perturbed_gold_at_10 = set(perturbed_paths[:10]) & set(gold_files) retention = ( None if not baseline_gold_at_10 else len(baseline_gold_at_10 & perturbed_gold_at_10) / len(baseline_gold_at_10) ) return { "file_recall_at_10_delta": perturbed["file_recall_at_10"] - baseline["file_recall_at_10"], "mrr_delta": perturbed["mrr"] - baseline["mrr"], "ndcg_at_10_delta": perturbed["ndcg_at_10"] - baseline["ndcg_at_10"], "first_gold_rank_displacement": ( None if baseline_rank is None or perturbed_rank is None else perturbed_rank - baseline_rank ), "first_gold_rank_displacement_censored": perturbed_censored - baseline_censored, "baseline_gold_top_10_retention": retention, "lost_all_top_10_gold": bool(baseline_gold_at_10 and not perturbed_gold_at_10), } def load_iterative_query(root: Path, task_id: str) -> dict[str, Any]: path = ( root / "results" / "staging" / "E02" / E02_QUERY_REVISION / "H010" / task_id / "query_stage.json" ) if not path.exists(): raise RobustnessExperimentError(f"missing pinned E02 query artifact: {path}") value = json.loads(path.read_text(encoding="utf-8")) if value.get("task_id") != task_id or value.get("harness_id") != "H010": raise RobustnessExperimentError(f"mismatched E02 query artifact: {path}") if value.get("query_source") not in {"model", "issue_fallback"}: raise RobustnessExperimentError(f"unamended E02 query artifact: {path}") if not isinstance(value.get("query"), str) or not value["query"].strip(): raise RobustnessExperimentError(f"empty E02 query artifact: {path}") return value def _build_index( chunks: Sequence[SourceChunk], symbols: Sequence[GoSymbol], embedding_spec: Any, embedding_client: LMStudioEmbeddingClient, embedding_cache: SQLiteEmbeddingCache, ) -> tuple[ExactRetriever, BM25FuzzyRetriever, SyntaxRetriever, FaissFlatRetriever, SymbolGraph, dict[str, Any]]: started = time.monotonic() dense_base, dense_stats = DenseRetriever.build( chunks, embedding_spec, embedding_client, embedding_cache ) dense = FaissFlatRetriever(dense_base) return ( ExactRetriever(chunks), BM25FuzzyRetriever(chunks), SyntaxRetriever(symbols), dense, SymbolGraph(symbols), { "source_file_count": len({chunk.path for chunk in chunks}), "source_chunk_count": len(chunks), "symbol_count": len(symbols), "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, "total_index_seconds": time.monotonic() - started, }, ) def _retrieve( harness: HarnessSpec, query: str, index: tuple[ExactRetriever, BM25FuzzyRetriever, SyntaxRetriever, FaissFlatRetriever, SymbolGraph, dict[str, Any]], limit: int, ) -> tuple[Candidate, ...]: return retrieve_treatment(harness, query, *index[:5], limit) def _ranking(candidates: Sequence[Candidate], gold_files: Sequence[str]) -> list[dict[str, Any]]: gold = set(gold_files) return [ { "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, "is_synthetic_distractor": candidate.path.startswith(f"{DISTRACTOR_PREFIX}/"), } for rank, candidate in enumerate(candidates, start=1) ] def _metrics( candidates: Sequence[Candidate], task: TaskSpec, evaluation_symbols: Sequence[GoSymbol], tokenizer: QwenTokenCounter, context_budget: int, ) -> dict[str, Any]: return extended_metrics( candidates, task.gold_files, task.gold_symbols, evaluation_symbols, tokenizer, context_budget, ) def run_robustness_experiment( root: Path, repository: Path, experiment_id: str = "E04", task_filter: set[str] | None = None, harness_filter: set[str] | None = None, seed_filter: set[int] | None = None, candidate_limit: int = 200, ) -> dict[str, Any]: """Run or resume every selected E04 composite robustness cell.""" code_revision = research_code_revision(root) experiment = load_experiments(root).get(experiment_id) if experiment is None or experiment.mode != "robustness": raise RobustnessExperimentError("runner requires the frozen E04 robustness experiment") 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] harness_catalog = load_harnesses(root) harnesses = [ harness_catalog[item] for item in experiment.harness_ids if harness_filter is None or item in harness_filter ] seeds = [seed for seed in experiment.seeds if seed_filter is None or seed in seed_filter] if not tasks or not harnesses or not seeds: raise RobustnessExperimentError("filters selected no E04 cells") if set(experiment.seeds) != set(DISTRACTOR_SEVERITIES): raise RobustnessExperimentError("E04 seeds no longer match the frozen dose mapping") model_spec = load_models(root)[experiment.model_ids[0]] embedding_spec = load_embeddings(root)[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 resident_models != (embedding_spec.model_key,): raise RobustnessExperimentError( f"E04 requires exclusive embedding-model residency; observed {resident_models}" ) snapshot = GitSnapshot(repository) origin = subprocess_git(repository, ["remote", "get-url", "origin"]) if {task.repository_url for task in tasks} != {origin}: raise RobustnessExperimentError("repository origin does not match frozen tasks") tokenizer = QwenTokenCounter() 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 tasks: parent_commit = subprocess_git(repository, ["rev-parse", f"{task.base_commit}^"]) base_chunks = chunk_snapshot( snapshot, task.base_commit, embedding_spec.chunk_lines, embedding_spec.chunk_overlap_lines, embedding_spec.chunk_char_limit, ) base_symbols = parse_snapshot(snapshot, task.base_commit) parent_chunks = chunk_snapshot( snapshot, parent_commit, embedding_spec.chunk_lines, embedding_spec.chunk_overlap_lines, embedding_spec.chunk_char_limit, ) parent_symbols = parse_snapshot(snapshot, parent_commit) base_index = _build_index( base_chunks, base_symbols, embedding_spec, embedding_client, embedding_cache ) stale_index = _build_index( parent_chunks, parent_symbols, embedding_spec, embedding_client, embedding_cache ) for harness in harnesses: query_record = ( load_iterative_query(root, task.task_id) if harness.harness_id == "H010" else { "query": task.statement, "query_source": "issue_statement", "protocol_violation": None, } ) query = query_record["query"] baseline_started = time.monotonic() baseline_candidates = _retrieve(harness, query, base_index, candidate_limit) baseline_query_seconds = time.monotonic() - baseline_started stale_started = time.monotonic() stale_candidates = _retrieve(harness, query, stale_index, candidate_limit) stale_query_seconds = time.monotonic() - stale_started baseline_metrics = _metrics( baseline_candidates, task, base_symbols, tokenizer, experiment.context_budgets[0] ) stale_metrics = _metrics( stale_candidates, task, base_symbols, tokenizer, experiment.context_budgets[0] ) stale_shift = robustness_shift( baseline_metrics, stale_metrics, [candidate.path for candidate in baseline_candidates], [candidate.path for candidate in stale_candidates], task.gold_files, candidate_limit + 1, ) for seed in seeds: 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=seed, 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 RobustnessExperimentError(f"incomplete pre-existing run: {directory}") rows.append(json.loads(final_path.read_text(encoding="utf-8"))) continue severity = distractor_severity(seed) sources = synthetic_distractor_sources(task, severity) distractor_chunks: list[SourceChunk] = list(base_chunks) distractor_symbols: list[GoSymbol] = list(base_symbols) for source in sources: distractor_chunks.extend( chunk_file( source, embedding_spec.chunk_lines, embedding_spec.chunk_overlap_lines, embedding_spec.chunk_char_limit, ) ) distractor_symbols.extend(parse_go_file(source.path, source.text)) distractor_index = _build_index( distractor_chunks, distractor_symbols, embedding_spec, embedding_client, embedding_cache, ) distractor_started = time.monotonic() distractor_candidates = _retrieve( harness, query, distractor_index, candidate_limit ) distractor_query_seconds = time.monotonic() - distractor_started distractor_metrics = _metrics( distractor_candidates, task, base_symbols, tokenizer, experiment.context_budgets[0], ) distractor_paths = [candidate.path for candidate in distractor_candidates] distractor_shift = robustness_shift( baseline_metrics, distractor_metrics, [candidate.path for candidate in baseline_candidates], distractor_paths, task.gold_files, candidate_limit + 1, ) synthetic_top_10 = [ path for path in distractor_paths[:10] if path.startswith(f"{DISTRACTOR_PREFIX}/") ] final = { "run_id": identity.run_id, "experiment_id": "E04", "task_id": task.task_id, "harness_id": harness.harness_id, "seed": seed, "query": query, "query_source": query_record["query_source"], "query_protocol_violation": query_record.get("protocol_violation"), "baseline": { "index_commit": task.base_commit, "metrics": baseline_metrics, "query_seconds": baseline_query_seconds, "index_stats": base_index[5], }, "stale_index": { "scenario_id": "S001", "index_commit": parent_commit, "evaluation_commit": task.base_commit, "shared_key": f"{task.task_id}:{harness.harness_id}:parent_commit", "shared_across_seeds": True, "metrics": stale_metrics, "shift": stale_shift, "query_seconds": stale_query_seconds, "index_stats": stale_index[5], }, "plausible_distractors": { "scenario_id": "S002", "severity": severity, "nested_dose": True, "synthetic_file_count": len(sources), "synthetic_in_top_10_count": len(synthetic_top_10), "synthetic_paths_in_top_10": synthetic_top_10, "metrics": distractor_metrics, "shift": distractor_shift, "query_seconds": distractor_query_seconds, "index_stats": distractor_index[5], }, } manifest = [ { "path": source.path, "sha256": sha256(source.text.encode("utf-8")).hexdigest(), "text": source.text, } for source in sources ] 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, "composite_scenarios": ["S001", "S002"], "stale_shared_across_seeds": True, "distractor_severity": severity, "candidate_limit": candidate_limit, "task_config_hash": task.config_hash, }, ) writer.emit("resource_sample", memory_sample()) writer.write_artifact( "baseline_ranking.json", json.dumps(_ranking(baseline_candidates, task.gold_files), indent=2) + "\n", ) writer.write_artifact( "stale_ranking.json", json.dumps(_ranking(stale_candidates, task.gold_files), indent=2) + "\n", ) writer.write_artifact( "distractor_ranking.json", json.dumps(_ranking(distractor_candidates, task.gold_files), indent=2) + "\n", ) writer.write_artifact( "distractor_sources.json", json.dumps(manifest, indent=2) + "\n" ) writer.write_artifact( "query_record.json", json.dumps(query_record, indent=2) + "\n" ) writer.write_artifact( "final_metrics.json", json.dumps(final, indent=2) + "\n" ) writer.emit("resource_sample", memory_sample()) writer.emit("run_finished", {"status": "completed", "metrics": final}) rows.append(final) summary = { "schema_version": 1, "experiment_id": experiment.experiment_id, "confirmatory": True, "code_revision": code_revision, "repository": str(repository.resolve()), "repository_origin": origin, "embedding_config_hash": embedding_spec.config_hash, "tokenizer_sha256": tokenizer.sha256, "resident_models": resident_models, "task_count": len(tasks), "harness_count": len(harnesses), "seed_count": len(seeds), "run_count": len(rows), "stale_unique_units": len(tasks) * len(harnesses), "stale_inference_rule": "deduplicate by stale_index.shared_key", "distractor_seed_dose_mapping": DISTRACTOR_SEVERITIES, "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