training/data/ai_version_generator.py

"""OpenRouter-backed AI rewriter used by the Stage 1 dataset builder.

Given a human-authored German paragraph, :class:`AI_Version_Generator`
asks an ``openai/gpt-5.4`` deployment on OpenRouter to produce an
AI-styled rewrite that must preserve every citation verbatim. The
result is a tuple ``(ai_text, citations, attempts)`` consumed by
``scripts/build_dataset_v2.py``.

Citation detection is delegated to
:mod:`training.data.citation_utils` — this module never redefines the
regex, so the TS-equivalent invariant documented in ``citation_utils``
continues to hold for the dataset pipeline.

Retry policy (Requirement 1.10)
-------------------------------
* Maximum 3 HTTP attempts per paragraph.
* Retries fire on any of: connect/read timeout, generic transport
  error, HTTP 5xx, HTTP 429 (rate limited), or an empty ``content``
  field in the OpenAI-style response body.
* 4xx responses other than 429 are treated as unrecoverable
  (bad request, invalid key, insufficient credits, etc.) and short-
  circuit to ``None`` immediately.
* Backoff between attempts is ``2 ** attempt`` seconds, so 2s after
  attempt 1, 4s after attempt 2, 8s after attempt 3 — matching the
  "2-4-8" pattern in design.md.

The class is usable either manually (``gen = AI_Version_Generator(key);
... ; gen.close()``) or as a context manager
(``with AI_Version_Generator(key) as gen: ...``).

Relates to requirements 1.4, 1.10 and 2.1 of spec
``grpo-humanizer-training-v2``.
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass
from typing import Any

import httpx

from training.data.citation_utils import extract_citations

__all__ = [
    "AI_MODEL",
    "AIVersionResult",
    "AI_Version_Generator",
]


logger = logging.getLogger(__name__)


#: OpenRouter model identifier used for all AI-version generations.
#: Fixed by Requirement 1.4 — do not parameterise without updating the
#: spec first.
AI_MODEL: str = "openai/gpt-5.4"


#: OpenRouter chat completions endpoint (OpenAI-compatible).
_OPENROUTER_URL: str = "https://openrouter.ai/api/v1/chat/completions"

#: Maximum number of HTTP attempts per paragraph (Requirement 1.10).
_MAX_ATTEMPTS: int = 3


@dataclass
class AIVersionResult:
    """Outcome of a successful OpenRouter generation.

    Attributes
    ----------
    ai_text:
        The AI-rewritten paragraph, stripped of leading/trailing
        whitespace. Guaranteed non-empty (empty responses trigger a
        retry).
    citations:
        Citations extracted from the *original* human text — i.e.
        the citations OpenRouter was instructed to preserve. Not the
        citations found in ``ai_text``. Callers that need to validate
        citation preservation should run
        :func:`training.data.citation_utils.passes_citation_check`
        against this list via a fresh extraction of ``ai_text``.
    attempts:
        Number of HTTP attempts made to obtain the result. ``1`` when
        the first request succeeded; up to ``3`` when earlier attempts
        were retried per Requirement 1.10.
    """

    ai_text: str
    citations: list[str]
    attempts: int


class AI_Version_Generator:
    """Thin synchronous OpenRouter client for AI-version generation.

    One instance wraps a single :class:`httpx.Client` with the required
    authorization header. The client is reused across many
    ``generate`` calls so HTTP keep-alive reduces per-request overhead.
    Always close the client via :meth:`close` or the context-manager
    protocol when done.

    Parameters
    ----------
    api_key:
        OpenRouter API key (``sk-or-v1-…``). Sent in the
        ``Authorization: Bearer`` header on every request.
    timeout:
        Per-request timeout in seconds, passed straight to
        :class:`httpx.Client`. Default ``60.0`` matches design.md.
    """

    def __init__(self, api_key: str, timeout: float = 60.0) -> None:
        self._client: httpx.Client = httpx.Client(
            timeout=timeout,
            headers={
                "Authorization": f"Bearer {api_key}",
                "HTTP-Referer": "https://github.com/LevArtesa/ghostwriter",
                "Content-Type": "application/json",
            },
        )

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    def close(self) -> None:
        """Close the underlying HTTP client. Safe to call multiple times."""
        self._client.close()

    def __enter__(self) -> "AI_Version_Generator":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    # ------------------------------------------------------------------
    # Prompt building
    # ------------------------------------------------------------------

    def _build_messages(
        self, human_text: str, citations: list[str]
    ) -> list[dict[str, str]]:
        """Build the OpenAI-style chat messages for one paragraph.

        The system prompt instructs GPT-5.4 (in German) to rewrite the
        paragraph in a distinctly AI-typical style while copy-pasting
        every listed citation verbatim. When ``citations`` is empty we
        still emit a non-empty fallback token so the model sees a
        well-formed list. The fallback text (``"(нет цитат)"``) comes
        from design.md verbatim — its Russian wording is harmless
        because the model is instructed to preserve *only* citations
        from the list, and an empty Russian tag signals "nothing to
        preserve".
        """
        citation_list = (
            "\n".join(f"- {c}" for c in citations) if citations else "(нет цитат)"
        )
        system_content = (
            "Du bist ein KI-Textgenerator, der akademische Texte im "
            "typischen KI-Stil produziert: gleichmäßige Satzlänge, "
            "formelhafte Konnektoren (\"darüber hinaus\", \"in diesem "
            "Zusammenhang\"), abstrakte Verben, Redundanz. Deutsch.\n\n"
            "WICHTIG: Übernimm ALLE folgenden Zitate DORTGENAU "
            "(copy-paste). Verändere keine Autorennamen, Jahreszahlen, "
            "Seitenzahlen:\n"
            f"{citation_list}\n\n"
            "Gib NUR den umgeschriebenen Text zurück."
        )
        user_content = (
            "Schreibe den folgenden Absatz im KI-Stil neu:\n\n" + human_text
        )
        return [
            {"role": "system", "content": system_content},
            {"role": "user", "content": user_content},
        ]

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------

    def generate(self, human_text: str) -> AIVersionResult | None:
        """Request an AI-styled rewrite from OpenRouter.

        Parameters
        ----------
        human_text:
            Source paragraph in German. Citations are extracted up-
            front and forwarded to the model so it can preserve them
            verbatim.

        Returns
        -------
        AIVersionResult | None
            The rewrite and metadata on success, or ``None`` when all
            three retry attempts fail. A ``None`` return is logged at
            WARNING level with the last observed error so callers can
            surface it in checkpoint status.
        """
        citations = extract_citations(human_text)
        messages = self._build_messages(human_text, citations)
        payload: dict[str, Any] = {
            "model": AI_MODEL,
            "messages": messages,
            "temperature": 0.7,
            "top_p": 1.0,
        }

        last_error: str | None = None

        for attempt in range(1, _MAX_ATTEMPTS + 1):
            retryable, error_msg, ai_text = self._attempt_once(payload)

            if ai_text is not None:
                # Success path. `attempt` is the final attempt count.
                return AIVersionResult(
                    ai_text=ai_text,
                    citations=citations,
                    attempts=attempt,
                )

            # Failure path. Remember the reason for the final warning
            # log, then decide whether to keep retrying.
            last_error = error_msg
            if not retryable:
                logger.warning(
                    "OpenRouter: non-retryable error on attempt %d: %s",
                    attempt,
                    error_msg,
                )
                return None

            # Exponential backoff: 2s after attempt 1, 4s after
            # attempt 2, 8s after attempt 3. The last sleep is only
            # executed if another attempt would follow it; after
            # the final attempt we fall through to the warning and
            # return None.
            backoff = 2 ** attempt
            if attempt < _MAX_ATTEMPTS:
                logger.info(
                    "OpenRouter: retrying after attempt %d (%s); "
                    "sleeping %ds",
                    attempt,
                    error_msg,
                    backoff,
                )
                time.sleep(backoff)

        logger.warning(
            "OpenRouter: giving up after %d attempts, last error: %s",
            _MAX_ATTEMPTS,
            last_error,
        )
        return None

    # ------------------------------------------------------------------
    # Single-attempt helper
    # ------------------------------------------------------------------

    def _attempt_once(
        self, payload: dict[str, Any]
    ) -> tuple[bool, str | None, str | None]:
        """Run a single HTTP attempt.

        Returns a 3-tuple ``(retryable, error_msg, ai_text)``:

        * On success, ``ai_text`` is the stripped content and
          ``error_msg`` is ``None``.
        * On retryable failure (timeout, transport error, 5xx, 429,
          empty content), ``retryable`` is ``True`` and ``ai_text`` is
          ``None``.
        * On unrecoverable failure (4xx other than 429, malformed
          JSON, unexpected shape), ``retryable`` is ``False``.
        """
        try:
            response = self._client.post(_OPENROUTER_URL, json=payload)
        except httpx.TimeoutException as exc:
            return True, f"timeout: {exc!r}", None
        except httpx.RequestError as exc:
            return True, f"request error: {exc!r}", None

        status = response.status_code
        if status >= 500 or status == 429:
            return True, f"HTTP {status}", None
        if status >= 400:
            # 4xx other than 429 — auth / bad request. Surface body
            # snippet (truncated) in the error message so operators
            # can diagnose without retrying pointlessly.
            body_snippet = response.text[:200]
            return False, f"HTTP {status}: {body_snippet}", None

        try:
            data = response.json()
        except ValueError as exc:
            # 2xx with non-JSON body — treat as unrecoverable, the
            # endpoint is misbehaving in a way retries won't fix.
            return False, f"invalid JSON: {exc!r}", None

        try:
            content = data["choices"][0]["message"]["content"]
        except (KeyError, IndexError, TypeError) as exc:
            return False, f"unexpected response shape: {exc!r}", None

        if not isinstance(content, str):
            return False, f"non-string content: {type(content).__name__}", None

        ai_text = content.strip()
        if not ai_text:
            # Empty content is treated as retryable per task contract
            # — a transient quirk of the upstream model rather than a
            # permanent misconfiguration.
            return True, "empty content", None

        return True, None, ai_text