UmbraTrans/tools/translation/llm.py

# Copyright 2026 Umbra Tools (umbratools.com)
# SPDX-License-Identifier: Apache-2.0

"""Minimal OpenAI-compatible client for the local LM Studio server.

Only the standard library is used. The Qwen models on BigBox emit reasoning
tokens unless told not to; ``reasoning_effort="none"`` is what makes a
translation call return prose instead of a thinking trace.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any

DEFAULT_BASE_URL = os.environ.get("UMBRA_TRANSLATION_BASE_URL", "http://192.168.123.16:1234/v1")
DEFAULT_MODEL = os.environ.get("UMBRA_TRANSLATION_MODEL", "qwen/qwen3.5-9b")
DEFAULT_ANALYST = os.environ.get(
    "UMBRA_TRANSLATION_ANALYST_MODEL",
    "lfm2.5-2.6b-qwen3.8-turbo-brilliance-power-x12-neo-max",
)
DEFAULT_CRITIC = os.environ.get("UMBRA_TRANSLATION_CRITIC_MODEL", "deepseek/deepseek-r1-0528-qwen3-8b")


@dataclass
class Usage:
    model: str
    prompt_tokens: int
    completion_tokens: int
    elapsed_s: float


class LLMError(RuntimeError):
    pass


class LLM:
    def __init__(
        self,
        base_url: str = DEFAULT_BASE_URL,
        cache_dir: Path | None = None,
        usage_log: Path | None = None,
        timeout: int = 600,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.cache_dir = cache_dir
        self.usage_log = usage_log
        self.timeout = timeout
        if cache_dir:
            cache_dir.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------ HTTP
    def _post(self, payload: dict[str, Any]) -> dict[str, Any]:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        req = urllib.request.Request(
            self.base_url + "/chat/completions",
            data=body,
            headers={"Content-Type": "application/json"},
        )
        try:
            with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                return json.loads(resp.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            detail = exc.read().decode("utf-8", "replace")[:800]
            raise LLMError(f"HTTP {exc.code}: {detail}") from exc
        except urllib.error.URLError as exc:
            raise LLMError(f"cannot reach {self.base_url}: {exc.reason}") from exc

    def chat(
        self,
        system: str,
        user: str,
        *,
        model: str = DEFAULT_MODEL,
        temperature: float = 0.3,
        max_tokens: int = 4096,
        reasoning: bool = False,
        json_mode: bool = False,
        cache_key: str | None = None,
        retries: int = 3,
    ) -> tuple[str, Usage]:
        digest = hashlib.sha256(
            json.dumps(
                {
                    "model": model,
                    "temperature": temperature,
                    "system": system,
                    "user": user,
                    "json_mode": json_mode,
                },
                ensure_ascii=False,
                sort_keys=True,
            ).encode("utf-8")
        ).hexdigest()
        cache_path = None
        if self.cache_dir:
            # Tie the cache file to the prompt hash as well as the logical key.
            # Otherwise a prompt edit silently reuses the old answer.
            safe_key = re.sub(r"[^A-Za-z0-9_.-]+", "_", cache_key or "chat")
            cache_path = self.cache_dir / f"{safe_key}.{digest[:12]}.json"
            if cache_path.exists():
                cached = json.loads(cache_path.read_text(encoding="utf-8"))
                return cached["content"], Usage(
                    model=model,
                    prompt_tokens=cached.get("prompt_tokens", 0),
                    completion_tokens=cached.get("completion_tokens", 0),
                    elapsed_s=0.0,
                )

        payload: dict[str, Any] = {
            "model": model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "temperature": temperature,
            "max_tokens": max_tokens,
        }
        if not reasoning:
            payload["reasoning_effort"] = "none"
        if json_mode:
            # LM Studio rejects the OpenAI "json_object" shorthand but accepts a
            # permissive json_schema. This keeps the model from wrapping the
            # answer in prose, which used to cost a full repair call.
            payload["response_format"] = {
                "type": "json_schema",
                "json_schema": {
                    "name": "translation_payload",
                    "schema": {"type": "object", "additionalProperties": True},
                },
            }

        last: Exception | None = None
        for attempt in range(retries):
            started = time.time()
            try:
                data = self._post(payload)
                message = data["choices"][0]["message"]
                content = (message.get("content") or "").strip()
                if not content:
                    # Reasoning models can spend the whole budget thinking.
                    raise LLMError(
                        "empty content "
                        f"(reasoning_tokens={data.get('usage', {}).get('completion_tokens_details', {}).get('reasoning_tokens')})"
                    )
                usage = data.get("usage", {})
                elapsed = time.time() - started
                u = Usage(
                    model=model,
                    prompt_tokens=int(usage.get("prompt_tokens", 0)),
                    completion_tokens=int(usage.get("completion_tokens", 0)),
                    elapsed_s=elapsed,
                )
                if cache_path:
                    cache_path.write_text(
                        json.dumps(
                            {
                                "content": content,
                                "prompt_tokens": u.prompt_tokens,
                                "completion_tokens": u.completion_tokens,
                                "model": model,
                            },
                            ensure_ascii=False,
                            indent=2,
                        )
                        + "\n",
                        encoding="utf-8",
                    )
                if self.usage_log:
                    self.usage_log.parent.mkdir(parents=True, exist_ok=True)
                    with self.usage_log.open("a", encoding="utf-8") as fh:
                        fh.write(
                            json.dumps(
                                {
                                    "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                                    "model": model,
                                    "prompt_tokens": u.prompt_tokens,
                                    "completion_tokens": u.completion_tokens,
                                    "elapsed_s": round(elapsed, 2),
                                    "cache_key": cache_key,
                                }
                            )
                            + "\n"
                        )
                return content, u
            except Exception as exc:  # noqa: BLE001 - retry boundary
                last = exc
                if attempt < retries - 1:
                    time.sleep(2**attempt)
        raise LLMError(f"chat failed after {retries} attempts: {last}")

    # ------------------------------------------------------------------ JSON
    def chat_json(
        self,
        system: str,
        user: str,
        *,
        model: str = DEFAULT_MODEL,
        temperature: float = 0.2,
        max_tokens: int = 4096,
        cache_key: str | None = None,
        json_mode: bool = True,
    ) -> tuple[Any, Usage]:
        content, usage = self.chat(
            system,
            user,
            model=model,
            temperature=temperature,
            max_tokens=max_tokens,
            reasoning=False,
            json_mode=json_mode,
            cache_key=cache_key,
        )
        try:
            return parse_json(content), usage
        except ValueError:
            # One repair attempt, with the malformed output in front of the model.
            repair = (
                "The previous answer was not valid JSON. Return only valid JSON, "
                "no code fence, no commentary. The required shape was:\n"
                f"{user}\n\nPrevious answer:\n{content[:6000]}"
            )
            fixed, usage2 = self.chat(
                system,
                repair,
                model=model,
                temperature=0.0,
                max_tokens=max_tokens,
                reasoning=False,
                json_mode=json_mode,
                cache_key=(cache_key + "-repair") if cache_key else None,
            )
            return parse_json(fixed), usage2


def parse_json(text: str) -> Any:
    """Parse a JSON object/array from a model answer.

    Accepts a code fence, leading prose, or a trailing sentence. Raises
    ValueError when no complete JSON value can be decoded.
    """
    text = text.strip()
    if text.startswith("```"):
        text = re.sub(r"^```[a-zA-Z0-9_-]*\s*", "", text)
        text = re.sub(r"\s*```$", "", text).strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    decoder = json.JSONDecoder()
    for i, ch in enumerate(text):
        if ch in "{[":
            try:
                value, _ = decoder.raw_decode(text[i:])
                return value
            except json.JSONDecodeError:
                continue
    raise ValueError(f"no JSON value found in answer: {text[:300]!r}")