src/agent_harness/llm_localization.py

"""Blinded, single-call LLM localization over a frozen retrieval ranking."""

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 .lm_studio import LMStudioClient, LMStudioError
from .pilot import PilotError, research_code_revision, retrieval_metrics
from .repository import GitSnapshot
from .specs import load_experiments, load_harnesses, load_models, load_tasks
from .telemetry import EventWriter, RunIdentity


SYSTEM_PROMPT = """You are performing blinded bug localization in a large Go repository.
Use only the issue and candidate snippets supplied by the harness. Select the files that
would most likely need source-code changes. Do not propose a patch. Return one JSON object
with exactly these keys: {"files":["path/to/file.go"],"reasoning":"brief rationale"}.
The files array must contain 1-10 distinct paths copied exactly from the candidates. Keep
the rationale below 200 words and do not wrap the JSON in Markdown."""


class LocalizationError(RuntimeError):
    """Raised when a blinded localization run is invalid or cannot be completed."""


def load_ranking(path: Path, limit: int = 10) -> tuple[list[dict[str, Any]], str]:
    raw = path.read_bytes()
    try:
        value = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise LocalizationError(f"Invalid ranking JSON: {path}") from exc
    if not isinstance(value, list) or not value:
        raise LocalizationError("Ranking must be a non-empty JSON array")

    candidates: list[dict[str, Any]] = []
    seen: set[str] = set()
    for record in value:
        if not isinstance(record, dict):
            raise LocalizationError("Every ranking entry must be an object")
        try:
            candidate = {
                "rank": int(record["rank"]),
                "path": str(record["path"]),
                "line_start": int(record["line_start"]),
                "line_end": int(record["line_end"]),
                "score": float(record["score"]),
                "source": str(record["source"]),
            }
        except (KeyError, TypeError, ValueError) as exc:
            raise LocalizationError(f"Malformed ranking entry: {record!r}") from exc
        path_value = candidate["path"]
        if Path(path_value).is_absolute() or ".." in Path(path_value).parts:
            raise LocalizationError(f"Unsafe candidate path: {path_value}")
        if path_value not in seen:
            candidates.append(candidate)
            seen.add(path_value)
        if len(candidates) >= limit:
            break
    if not candidates:
        raise LocalizationError("Ranking has no usable candidate files")
    return candidates, sha256(raw).hexdigest()


def build_prompt(
    statement: str,
    snapshot: GitSnapshot,
    commit: str,
    candidates: Sequence[dict[str, Any]],
) -> str:
    blocks = [f"ISSUE:\n{statement}\n\nCANDIDATE SNIPPETS:"]
    for candidate in candidates:
        source = snapshot.read_file(commit, candidate["path"])
        lines = source.text.splitlines()
        start = max(candidate["line_start"], 1)
        end = min(candidate["line_end"], len(lines))
        numbered = "\n".join(
            f"{line_number:>6}: {lines[line_number - 1]}"
            for line_number in range(start, end + 1)
        )
        blocks.append(
            f"\n--- Candidate {candidate['rank']}: {candidate['path']} "
            f"(lines {start}-{end}) ---\n{numbered}"
        )
    return "\n".join(blocks)


def parse_selection(response: dict[str, Any], allowed_paths: set[str]) -> dict[str, Any]:
    try:
        message = response["choices"][0]["message"]
        content = message["content"]
    except (KeyError, IndexError, TypeError) as exc:
        raise LocalizationError("Chat completion has no assistant content") from exc
    if not isinstance(content, str):
        raise LocalizationError("Assistant content is not text")
    stripped = content.strip()
    if stripped.startswith("```"):
        stripped = stripped.removeprefix("```json").removeprefix("```")
        stripped = stripped.removesuffix("```").strip()
    try:
        value = json.loads(stripped)
    except json.JSONDecodeError:
        start, end = stripped.find("{"), stripped.rfind("}")
        if start < 0 or end <= start:
            raise LocalizationError(f"Assistant did not return JSON: {content!r}")
        try:
            value = json.loads(stripped[start : end + 1])
        except json.JSONDecodeError as exc:
            raise LocalizationError(f"Assistant returned invalid JSON: {content!r}") from exc
    if not isinstance(value, dict) or set(value) != {"files", "reasoning"}:
        raise LocalizationError("Assistant JSON must contain exactly files and reasoning")
    files = value["files"]
    if (
        not isinstance(files, list)
        or not 1 <= len(files) <= 10
        or not all(isinstance(item, str) for item in files)
        or len(files) != len(set(files))
    ):
        raise LocalizationError("Assistant files must be 1-10 distinct path strings")
    unknown = set(files) - allowed_paths
    if unknown:
        raise LocalizationError(f"Assistant selected paths outside the candidates: {sorted(unknown)}")
    if not isinstance(value["reasoning"], str):
        raise LocalizationError("Assistant reasoning must be text")
    return {"files": files, "reasoning": value["reasoning"]}


def _exclusive_agent_residency(discovery: Any, expected_key: str) -> tuple[str, ...]:
    loaded = tuple(
        str(record.get("key"))
        for record in discovery.native_models
        if record.get("loaded_instances")
    )
    if loaded != (expected_key,):
        raise LocalizationError(
            "LLM localization requires exclusive agent-model residency; "
            f"expected {(expected_key,)}, observed {loaded}"
        )
    return loaded


def run_localization(
    root: Path,
    repository: Path,
    ranking_path: Path,
    task_id: str,
    harness_id: str,
    experiment_id: str = "E06",
    candidate_limit: int = 10,
    timeout_seconds: float = 900.0,
) -> dict[str, Any]:
    revision = research_code_revision(root)
    experiments = load_experiments(root)
    harnesses = load_harnesses(root)
    models = load_models(root)
    tasks = load_tasks(root)
    try:
        experiment = experiments[experiment_id]
        harness = harnesses[harness_id]
        model = models[experiment.model_ids[0]]
        task = tasks[task_id]
    except KeyError as exc:
        raise LocalizationError(f"Unknown experiment, harness, model, or task: {exc}") from exc
    if harness_id not in experiment.harness_ids:
        raise LocalizationError(f"{harness_id} is not assigned to {experiment_id}")

    candidates, ranking_hash = load_ranking(ranking_path, candidate_limit)
    snapshot = GitSnapshot(repository)
    snapshot.verify_commit(task.base_commit)
    prompt = build_prompt(task.statement, snapshot, task.base_commit, candidates)
    prompt_hash = sha256(prompt.encode("utf-8")).hexdigest()

    client = LMStudioClient(model, timeout_seconds=timeout_seconds)
    discovery, resolved = client.resolve()
    loaded_models = _exclusive_agent_residency(discovery, resolved.inference_key)
    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.model_id,
        model_key=resolved.inference_key,
        model_config_hash=model.config_hash,
        context_budget=experiment.context_budgets[0],
        seed=model.seed,
        repetition=1,
        repository_sha=task.base_commit,
        code_revision=revision,
    )
    resolved_model = resolved.to_dict()
    resolved_model.update(
        {
            "exclusive_loaded_models": loaded_models,
            "ranking_sha256": ranking_hash,
            "prompt_sha256": prompt_hash,
            "candidate_limit": len(candidates),
        }
    )
    with EventWriter(
        root / "results",
        identity,
        asdict(harness),
        resolved_model,
    ) as writer:
        writer.emit(
            "run_started",
            {
                "development_only": True,
                "blinded_prompt": True,
                "ranking_path": str(ranking_path.resolve()),
                "prompt_chars": len(prompt),
                "candidate_count": len(candidates),
            },
        )
        started = time.monotonic()
        try:
            response = client.chat_completions(
                resolved.inference_key,
                [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=model.max_tokens,
            )
            elapsed = time.monotonic() - started
            writer.emit(
                "model_call",
                {
                    "elapsed_seconds": elapsed,
                    "usage": response.get("usage", {}),
                    "finish_reason": response.get("choices", [{}])[0].get("finish_reason"),
                },
            )
            selection = parse_selection(response, {item["path"] for item in candidates})
            metrics = retrieval_metrics(selection["files"], task.gold_files)
            final = {
                "run_id": identity.run_id,
                "experiment_id": experiment.experiment_id,
                "task_id": task.task_id,
                "harness_id": harness.harness_id,
                "selected_files": selection["files"],
                "reasoning": selection["reasoning"],
                "metrics": metrics,
                "usage": response.get("usage", {}),
                "elapsed_seconds": elapsed,
                "prompt_chars": len(prompt),
                "prompt_sha256": prompt_hash,
                "ranking_sha256": ranking_hash,
            }
            writer.write_artifact("model_response.json", json.dumps(response, indent=2) + "\n")
            writer.write_artifact("selection.json", json.dumps(selection, indent=2) + "\n")
            writer.write_artifact("final_metrics.json", json.dumps(final, indent=2) + "\n")
            writer.emit("run_finished", final)
            return {**final, "run_directory": str(writer.directory)}
        except (LMStudioError, LocalizationError) as exc:
            writer.emit("run_finished", {"status": "failed", "error": str(exc)})
            raise