llm-in-text/backend/llm_policy.py

from dataclasses import dataclass
from typing import Any

from risk_config import RiskConfig


@dataclass(frozen=True)
class LLMPolicy:
    job_type: str
    model: str
    profile: str
    max_input_chars: int
    max_output_tokens: int
    temperature: float
    thinking: str | None


def _normalize_thinking(value: str | None, *, allow_high: bool) -> str | None:
    candidate = (value or "").strip().lower()
    if candidate in {"", "none", "off"}:
        return None
    if candidate not in {"low", "medium", "high"}:
        return "low"
    if candidate == "high" and not allow_high:
        return "medium"
    return candidate


def resolve_llm_policy(job_type: str, request_payload: dict[str, Any], config: RiskConfig) -> LLMPolicy:
    if job_type == "completion":
        return LLMPolicy(
            job_type=job_type,
            model=config.completion_model,
            profile="completion",
            max_input_chars=config.completion_max_input_chars,
            max_output_tokens=config.completion_max_output_tokens,
            temperature=config.completion_temperature,
            thinking=_normalize_thinking(request_payload.get("model_thinking"), allow_high=False),
        )
    if job_type == "pro_completion":
        return LLMPolicy(
            job_type=job_type,
            model=config.pro_model,
            profile="pro",
            max_input_chars=config.pro_max_input_chars,
            max_output_tokens=config.pro_max_output_tokens,
            temperature=config.pro_temperature,
            thinking=_normalize_thinking(request_payload.get("pro_thinking"), allow_high=True) or "medium",
        )
    if job_type == "compress":
        return LLMPolicy(
            job_type=job_type,
            model=config.completion_model,
            profile="completion",
            max_input_chars=config.compress_max_input_chars,
            max_output_tokens=config.compress_max_output_tokens,
            temperature=0.2,
            thinking="low",
        )
    if job_type == "ocr":
        return LLMPolicy(
            job_type=job_type,
            model=config.vision_model,
            profile="vision",
            max_input_chars=config.ocr_max_input_bytes,
            max_output_tokens=config.completion_max_output_tokens,
            temperature=0.0,
            thinking=None,
        )
    raise ValueError(f"unsupported llm policy job type: {job_type}")