soup/soup_cli/utils/structured_output.py

"""Structured output constraints for inference (v0.30.0).

Supports JSON schema and regex. Backed by ``outlines`` or
``lm-format-enforcer`` if installed; otherwise returns an inert constraint
descriptor so the server can degrade gracefully.
"""

from __future__ import annotations

import re
from typing import Any, Literal, Optional

Mode = Literal["off", "json", "regex"]

_VALID_MODES = {"off", "json", "regex"}
_MAX_REGEX_LEN = 2048
_MAX_SCHEMA_STR_LEN = 65536


def validate_mode(mode: Optional[str]) -> Mode:
    """Normalise and validate the structured-output mode string."""
    if mode is None:
        return "off"
    normalised = mode.strip().lower()
    if normalised not in _VALID_MODES:
        raise ValueError(
            f"Unknown structured-output mode: {mode!r}. "
            f"Valid: {sorted(_VALID_MODES)}"
        )
    return normalised  # type: ignore[return-value]


def validate_regex_pattern(pattern: str) -> str:
    """Validate a regex pattern: length-bounded, must compile.

    We don't try to detect catastrophic-backtracking patterns here — that is
    the library (outlines / lm-format-enforcer) problem at runtime. We just
    cap length to stop attacker-supplied giant strings.
    """
    if not isinstance(pattern, str):
        raise ValueError("regex pattern must be a string")
    if "\x00" in pattern:
        raise ValueError("regex pattern contains null byte")
    if len(pattern) > _MAX_REGEX_LEN:
        raise ValueError(
            f"regex pattern length {len(pattern)} exceeds max {_MAX_REGEX_LEN}"
        )
    try:
        re.compile(pattern)
    except re.error as exc:
        raise ValueError(f"invalid regex: {exc}") from exc
    return pattern


def validate_json_schema(schema: Any) -> dict:
    """Accept a dict-shaped JSON schema. Must be JSON-serialisable and
    declare a top-level ``type`` field. Total serialised size is capped at
    64KB to prevent ReDoS-style schemas."""
    import json

    if not isinstance(schema, dict):
        raise ValueError("JSON schema must be a dict")
    try:
        serialised = json.dumps(schema)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"JSON schema is not serialisable: {exc}") from exc
    if len(serialised) > _MAX_SCHEMA_STR_LEN:
        raise ValueError(
            f"JSON schema size {len(serialised)} exceeds max {_MAX_SCHEMA_STR_LEN}"
        )
    # Minimum viable shape
    if "type" not in schema:
        raise ValueError("JSON schema must declare a 'type' field")
    return schema


def is_outlines_available() -> bool:
    try:
        import outlines  # noqa: F401

        return True
    except ImportError:
        return False


def is_lmfe_available() -> bool:
    try:
        import lmformatenforcer  # noqa: F401

        return True
    except ImportError:
        return False


def build_logits_processors(
    constraint: Optional[dict], tokenizer: Any,
) -> list:
    """Build a list of HF ``LogitsProcessor`` instances for ``constraint``.

    Returns an empty list when:
      - constraint is None / off
      - neither ``outlines`` nor ``lm-format-enforcer`` is installed
      - the chosen library cannot construct a processor for the given kind
        (we degrade to free-form rather than crashing the request)

    The returned list can be passed directly to
    ``model.generate(..., logits_processor=...)``.

    Security: this function never executes user-supplied code. The schema
    and regex are already validated upstream by ``validate_json_schema`` /
    ``validate_regex_pattern``.
    """
    if constraint is None:
        return []
    kind = constraint.get("kind")
    if kind not in ("json_schema", "regex"):
        return []

    # Prefer outlines (broader coverage); fall back to lm-format-enforcer.
    if is_outlines_available():
        try:
            return _build_outlines_processors(constraint, tokenizer)
        except Exception:  # noqa: BLE001 — degrade to free-form rather than 500
            return []
    if is_lmfe_available():
        try:
            return _build_lmfe_processors(constraint, tokenizer)
        except Exception:  # noqa: BLE001
            return []
    return []


def _build_outlines_processors(constraint: dict, tokenizer: Any) -> list:
    """Best-effort outlines integration. Schema-driver may be missing on
    older outlines builds so we try multiple entry points."""
    import outlines  # type: ignore

    kind = constraint["kind"]
    if kind == "json_schema":
        builder = (
            getattr(outlines, "JsonSchema", None)
            or getattr(outlines, "regex", None)
        )
        if builder is None:
            return []
        # outlines >= 0.1: outlines.processors.JSONLogitsProcessor
        proc_factory = getattr(
            __import__("outlines.processors", fromlist=["JSONLogitsProcessor"]),
            "JSONLogitsProcessor", None,
        )
        if proc_factory is None:
            return []
        return [proc_factory(constraint["schema"], tokenizer)]
    if kind == "regex":
        proc_factory = getattr(
            __import__("outlines.processors", fromlist=["RegexLogitsProcessor"]),
            "RegexLogitsProcessor", None,
        )
        if proc_factory is None:
            return []
        return [proc_factory(constraint["pattern"], tokenizer)]
    return []


def _build_lmfe_processors(constraint: dict, tokenizer: Any) -> list:
    """lm-format-enforcer integration."""
    from lmformatenforcer import (  # type: ignore
        JsonSchemaParser,
        RegexParser,
    )
    from lmformatenforcer.integrations.transformers import (  # type: ignore
        build_transformers_prefix_allowed_tokens_fn,
    )
    from transformers import LogitsProcessorList

    kind = constraint["kind"]
    if kind == "json_schema":
        parser = JsonSchemaParser(constraint["schema"])
    elif kind == "regex":
        parser = RegexParser(constraint["pattern"])
    else:
        return []

    fn = build_transformers_prefix_allowed_tokens_fn(tokenizer, parser)
    # PrefixConstrainedLogitsProcessor wants num_beams. We use 1 (greedy /
    # standard sampling) for chat completions.
    from transformers import PrefixConstrainedLogitsProcessor

    proc = PrefixConstrainedLogitsProcessor(fn, 1)
    processors = LogitsProcessorList()
    processors.append(proc)
    return list(processors)


def build_constraint(
    mode: Mode,
    json_schema: Optional[dict],
    regex_pattern: Optional[str],
) -> Optional[dict]:
    """Build a constraint descriptor for the server.

    Return shape: a dict with ``kind`` + kind-specific fields, or None if
    constraint is off / unsupported.

    Callers should treat None as "no constraint" (free-form generation).
    """
    if mode == "off":
        return None
    if mode == "json":
        if json_schema is not None:
            schema = validate_json_schema(json_schema)
            return {"kind": "json_schema", "schema": schema}
        return None  # free-form JSON is library-dependent; fall through
    if mode == "regex":
        if not regex_pattern:
            raise ValueError("regex mode requires a non-empty pattern")
        pattern = validate_regex_pattern(regex_pattern)
        return {"kind": "regex", "pattern": pattern}
    # Unreachable — validate_mode would have rejected
    raise ValueError(f"unexpected mode: {mode!r}")