src/agent_harness/backend_experiment.py

"""E05 dense-index backend systems experiment."""

from __future__ import annotations

from dataclasses import asdict
from hashlib import sha256
import json
from pathlib import Path
import statistics
import time
from typing import Any, Sequence

import faiss
import psutil

from .components import Candidate
from .confirmatory_retrieval import extended_metrics, subprocess_git
from .fusion import reciprocal_rank_fusion, unique_files
from .lm_studio_embeddings import LMStudioEmbeddingClient
from .pilot import research_code_revision
from .repository import GitSnapshot, chunk_snapshot
from .retrieval import BM25FuzzyRetriever, DenseRetriever, SQLiteEmbeddingCache
from .specs import (
    BackendSpec,
    HarnessSpec,
    load_backends,
    load_embeddings,
    load_experiments,
    load_harnesses,
    load_models,
    load_task_split,
    load_tasks,
)
from .syntax_index import SyntaxRetriever, parse_snapshot
from .telemetry import EventWriter, RunIdentity, run_directory
from .tokenization import QwenTokenCounter
from .vector_backends import FaissFlatRetriever, FaissHNSWRetriever, SQLiteVecRetriever


class BackendExperimentError(RuntimeError):
    """Raised when E05 cannot execute its frozen backend protocol."""


def percentile(values: Sequence[float], fraction: float) -> float:
    ordered = sorted(values)
    if not ordered:
        raise ValueError("percentile requires observations")
    position = (len(ordered) - 1) * fraction
    lower = int(position)
    upper = min(lower + 1, len(ordered) - 1)
    weight = position - lower
    return ordered[lower] * (1.0 - weight) + ordered[upper] * weight


def build_backend(
    backend: BackendSpec,
    dense: DenseRetriever,
    index_path: Path,
) -> tuple[Any, int, int]:
    process = psutil.Process()
    before = process.memory_info().rss
    if backend.backend_id == "B001":
        instance = FaissFlatRetriever(dense)
        faiss.write_index(instance.index, str(index_path))
    elif backend.backend_id == "B002":
        instance = FaissHNSWRetriever(
            dense,
            neighbors=int(backend.neighbors or 32),
            ef_construction=int(backend.ef_construction or 80),
            ef_search=int(backend.ef_search or 64),
        )
        faiss.write_index(instance.index, str(index_path))
    elif backend.backend_id == "B003":
        instance = SQLiteVecRetriever(dense, index_path)
    else:
        raise BackendExperimentError(f"unsupported backend {backend.backend_id}")
    after = process.memory_info().rss
    return instance, max(after - before, 0), index_path.stat().st_size


def treatment_ranking(
    harness: HarnessSpec,
    dense_ranking: Sequence[Candidate],
    lexical_ranking: Sequence[Candidate],
    syntax_ranking: Sequence[Candidate],
    limit: int,
) -> tuple[Candidate, ...]:
    if harness.harness_id == "H003":
        return unique_files(tuple(dense_ranking))[:limit]
    if harness.harness_id == "H005":
        return reciprocal_rank_fusion([lexical_ranking, dense_ranking], limit)
    if harness.harness_id == "H007":
        return reciprocal_rank_fusion([lexical_ranking, syntax_ranking, dense_ranking], limit)
    raise BackendExperimentError(f"E05 does not implement {harness.harness_id}")


def run_backend_experiment(
    root: Path,
    repository: Path,
    experiment_id: str = "E05",
    task_filter: set[str] | None = None,
    backend_filter: set[str] | None = None,
    harness_filter: set[str] | None = None,
    candidate_limit: int = 200,
) -> dict[str, Any]:
    revision = research_code_revision(root)
    experiments = load_experiments(root)
    experiment = experiments.get(experiment_id)
    if experiment is None or experiment.mode != "index_backend":
        raise BackendExperimentError("runner requires the frozen E05 index_backend experiment")
    harness_catalog = load_harnesses(root)
    backend_catalog = load_backends(root)
    model = load_models(root)[experiment.model_ids[0]]
    embedding = load_embeddings(root)[experiment.embedding_id]
    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]
    harnesses = [
        harness_catalog[item]
        for item in experiment.harness_ids
        if harness_filter is None or item in harness_filter
    ]
    backends = [
        backend_catalog[item]
        for item in experiment.backend_ids
        if backend_filter is None or item in backend_filter
    ]
    if not tasks or not harnesses or not backends:
        raise BackendExperimentError("filters selected no E05 cells")

    client = LMStudioEmbeddingClient(embedding, timeout_seconds=120.0)
    runtime = client.resolve()
    resident = client.loaded_model_keys()
    if tuple(resident) != (embedding.model_key,):
        raise BackendExperimentError(f"E05 requires exclusive embedding residency; observed {resident}")
    tokenizer = QwenTokenCounter()
    snapshot = GitSnapshot(repository)
    origin = subprocess_git(repository, ["remote", "get-url", "origin"])
    rows: list[dict[str, Any]] = []
    cache_path = root / "indexes" / "embeddings" / f"{embedding.config_hash}.sqlite3"
    with SQLiteEmbeddingCache(cache_path, embedding) as cache:
        for task in tasks:
            chunks = chunk_snapshot(
                snapshot,
                task.base_commit,
                embedding.chunk_lines,
                embedding.chunk_overlap_lines,
                embedding.chunk_char_limit,
            )
            symbols = parse_snapshot(snapshot, task.base_commit)
            dense_base, dense_stats = DenseRetriever.build(chunks, embedding, client, cache)
            lexical_ranking = BM25FuzzyRetriever(chunks).retrieve(task.statement, candidate_limit)
            syntax_ranking = SyntaxRetriever(symbols).retrieve(task.statement, candidate_limit)
            index_dir = root / "indexes" / "e05" / task.base_commit
            index_dir.mkdir(parents=True, exist_ok=True)
            instances: dict[str, tuple[Any, int, int]] = {}
            for backend in backends:
                suffix = ".sqlite3" if backend.backend_id == "B003" else ".faiss"
                instances[backend.backend_id] = build_backend(
                    backend,
                    dense_base,
                    index_dir / f"{backend.backend_id}{suffix}",
                )
            flat_paths = [
                candidate.path
                for candidate in unique_files(
                    tuple(instances["B001"][0].retrieve(task.statement, candidate_limit))
                )[:10]
            ] if "B001" in instances else []

            for backend in backends:
                instance, ram_delta, disk_bytes = instances[backend.backend_id]
                for seed in experiment.seeds:
                    timings: list[float] = []
                    dense_ranking: Sequence[Candidate] = ()
                    for _ in range(backend.query_repetitions):
                        started = time.perf_counter()
                        dense_ranking = instance.retrieve(task.statement, candidate_limit)
                        timings.append((time.perf_counter() - started) * 1000.0)
                    for harness in harnesses:
                        treatment_id = f"{harness.harness_id}_{backend.backend_id}"
                        treatment_hash = sha256(
                            f"{harness.config_hash}\0{backend.config_hash}".encode("utf-8")
                        ).hexdigest()
                        identity = RunIdentity(
                            experiment_id=experiment.experiment_id,
                            task_id=task.task_id,
                            harness_id=treatment_id,
                            harness_hash=treatment_hash,
                            model_id=model.model_id,
                            model_key=model.expected_inference_key,
                            model_config_hash=model.config_hash,
                            context_budget=experiment.context_budgets[0],
                            seed=seed,
                            repetition=0,
                            repository_sha=task.base_commit,
                            code_revision=revision,
                        )
                        directory = run_directory(root / "results", identity)
                        if directory.exists():
                            final_path = directory / "final_metrics.json"
                            if not final_path.exists():
                                raise BackendExperimentError(f"incomplete E05 run: {directory}")
                            final = json.loads(final_path.read_text(encoding="utf-8"))
                            rows.append({"run_id": identity.run_id, **final})
                            continue
                        ranking = treatment_ranking(
                            harness,
                            dense_ranking,
                            lexical_ranking,
                            syntax_ranking,
                            candidate_limit,
                        )
                        metrics = extended_metrics(
                            ranking,
                            task.gold_files,
                            task.gold_symbols,
                            symbols,
                            tokenizer,
                            experiment.context_budgets[0],
                        )
                        backend_top = [candidate.path for candidate in unique_files(tuple(dense_ranking))[:10]]
                        metrics.update(
                            {
                                "experiment_id": experiment.experiment_id,
                                "task_id": task.task_id,
                                "harness_id": harness.harness_id,
                                "backend_id": backend.backend_id,
                                "seed": seed,
                                "backend_recall_at_10_vs_flat": (
                                    len(set(backend_top) & set(flat_paths)) / 10.0 if flat_paths else None
                                ),
                                "index_build_seconds": instance.stats.build_seconds,
                                "index_ram_bytes_delta": ram_delta,
                                "index_disk_bytes": disk_bytes,
                                "query_repetitions": backend.query_repetitions,
                                "query_mean_ms": statistics.fmean(timings),
                                "query_p50_ms": percentile(timings, 0.50),
                                "query_p95_ms": percentile(timings, 0.95),
                                "dense_cached_chunks": dense_stats.cached_chunks,
                                "dense_embedded_chunks": dense_stats.embedded_chunks,
                            }
                        )
                        with EventWriter(
                            root / "results",
                            identity,
                            {"harness": asdict(harness), "backend": asdict(backend)},
                            {
                                "agent_model_not_loaded": asdict(model),
                                "embedding_model": asdict(embedding),
                                "embedding_runtime": runtime,
                            },
                        ) as writer:
                            writer.emit("run_started", {"confirmatory": True, "candidate_limit": candidate_limit})
                            for rank, candidate in enumerate(ranking, 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 set(task.gold_files),
                                    },
                                )
                            writer.write_artifact(
                                "final_metrics.json", json.dumps(metrics, indent=2) + "\n"
                            )
                            writer.emit("run_finished", {"status": "completed", "metrics": metrics})
                        rows.append({"run_id": identity.run_id, **metrics})
            for instance, _, _ in instances.values():
                if isinstance(instance, SQLiteVecRetriever):
                    instance.close()

    summary = {
        "schema_version": 1,
        "experiment_id": experiment.experiment_id,
        "confirmatory": True,
        "repository_origin": origin,
        "code_revision": revision,
        "task_count": len(tasks),
        "harness_count": len(harnesses),
        "backend_count": len(backends),
        "seed_count": len(experiment.seeds),
        "run_count": len(rows),
        "rows": rows,
    }
    report = root / "results" / "reports" / f"E05_{revision[:12]}_{int(time.time())}.json"
    report.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    summary["report_path"] = str(report)
    return summary