office_translator/services/quality/llm_judge.py

"""
L1 LLM Judge — uses a cheap LLM via the OpenAI-compatible API to validate
the quality of a translation.

Why a SEPARATE LLM (not the one that did the translation)?
  - Self-evaluation bias: a model tends to rate its own output as good
    (Meta/Stanford papers on LLM-as-judge bias, 2024-2025).
  - Independence gives a more reliable signal.

Design constraints:
  - 100% API-based. No local models, no GPU.
  - Async, with a hard timeout.
  - Defensive: never raises. Returns a "skip" verdict on any internal error.
  - Output is structured JSON for reliable parsing.
  - Sampled (we never judge all chunks, just a small subset).

The LLM is asked a binary question per chunk: "Is this translation
accurate and fluent, in the correct language?" with a short reason.
We do NOT ask for nuanced MQM scores — binary is more reliable and
cheaper.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
import time
from dataclasses import dataclass, field, asdict
from typing import List, Optional, Tuple

from core.logging import get_logger

logger = get_logger(__name__)


# ---------- Result dataclasses ----------

@dataclass
class L1ChunkVerdict:
    """Verdict for a single (source, translation) pair."""
    accurate: bool
    fluent: bool
    correct_language: bool
    no_leaks: bool
    reason: str = ""

    @property
    def passed(self) -> bool:
        return self.accurate and self.fluent and self.correct_language and self.no_leaks

    def to_log_dict(self) -> dict:
        return asdict(self)


@dataclass
class L1Result:
    """Aggregate result of an L1 check on a sample of chunks."""
    verdict: str                  # "pass", "fail", "skip"
    chunks_evaluated: int
    chunks_passed: int
    chunks_failed: int
    failure_rate: float           # 0.0 to 1.0
    samples: List[dict] = field(default_factory=list)
    model_used: str = ""
    elapsed_ms: float = 0.0
    cost_estimate_usd: float = 0.0
    error: str = ""

    def to_log_dict(self) -> dict:
        return asdict(self)


# ---------- Prompt template ----------

JUDGE_SYSTEM_PROMPT = """You are a strict translation quality evaluator. Your job is to detect translation failures.

For each (SOURCE, TRANSLATION) pair, check these 4 criteria:

1. ACCURATE — Does the translation preserve the meaning of the source? (yes/no)
2. FLUENT — Is the translation grammatically correct in the target language? (yes/no)
3. CORRECT_LANGUAGE — Is the translation actually in {target_lang_name} (ISO code: {target_lang})? (yes/no)
4. NO_LEAKS — Is the translation free of prompt artifacts, source-language text, or meta-commentary? (yes/no)

A translation FAILS if ANY of the 4 checks is "no".

Respond with a JSON array, one object per pair, in the same order. NO other text, NO markdown fences:
[
  {{"accurate": "yes"|"no", "fluent": "yes"|"no", "correct_language": "yes"|"no", "no_leaks": "yes"|"no", "reason": "one short sentence"}}
]

Be honest: if the translation looks suspicious, say "no". The "reason" must be in English and ≤ 12 words.
"""


# ---------- LLM client ----------

class LLMJudge:
    """
    Calls a cheap LLM via the OpenAI-compatible API to judge translation quality.

    Default configuration targets deepseek-chat via the OpenAI-compatible
    DeepSeek API. The judge can be reconfigured to use any OpenAI-compatible
    endpoint (OpenAI, OpenRouter, etc.) by passing different params.
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = "https://api.deepseek.com/v1",
        model: str = "deepseek-chat",
        timeout_seconds: float = 8.0,
        max_retries: int = 1,
    ):
        if not api_key:
            raise ValueError("api_key is required for LLMJudge")
        self._api_key = api_key
        self._base_url = base_url.rstrip("/")
        self._model = model
        self._timeout = timeout_seconds
        self._max_retries = max_retries
        # Lazy import — keep services/quality free of the openai dep at import time
        self._client = None

    def _get_client(self):
        if self._client is None:
            try:
                import openai
                self._client = openai.AsyncOpenAI(
                    api_key=self._api_key,
                    base_url=self._base_url,
                    timeout=self._timeout,
                )
            except Exception as e:
                logger.warning("llm_judge_client_init_failed", error=str(e))
                return None
        return self._client

    async def judge_batch(
        self,
        pairs: List[Tuple[str, str]],
        target_lang: str,
        target_lang_name: str = "",
    ) -> L1Result:
        """
        Judge a batch of (source, translation) pairs.

        Returns an L1Result with verdict="skip" on any internal error.
        Never raises.
        """
        start = time.time()
        if not pairs:
            return L1Result(verdict="skip", chunks_evaluated=0, chunks_passed=0,
                           chunks_failed=0, failure_rate=0.0,
                           error="empty pairs")

        client = self._get_client()
        if client is None:
            return L1Result(verdict="skip", chunks_evaluated=0, chunks_passed=0,
                           chunks_failed=0, failure_rate=0.0,
                           error="client unavailable")

        # Build the user message
        user_lines = [f"Target language: {target_lang} ({target_lang_name})\n"]
        for i, (src, trans) in enumerate(pairs, 1):
            user_lines.append(f"\n--- Pair {i} ---")
            user_lines.append(f"SOURCE: {src}")
            user_lines.append(f"TRANSLATION: {trans}")
        user_msg = "\n".join(user_lines)

        system_prompt = JUDGE_SYSTEM_PROMPT.format(
            target_lang=target_lang,
            target_lang_name=target_lang_name or target_lang,
        )

        try:
            response = await asyncio.wait_for(
                self._call_with_retries(client, system_prompt, user_msg),
                timeout=self._timeout + 2.0,  # hard ceiling
            )
        except asyncio.TimeoutError:
            elapsed_ms = round((time.time() - start) * 1000, 2)
            logger.warning("llm_judge_timeout", timeout_s=self._timeout, elapsed_ms=elapsed_ms)
            return L1Result(verdict="skip", chunks_evaluated=0, chunks_passed=0,
                           chunks_failed=0, failure_rate=0.0,
                           error="timeout", elapsed_ms=elapsed_ms)
        except Exception as e:
            elapsed_ms = round((time.time() - start) * 1000, 2)
            logger.warning("llm_judge_error", error=str(e)[:200], elapsed_ms=elapsed_ms)
            return L1Result(verdict="skip", chunks_evaluated=0, chunks_passed=0,
                           chunks_failed=0, failure_rate=0.0,
                           error=str(e)[:200], elapsed_ms=elapsed_ms)

        # Parse the response
        verdicts = self._parse_response(response, len(pairs))

        passed = sum(1 for v in verdicts if v.passed)
        failed = len(verdicts) - passed
        failure_rate = (failed / len(verdicts)) if verdicts else 0.0

        # Aggregate verdict
        if failed == 0:
            verdict = "pass"
        elif failure_rate >= 0.5:
            verdict = "fail"
        else:
            # Some pass, some fail — degraded but not catastrophic.
            # Caller can decide what to do.
            verdict = "fail"  # conservative: any failure = fail

        elapsed_ms = round((time.time() - start) * 1000, 2)

        # Cost estimate: deepseek-chat at $0.14/M in, $0.28/M out
        # Rough estimate: 500 input tokens + 100 output tokens per call
        cost_estimate = self._estimate_cost(len(pairs))

        return L1Result(
            verdict=verdict,
            chunks_evaluated=len(verdicts),
            chunks_passed=passed,
            chunks_failed=failed,
            failure_rate=round(failure_rate, 3),
            samples=[v.to_log_dict() for v in verdicts],
            model_used=self._model,
            elapsed_ms=elapsed_ms,
            cost_estimate_usd=cost_estimate,
        )

    async def _call_with_retries(self, client, system_prompt: str, user_msg: str):
        """Call the LLM with retry on transient errors."""
        last_exc = None
        for attempt in range(self._max_retries + 1):
            try:
                response = await client.chat.completions.create(
                    model=self._model,
                    messages=[
                        {"role": "system", "content": system_prompt},
                        {"role": "user", "content": user_msg},
                    ],
                    temperature=0.0,           # deterministic
                    max_tokens=500,           # enough for ~5 verdicts
                    response_format={"type": "json_object"},  # if supported
                )
                return response
            except Exception as e:
                last_exc = e
                if attempt < self._max_retries:
                    await asyncio.sleep(0.5)
        raise last_exc

    def _parse_response(self, response, expected_count: int) -> List[L1ChunkVerdict]:
        """
        Parse the LLM response into a list of verdicts.
        Robust to JSON wrapped in code fences or with extra commentary.
        """
        # Extract the message content
        try:
            content = response.choices[0].message.content or ""
        except (AttributeError, IndexError) as e:
            logger.warning("llm_judge_bad_response", error=str(e))
            return []

        content = content.strip()

        # Strip markdown code fences if present
        if content.startswith("```"):
            content = re.sub(r"^```(?:json)?\s*\n?", "", content)
            content = re.sub(r"\n?```\s*$", "", content)

        # Try to parse as JSON object (deepseek with json_object format) or array
        try:
            data = json.loads(content)
        except json.JSONDecodeError as e:
            logger.warning("llm_judge_json_parse_error", error=str(e), content_preview=content[:200])
            return []

        # Handle both {"verdicts": [...]} and direct [...]
        if isinstance(data, dict):
            # Look for a list-typed value
            items = None
            for key in ("verdicts", "results", "translations", "data"):
                if key in data and isinstance(data[key], list):
                    items = data[key]
                    break
            if items is None:
                # Take the first list-typed value
                for v in data.values():
                    if isinstance(v, list):
                        items = v
                        break
            if items is None:
                logger.warning("llm_judge_no_list_in_response")
                return []
        elif isinstance(data, list):
            items = data
        else:
            logger.warning("llm_judge_unexpected_response_type", type_=type(data).__name__)
            return []

        # Parse each item
        verdicts: List[L1ChunkVerdict] = []
        for item in items:
            try:
                v = L1ChunkVerdict(
                    accurate=str(item.get("accurate", "")).lower() == "yes",
                    fluent=str(item.get("fluent", "")).lower() == "yes",
                    correct_language=str(item.get("correct_language", "")).lower() == "yes",
                    no_leaks=str(item.get("no_leaks", "")).lower() == "yes",
                    reason=str(item.get("reason", ""))[:200],
                )
                verdicts.append(v)
            except Exception as e:
                logger.warning("llm_judge_item_parse_error", error=str(e), item=str(item)[:200])
                # Continue with what we have

        return verdicts

    def _estimate_cost(self, num_pairs: int) -> float:
        """Rough USD cost estimate for the call. Conservative (rounded up)."""
        # Approximation: 250 input tokens per pair + 50 output tokens per pair
        # (rough based on the prompt template + JSON output)
        input_tokens = 200 + (num_pairs * 250)
        output_tokens = num_pairs * 50
        # deepseek-chat pricing
        if "deepseek" in self._model.lower():
            input_cost = input_tokens / 1_000_000 * 0.14
            output_cost = output_tokens / 1_000_000 * 0.28
        elif "gpt-4o-mini" in self._model.lower():
            input_cost = input_tokens / 1_000_000 * 0.15
            output_cost = output_tokens / 1_000_000 * 0.60
        elif "gemini" in self._model.lower() and "flash" in self._model.lower():
            # gemini-2.5-flash-lite ~ $0.10/M in, $0.40/M out
            input_cost = input_tokens / 1_000_000 * 0.10
            output_cost = output_tokens / 1_000_000 * 0.40
        else:
            # Generic conservative estimate
            input_cost = input_tokens / 1_000_000 * 0.50
            output_cost = output_tokens / 1_000_000 * 1.00
        return round(input_cost + output_cost, 6)


# ---------- Convenience factory ----------

def make_judge_from_env() -> Optional[LLMJudge]:
    """
    Build an LLMJudge from environment variables. Returns None if no
    API key is configured.

    Reads:
      - L1_JUDGE_API_KEY (required)
      - L1_JUDGE_BASE_URL (default: DeepSeek)
      - L1_JUDGE_MODEL (default: deepseek-chat)
      - L1_JUDGE_TIMEOUT (default: 8.0)
    """
    import os
    api_key = os.getenv("L1_JUDGE_API_KEY", "").strip()
    if not api_key:
        return None
    return LLMJudge(
        api_key=api_key,
        base_url=os.getenv("L1_JUDGE_BASE_URL", "https://api.deepseek.com/v1"),
        model=os.getenv("L1_JUDGE_MODEL", "deepseek-chat"),
        timeout_seconds=float(os.getenv("L1_JUDGE_TIMEOUT", "8.0")),
    )