siemens_ragas/scripts/build_judge_prompt_cache.py

"""One-off bootstrap: generate committed Chinese judge-prompt cache via RAGAS adapt().

Usage:
    python -m scripts.build_judge_prompt_cache --language zh [--judge-model <model>]

Runs BasePrompt.adapt("chinese", llm, adapt_instruction=True) for every judge
prompt of every LLM-scored metric and writes the result to
configs/judge_prompts/<language>/<metric>__<attr>.json.

Re-run after a RAGAS upgrade to refresh the cache (the localizer will warn
via 'stale cache' log messages when the source hash changes).
Requires a working judge LLM configured in settings / LLM profiles.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
from pathlib import Path
from typing import Any

from rag_eval.compat import ensure_ragas_import_compat
from rag_eval.metrics.factory import build_metric_registry, build_models
from rag_eval.metrics.judge_prompts import CACHE_ROOT, METRIC_PROMPT_ATTRS, prompt_source_hash
from rag_eval.settings import EvaluationSettings

ensure_ragas_import_compat()

logger = logging.getLogger("scripts.build_judge_prompt_cache")

# Map from judge_language short code to the natural-language name used by RAGAS adapt().
_ADAPT_LANGUAGE: dict[str, str] = {
    "zh": "chinese",
    "en": "english",
}


def serialize_prompt(
    adapted: Any,
    source_hash: str,
    metric: str,
    attr: str,
    language: str,
    ragas_version: str,
) -> dict:
    """Serialize an adapted prompt instance into the committed cache JSON schema."""
    return {
        "metric": metric,
        "prompt_attr": attr,
        "language": getattr(adapted, "language", language),
        "ragas_version": ragas_version,
        "source_hash": source_hash,
        "instruction": adapted.instruction,
        "examples": [
            {"input": inp.model_dump(), "output": out.model_dump()}
            for inp, out in adapted.examples
        ],
    }


async def build_cache(
    language: str,
    judge_model: str,
    embedding_model: str,
    settings: Any,
) -> list[Path]:
    """Generate and write the full prompt cache for one language; return written paths.

    Each prompt is adapted once, serialized, and written atomically (temp file rename
    is not used — if the script crashes mid-way, re-running overwrites partial files).
    """
    import ragas

    target = _ADAPT_LANGUAGE.get(language, language)
    llm, _embeddings = build_models(judge_model, embedding_model, settings)
    registry = build_metric_registry(llm, _embeddings)
    out_dir = CACHE_ROOT / language
    out_dir.mkdir(parents=True, exist_ok=True)

    written: list[Path] = []
    for metric_name, attrs in METRIC_PROMPT_ATTRS.items():
        for attr in attrs:
            metric = registry[metric_name]
            prompt = getattr(metric, attr)
            src_hash = prompt_source_hash(prompt)
            logger.info("adapting %s.%s …", metric_name, attr)
            adapted = await prompt.adapt(target, llm, adapt_instruction=True)
            data = serialize_prompt(adapted, src_hash, metric_name, attr, language, ragas.__version__)
            path = out_dir / f"{metric_name}__{attr}.json"
            path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
            written.append(path)
            logger.info("  written → %s", path.name)

    return written


def main() -> None:
    """CLI entry point: parse args, resolve models, and build the cache."""
    parser = argparse.ArgumentParser(
        description="Build the RAGAS judge-prompt translation cache for a target language."
    )
    parser.add_argument("--language", default="zh", help="Target language code (default: zh)")
    parser.add_argument("--judge-model", default=None, help="Override judge model from settings")
    parser.add_argument("--embedding-model", default=None, help="Override embedding model")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s  %(message)s")
    settings = EvaluationSettings()
    judge_model = args.judge_model or settings.ragas_judge_model
    embedding_model = args.embedding_model or settings.ragas_embedding_model

    paths = asyncio.run(build_cache(args.language, judge_model, embedding_model, settings))
    logger.info("Done: %d cache files written to configs/judge_prompts/%s/", len(paths), args.language)


if __name__ == "__main__":
    main()