open-swe/evals/reviewer/target.py

"""Target function for the reviewer eval.

Spawns the reviewer graph over `langgraph_sdk` for one PR, waits for
completion, and returns every `add_finding` tool call the agent made as the
structured output for the eval. Findings are normalized into the legacy
``{file, line, body, severity}`` shape so the judge prompt can stay the
verbatim form martian published.
"""

from __future__ import annotations

import os
import threading
from typing import Any, Literal, cast

from dotenv import load_dotenv
from langgraph_sdk import get_client

from agent.reviewer_findings import Finding, Severity, filter_findings_for_publish

load_dotenv()

DEFAULT_REVIEWER_ASSISTANT_ID = "reviewer"
DEFAULT_LANGGRAPH_URL = "http://localhost:2024"
ScoreMode = Literal["all_findings", "surfaced_findings"]
_VALID_SCORE_MODES: set[ScoreMode] = {"all_findings", "surfaced_findings"}
_VALID_SEVERITIES: set[Severity] = {"low", "medium", "high", "critical"}

_THREAD_IDS: set[str] = set()
_THREAD_IDS_LOCK = threading.Lock()


def _record_thread_id(thread_id: str) -> None:
    with _THREAD_IDS_LOCK:
        _THREAD_IDS.add(thread_id)


def drain_thread_ids() -> set[str]:
    """Return and clear thread IDs created by ``review_pr`` so far.

    Used by ``run_eval`` to delete threads after the experiment finishes.
    Underlying provider sandboxes time out via their own TTL — deleting the
    LangGraph thread frees the checkpoint/metadata records, not the sandbox.
    """
    with _THREAD_IDS_LOCK:
        snapshot = set(_THREAD_IDS)
        _THREAD_IDS.clear()
    return snapshot


def get_langgraph_url() -> str:
    return os.getenv("LANGGRAPH_URL", DEFAULT_LANGGRAPH_URL)


def get_reviewer_assistant_id() -> str:
    return os.getenv("REVIEWER_ASSISTANT_ID", DEFAULT_REVIEWER_ASSISTANT_ID)


def get_score_mode() -> ScoreMode:
    value = os.getenv("REVIEWER_EVAL_SCORE_MODE", "all_findings")
    if value in _VALID_SCORE_MODES:
        return cast(ScoreMode, value)
    return "all_findings"


def get_reviewer_model_id() -> str | None:
    value = os.getenv("REVIEWER_EVAL_MODEL_ID")
    return value if value else None


def get_reviewer_reasoning_effort() -> str | None:
    value = os.getenv("REVIEWER_EVAL_REASONING_EFFORT")
    return value if value else None


def _build_user_message(inputs: dict[str, Any]) -> str:
    return (
        f"Review pull request {inputs['pr_url']}.\n\n"
        f"- repo: {inputs['repo']}\n"
        f"- pr_number: {inputs['pr_number']}\n"
        f"- title: {inputs.get('pr_title', '')}\n"
        f"- base_sha: {inputs['base_sha']}\n"
        f"- head_sha: {inputs['head_sha']}\n"
        f"- base_ref: {inputs.get('base_ref', '')}\n"
        f"- head_ref: {inputs.get('head_ref', '')}\n\n"
        f"Record each issue you find with the `add_finding` tool, then call "
        f"`publish_review` once at the end."
    )


def _build_configurable(inputs: dict[str, Any]) -> dict[str, Any]:
    repo = inputs.get("repo", "")
    owner, _, name = repo.partition("/") if isinstance(repo, str) else ("", "", "")
    configurable: dict[str, Any] = {
        "__is_for_execution__": True,
        "reviewer_eval": True,
        "eval": True,
        "repo": {"owner": owner, "name": name},
        "pr_number": inputs.get("pr_number"),
        "pr_url": inputs.get("pr_url", ""),
        "base_sha": inputs.get("base_sha", ""),
        "head_sha": inputs.get("head_sha", ""),
        "branch_name": inputs.get("head_ref", ""),
    }
    model_id = get_reviewer_model_id()
    if model_id:
        configurable["reviewer_model_id"] = model_id
    reasoning_effort = get_reviewer_reasoning_effort()
    if reasoning_effort:
        configurable["reviewer_reasoning_effort"] = reasoning_effort
    return configurable


async def review_pr(inputs: dict[str, Any]) -> dict[str, Any]:
    """LangSmith target: run the reviewer agent on one PR."""
    client = get_client(url=get_langgraph_url())
    thread = await client.threads.create()
    thread_id: str = thread["thread_id"]
    _record_thread_id(thread_id)
    result = await client.runs.wait(
        thread_id,
        assistant_id=get_reviewer_assistant_id(),
        input={"messages": [{"role": "user", "content": _build_user_message(inputs)}]},
        config={"configurable": _build_configurable(inputs)},
    )
    if get_score_mode() == "surfaced_findings":
        return {"comments": await _extract_surfaced_comments(client, thread_id)}
    return {"comments": _extract_comments(result)}


def _extract_comments(result: Any) -> list[dict[str, Any]]:
    """Collect every ``add_finding`` tool call from the run's message stream.

    Normalizes the new finding shape (``start_line``/``end_line``/``description``)
    into the legacy ``{file, line, body, severity}`` shape the judge prompt
    consumes verbatim from martian's benchmark.
    """
    if not isinstance(result, dict):
        return []
    comments: list[dict[str, Any]] = []
    for msg in result.get("messages") or []:
        if not isinstance(msg, dict):
            continue
        for tc in msg.get("tool_calls") or []:
            if tc.get("name") != "add_finding":
                continue
            args = tc.get("args") or {}
            file = args.get("file")
            severity = args.get("severity")
            description = args.get("description") or args.get("body") or ""
            line = args.get("end_line")
            if line is None:
                line = args.get("start_line")
            if not file or not severity:
                continue
            comments.append(
                {
                    "file": file,
                    "line": line,
                    "body": description,
                    "severity": severity,
                }
            )
    return comments


async def _extract_surfaced_comments(client: Any, thread_id: str) -> list[dict[str, Any]]:
    thread = await client.threads.get(thread_id)
    metadata = thread.get("metadata") if isinstance(thread, dict) else None
    findings_value = metadata.get("findings") if isinstance(metadata, dict) else None
    findings = _coerce_findings(findings_value)
    surfaced = filter_findings_for_publish(
        findings,
        severity_threshold=_score_severity_threshold(),
        cap=_score_cap(),
    )
    return [_normalize_finding(finding) for finding in surfaced]


def _coerce_findings(value: Any) -> list[Finding]:
    if not isinstance(value, list):
        return []
    findings: list[Finding] = []
    for item in value:
        if not isinstance(item, dict):
            continue
        if not isinstance(item.get("id"), str):
            continue
        findings.append(cast(Finding, item))
    return findings


def _normalize_finding(finding: Finding) -> dict[str, Any]:
    line = finding.get("end_line")
    if line is None:
        line = finding.get("start_line")
    return {
        "file": finding.get("file"),
        "line": line,
        "body": finding.get("description", ""),
        "severity": finding.get("severity"),
    }


def _score_severity_threshold() -> Severity:
    value = os.getenv("REVIEWER_EVAL_SEVERITY_THRESHOLD", "medium")
    if value in _VALID_SEVERITIES:
        return cast(Severity, value)
    return "medium"


def _score_cap() -> int:
    raw = os.getenv("REVIEWER_EVAL_CAP", "4")
    try:
        cap = int(raw)
    except ValueError:
        return 4
    return max(cap, 0)