judge.py

"""Build and preview one-record-at-a-time inputs for the LLM judge.

The default command does not call an LLM. It loads records from
`subset_dataset.json`, applies URL preprocessing if needed, and prints the JSON
payload that should be sent as the judge's user message.

Example:
    python judge.py --preview --limit 1
"""

from __future__ import annotations

import argparse
import json
import traceback
from pathlib import Path
from typing import Any, Iterator

from pydantic import BaseModel, Field
from pydantic_ai import Agent

from schemas import DocumentationJudgeOutput
from tools.process_urls import process_record


PROJECT_ROOT = Path(__file__).resolve().parent
DEFAULT_SUBSET_PATH = PROJECT_ROOT / "subset_dataset.json"
DEFAULT_SYSTEM_PROMPT_PATH = PROJECT_ROOT / "prompts" / "judge_system.md"
DEFAULT_MODEL = "openai:gpt-4o"


class JudgeSource(BaseModel):
    title: str | None = Field(
        default=None,
        description="Human-readable title of the source document, if available.",
    )
    raw_url: str | None = Field(
        default=None,
        description="The original source URL or partial URL from the dataset.",
    )
    canonical_url: str | None = Field(
        default=None,
        description="The source URL after canonicalizing partial docs paths.",
    )
    normalized_url: str | None = Field(
        default=None,
        description="The source URL normalized for deduping and comparison.",
    )
    evaluation_url: str | None = Field(
        default=None,
        description=(
            "The URL actually fetched for evaluation evidence. For localized docs, "
            "this may be the English equivalent of the original source URL."
        ),
    )
    source_locale: str | None = Field(
        default=None,
        description=(
            "Locale inferred from the original source URL, such as en, ja, ko, or fr."
        ),
    )
    locale_rewritten_for_evaluation: bool | None = Field(
        default=None,
        description=(
            "Whether the original localized source URL was rewritten to an English "
            "evaluation URL before fetching page text."
        ),
    )
    original_canonical_url: str | None = Field(
        default=None,
        description=(
            "The canonicalized user-facing source URL before any locale rewrite for "
            "evaluation."
        ),
    )
    content_excerpt: str | None = Field(
        default=None,
        description=(
            "Optional retrieved source text shown to the support bot. This may "
            "be absent in the current subset."
        ),
    )


class JudgeResponseLink(BaseModel):
    raw_url: str | None = Field(
        default=None,
        description="The link exactly as it appeared in the bot response.",
    )
    canonical_url: str | None = Field(
        default=None,
        description="The response link after canonicalizing partial docs paths.",
    )
    normalized_url: str | None = Field(
        default=None,
        description="The response link normalized for deduping and comparison.",
    )


class DocumentationJudgeInput(BaseModel):
    unique_id: str = Field(description="Unique ID of the answer record.")
    timestamp: str | None = Field(
        default=None,
        description="Timestamp of the original support interaction.",
    )
    query: str = Field(description="The user's question.")
    response: str = Field(description="The support bot's answer to evaluate.")
    queryCategory: str | None = Field(
        default=None,
        description="Existing query category label from the dataset.",
    )
    sources: list[JudgeSource] = Field(
        description="Documentation sources the support bot used, after URL normalization.",
    )
    response_links: list[JudgeResponseLink] = Field(
        description="Documentation links extracted from the support bot response.",
    )
    source_url_validity: dict[str, int] = Field(
        description="Code-computed URL hygiene metrics for the supplied sources.",
    )
    response_link_source_match: str = Field(
        description=(
            "Code-computed relationship between response links and supplied "
            "sources after URL canonicalization."
        ),
    )


def load_records(path: Path) -> list[dict[str, Any]]:
    """Load a JSON array of records from the subset dataset."""
    with path.open(encoding="utf-8") as file:
        data = json.load(file)

    if not isinstance(data, list):
        raise ValueError(f"Expected {path} to contain a JSON array.")

    return data


def ensure_url_processed(record: dict[str, Any]) -> dict[str, Any]:
    """Return a record with normalized URL fields present."""
    required_fields = {
        "normalized_sources",
        "response_links",
        "source_url_validity",
        "response_link_source_match",
    }
    if required_fields.issubset(record):
        return record
    return process_record(record)


def build_judge_input(record: dict[str, Any]) -> DocumentationJudgeInput:
    """Build the exact per-record payload to send to the LLM judge."""
    processed = ensure_url_processed(record)

    sources = [
        JudgeSource(
            title=source.get("title"),
            raw_url=source.get("raw_url") or source.get("url"),
            canonical_url=source.get("canonical_url"),
            normalized_url=source.get("normalized_url"),
            evaluation_url=source.get("evaluation_url"),
            source_locale=source.get("source_locale"),
            locale_rewritten_for_evaluation=source.get("locale_rewritten_for_evaluation"),
            original_canonical_url=source.get("original_canonical_url"),
            content_excerpt=source.get("content_excerpt") or source.get("excerpt"),
        )
        for source in processed.get("normalized_sources", [])
    ]

    response_links = [
        JudgeResponseLink(
            raw_url=link.get("raw_url"),
            canonical_url=link.get("canonical_url"),
            normalized_url=link.get("normalized_url"),
        )
        for link in processed.get("response_links", [])
    ]

    return DocumentationJudgeInput(
        unique_id=processed["unique_id"],
        timestamp=processed.get("timestamp"),
        query=processed["query"],
        response=processed["response"],
        queryCategory=processed.get("queryCategory"),
        sources=sources,
        response_links=response_links,
        source_url_validity=processed["source_url_validity"],
        response_link_source_match=processed["response_link_source_match"],
    )


def iter_judge_inputs(records: list[dict[str, Any]]) -> Iterator[DocumentationJudgeInput]:
    """Yield DocumentationJudgeInput objects for each record."""
    for record in records:
        yield build_judge_input(record)


def build_user_message(judge_input: DocumentationJudgeInput) -> str:
    """Serialize the judge input as the user message for the LLM judge."""
    payload = judge_input.model_dump(mode="json", exclude_none=True)
    return (
        "Evaluate this documentation-agent answer. Return only a structured "
        "DocumentationJudgeOutput object matching the provided Pydantic schema.\n\n"
        "Do not use resolutionStatus; it has intentionally been omitted.\n\n"
        f"{json.dumps(payload, ensure_ascii=False, indent=2)}"
    )


def load_system_prompt(path: Path = DEFAULT_SYSTEM_PROMPT_PATH) -> str:
    return path.read_text(encoding="utf-8")


def create_judge_agent(model: str):
    """Create the PydanticAI judge agent."""
    return Agent(
        model,
        output_type=DocumentationJudgeOutput,
        system_prompt=load_system_prompt(),
    )


async def judge_one_record(model: str, judge_input: DocumentationJudgeInput) -> DocumentationJudgeOutput:
    """Run the LLM judge for one record and return validated output."""
    agent = create_judge_agent(model)
    result = await agent.run(build_user_message(judge_input))
    return result.output


def judge_one_record_sync(model: str, judge_input: DocumentationJudgeInput) -> DocumentationJudgeOutput:
    """Run the LLM judge synchronously for one record and return validated output."""
    agent = create_judge_agent(model)
    result = agent.run_sync(build_user_message(judge_input))
    return result.output


def write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as file:
        for row in rows:
            file.write(json.dumps(row, ensure_ascii=False, sort_keys=True))
            file.write("\n")


def append_jsonl(path: Path, row: dict[str, Any]) -> None:
    """Append one JSON object and flush it so completed evals survive crashes."""
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as file:
        file.write(json.dumps(row, ensure_ascii=False, sort_keys=True))
        file.write("\n")
        file.flush()


def load_completed_ids(path: Path | None) -> set[str]:
    """Load successfully judged IDs from an existing JSONL output file."""
    completed: set[str] = set()
    if path is None or not path.exists():
        return completed

    with path.open(encoding="utf-8") as file:
        for line in file:
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                continue
            unique_id = row.get("unique_id")
            if isinstance(unique_id, str):
                completed.add(unique_id)

    return completed


def run_llm_judge(
    records: list[dict[str, Any]],
    *,
    model: str,
    limit: int,
    output_path: Path | None,
    errors_output_path: Path | None,
    resume: bool,
    continue_on_error: bool,
) -> list[DocumentationJudgeOutput]:
    """Run the judge over records one at a time."""
    judgments: list[DocumentationJudgeOutput] = []
    completed_ids = load_completed_ids(output_path) if resume else set()

    if output_path is not None and not resume:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        output_path.write_text("", encoding="utf-8")

    if errors_output_path is not None and not resume:
        errors_output_path.parent.mkdir(parents=True, exist_ok=True)
        errors_output_path.write_text("", encoding="utf-8")

    for index, judge_input in enumerate(iter_judge_inputs(records)):
        if index >= limit:
            break

        if judge_input.unique_id in completed_ids:
            print(
                json.dumps(
                    {
                        "event": "skip_completed",
                        "index": index,
                        "unique_id": judge_input.unique_id,
                    },
                    ensure_ascii=False,
                    sort_keys=True,
                )
            )
            continue

        try:
            judgment = judge_one_record_sync(model, judge_input)
        except Exception as error:
            error_row = {
                "event": "judge_error",
                "index": index,
                "unique_id": judge_input.unique_id,
                "query": judge_input.query,
                "error_type": type(error).__name__,
                "error": str(error),
                "traceback": traceback.format_exc(),
            }
            if errors_output_path is not None:
                append_jsonl(errors_output_path, error_row)
            print(json.dumps(error_row, ensure_ascii=False, sort_keys=True))
            if continue_on_error:
                continue
            raise

        judgment_row = judgment.model_dump(mode="json")
        judgments.append(judgment)
        if output_path is not None:
            append_jsonl(output_path, judgment_row)
        print(json.dumps(judgment_row, ensure_ascii=False, sort_keys=True))

    return judgments


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--dataset",
        type=Path,
        default=DEFAULT_SUBSET_PATH,
        help="Path to a JSON array of records to evaluate.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        default=1,
        help="Number of records to preview.",
    )
    parser.add_argument(
        "--preview",
        action="store_true",
        help="Print judge input user messages. This is the default behavior.",
    )
    parser.add_argument(
        "--schema",
        action="store_true",
        help="Print the expected DocumentationJudgeOutput JSON schema.",
    )
    parser.add_argument(
        "--run-llm",
        action="store_true",
        help="Call the PydanticAI judge instead of only previewing judge inputs.",
    )
    parser.add_argument(
        "--model",
        default=DEFAULT_MODEL,
        help=(
            "PydanticAI model string for --run-llm. Defaults to "
            f"{DEFAULT_MODEL}. Use openai:gpt-4 if you want that exact model."
        ),
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=None,
        help="Optional JSONL path for validated judge outputs from --run-llm.",
    )
    parser.add_argument(
        "--errors-output",
        type=Path,
        default=None,
        help=(
            "Optional JSONL path for per-record judge failures. Defaults to "
            "<output_stem>_errors.jsonl when --output is provided."
        ),
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Skip records whose unique_id already appears in --output.",
    )
    parser.add_argument(
        "--continue-on-error",
        action="store_true",
        help="Log per-record judge failures and continue with later records.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    if args.schema:
        print(json.dumps(DocumentationJudgeOutput.model_json_schema(), indent=2))
        return

    records = load_records(args.dataset)
    if args.run_llm:
        errors_output = args.errors_output
        if errors_output is None and args.output is not None:
            errors_output = args.output.with_name(f"{args.output.stem}_errors.jsonl")

        run_llm_judge(
            records,
            model=args.model,
            limit=args.limit,
            output_path=args.output,
            errors_output_path=errors_output,
            resume=args.resume,
            continue_on_error=args.continue_on_error,
        )
        return

    for index, judge_input in enumerate(iter_judge_inputs(records)):
        if index >= args.limit:
            break
        print(build_user_message(judge_input))
        if index < args.limit - 1:
            print("\n" + "=" * 80 + "\n")


if __name__ == "__main__":
    main()