llm.py

"""Thin client to the MiniCPM endpoint (vLLM on Modal, OpenAI-compatible).

`chat_json` sends a guided-JSON request so the model is forced to emit JSON that
matches our schema — no fence-stripping or repair needed. If the endpoint isn't
configured, callers fall back to the deterministic story (see narrate.py), so the
app never hard-crashes on a missing backend.
"""
from __future__ import annotations

import json

import config

_client = None
_NO_THINK_SUFFIX = "/no_think"


class LLMUnavailable(RuntimeError):
    """Raised when no model backend is configured/reachable."""


def available() -> bool:
    return bool(config.MODAL_ENDPOINT_URL and config.MODAL_API_KEY) or config.USE_ZEROGPU_FALLBACK


def _get_client():
    global _client
    if _client is None:
        from openai import OpenAI
        _client = OpenAI(base_url=config.MODAL_ENDPOINT_URL,
                         api_key=config.MODAL_API_KEY or "x",
                         timeout=config.LLM_TIMEOUT_S)
    return _client


def chat_json(messages: list[dict], schema: dict, max_tokens: int = 1024) -> dict:
    """Return a JSON object from the model, constrained to `schema`."""
    if config.USE_ZEROGPU_FALLBACK and not config.MODAL_ENDPOINT_URL:
        import zerogpu_backend
        return zerogpu_backend.chat_json(messages, schema)

    if not (config.MODAL_ENDPOINT_URL and config.MODAL_API_KEY):
        raise LLMUnavailable("No MODAL_ENDPOINT_URL / MODAL_API_KEY configured.")

    client = _get_client()
    # vLLM accepts guided JSON via extra_body; xgrammar backend enforces the schema.
    resp = client.chat.completions.create(
        model=config.MODEL_ID,
        messages=_without_thinking(messages),
        temperature=0.4,
        max_tokens=max_tokens,
        extra_body={
            "guided_json": schema,
            "guided_decoding_backend": "xgrammar",
            "add_special_tokens": True,
        },
    )
    content = resp.choices[0].message.content or "{}"
    return _loads(content)


def _loads(content: str) -> dict:
    try:
        return json.loads(content)
    except (json.JSONDecodeError, ValueError):
        # Guided decoding should make this unreachable, but stay defensive.
        s = content.strip().removeprefix("```json").removeprefix("```").removesuffix("```")
        # The model sometimes outputs \n and \" as literal text in thinking mode.
        s = s.replace("\\n", "\n").replace('\\"', '"')
        start, end = s.find("{"), s.rfind("}")
        if start != -1 and end != -1:
            try:
                return json.loads(s[start:end + 1])
            except (json.JSONDecodeError, ValueError):
                pass
        return {}


def _without_thinking(messages: list[dict]) -> list[dict]:
    """Force MiniCPM4.1 into non-reasoning mode without mutating callers' prompts."""
    guarded = [dict(m) for m in messages]
    for msg in reversed(guarded):
        if msg.get("role") == "user" and isinstance(msg.get("content"), str):
            content = msg["content"].rstrip()
            if not content.endswith(_NO_THINK_SUFFIX):
                msg["content"] = f"{content}\n\n{_NO_THINK_SUFFIX}"
            break
    return guarded