src/agent_harness/pilot.py

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