src/agent_harness/confirmatory_retrieval.py

"""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()