deployment/contracts.py

"""Wire validation and deterministic semantic bridge to existing decision prompts."""

import json
import math
from dataclasses import dataclass

PROMPT_VERSION = "jev.dynamic.prompt.v2"
LABELS = tuple("ABCDEFGHIJKLMNOP")
BINARY_CRITERIA = [
    {"id": "yes", "description": "The stated proposition is true."},
    {"id": "no", "description": "The stated proposition is false."},
]


def validate_request(record):
    for key in ("id", "group_id", "state", "instructions"):
        if not isinstance(record.get(key), str) or not record[key].strip():
            raise ValueError(f"{key} must be a nonempty string")
    if record.get("primitive") not in ("choice", "noul", "score_level"):
        raise ValueError("unsupported primitive")
    criteria = record.get("criteria")
    if not isinstance(criteria, list) or not 2 <= len(criteria) <= len(LABELS):
        raise ValueError("criteria must contain 2..16 candidates")
    ids = []
    for candidate in criteria:
        if not isinstance(candidate, dict):
            raise TypeError("candidate must be an object")
        for key in ("id", "description"):
            if not isinstance(candidate.get(key), str) or not candidate[key].strip():
                raise ValueError(f"candidate {key} must be nonempty")
        ids.append(candidate["id"])
    if len(set(ids)) != len(ids):
        raise ValueError("duplicate candidate ids")
    if record["primitive"] != "choice" and criteria != BINARY_CRITERIA:
        raise ValueError("noul and score_level require canonical yes/no criteria")


def validate_record(record):
    validate_request(record)
    if record.get("gold") not in [c["id"] for c in record["criteria"]]:
        raise ValueError("gold must be a candidate id")
    if not isinstance(record.get("provenance"), dict):
        raise TypeError("provenance must be an object")


def label_mapping(record):
    validate_request(record)
    return dict(zip(LABELS, (c["id"] for c in record["criteria"])))


def render_prompt_parts(record):
    """Text prefix/suffix; callers MUST check tokenizer boundary equivalence."""
    validate_request(record)
    prefix = "Shared state:\n" + record["state"] + "\n\n"
    task = {
        "primitive": record["primitive"],
        "instructions": record["instructions"],
        "criteria": [
            {"label": label, "description": candidate["description"]}
            for label, candidate in zip(LABELS, record["criteria"])
        ],
    }
    suffix = json.dumps(task, ensure_ascii=False, sort_keys=True)
    suffix += (
        "\nReturn only the selected letter: "
        + ", ".join(LABELS[: len(record["criteria"])])
        + ".\nAnswer:"
    )
    return prefix, suffix


def render_prompt(record):
    return "".join(render_prompt_parts(record))


COMPILER_VERSION = "jev.typesafe.compiler.v1"
CONFIDENCE_POLICY = "normalized-entropy.v1"


class ServingError(Exception):
    def __init__(self, status, code, message):
        super().__init__(message)
        self.status, self.code, self.message = status, code, message


def invalid(message):
    raise ServingError(422, "invalid_request", message)


def text(value):
    if isinstance(value, str) and value.strip():
        return value
    return json.dumps(value, ensure_ascii=False, sort_keys=True, allow_nan=False)


def content(value):
    return isinstance(value, (str, dict, list))


@dataclass(frozen=True)
class Question:
    key: str
    kind: str
    record: dict
    legend: dict | None = None


def prepare_request(body, *, model_name, max_questions=16, allow_score=True):
    if not isinstance(body, dict) or set(body) - {
        "state",
        "questions",
        "model",
        "effort",
    }:
        invalid("Expected state, questions, model and optional effort only")
    if not isinstance(body.get("model", "jev-latest"), str) or body.get(
        "model", "jev-latest"
    ) not in {model_name, "jev-latest"}:
        raise ServingError(404, "unknown_model", "Model is not served")
    if not content(body.get("state")):
        invalid("state must be a string, object or array")
    questions = body.get("questions")
    if not isinstance(questions, dict) or not 1 <= len(questions) <= max_questions:
        invalid("questions must be a nonempty map within the configured limit")
    result = []
    for key, question in questions.items():
        if not isinstance(key, str) or not key.strip():
            invalid("Question IDs must be nonempty strings")
        if not isinstance(question, dict) or set(question) - {
            "type",
            "instructions",
            "criteria",
        }:
            invalid("Invalid question fields")
        kind = question.get("type")
        instruction = question.get("instructions")
        if instruction is not None and not content(instruction):
            invalid("instructions must be text, object, array or null")
        instruction = (
            text(instruction)
            if instruction is not None
            else "Evaluate the state against the supplied criteria."
        )
        criteria = question.get("criteria")
        legend = None
        internal_kind = kind
        if kind == "choice":
            if not isinstance(criteria, dict) or not 2 <= len(criteria) <= len(LABELS):
                invalid(
                    "This model supports 2..16 Choice candidates; none are truncated"
                )
            candidates = []
            for name, description in criteria.items():
                if not isinstance(name, str) or not name.strip():
                    invalid("Choice names must be nonempty strings")
                if description is not None and not content(description):
                    invalid("Invalid Choice description")
                candidates.append(
                    {
                        "id": name,
                        "description": name
                        if description is None
                        else text(description),
                    }
                )
        elif kind == "noul":
            if criteria is not None:
                if not isinstance(criteria, dict) or set(criteria) - {"true", "false"}:
                    invalid("Noul criteria accepts true and false descriptions")
                if any(v is not None and not content(v) for v in criteria.values()):
                    invalid("Invalid Noul criterion")
                if any(v is not None for v in criteria.values()) and criteria != {
                    "true": BINARY_CRITERIA[0]["description"],
                    "false": BINARY_CRITERIA[1]["description"],
                }:
                    instruction = text(
                        {"question": instruction, "meaning_of_answers": criteria}
                    )
            candidates = [dict(c) for c in BINARY_CRITERIA]
        elif kind == "score":
            if not allow_score:
                raise ServingError(
                    422,
                    "unsupported_score",
                    "Ordinal Score is experimental and disabled; it is not score_level",
                )
            if not isinstance(criteria, list) or not 1 <= len(criteria) <= 10:
                invalid("Score requires 1..10 ordered levels")
            if any(not content(v) for v in criteria):
                invalid("Invalid Score level")
            legend = {str(i): v for i, v in enumerate(criteria)}
            candidates = [{"id": k, "description": text(v)} for k, v in legend.items()]
            internal_kind = "choice"
            instruction = text(
                {
                    "question": instruction,
                    "task": "Select the best matching ordered rubric level. Levels run from lowest to highest.",
                }
            )
        else:
            invalid("Unsupported question type")
        record = {
            "id": key,
            "group_id": "serving",
            "state": text(body["state"]),
            "primitive": internal_kind,
            "instructions": instruction,
            "criteria": candidates,
        }
        result.append(Question(key, kind, record, legend))
    return result


@dataclass(frozen=True)
class Compiled:
    question: Question
    prompt_ids: list[int]
    candidate_token_ids: list[int]


def compile_question(question, tokenizer, max_tokens=8192):
    if question.kind == "score" and len(question.record["criteria"]) == 1:
        return Compiled(question, [], [])
    prompt = tokenizer.apply_chat_template(
        [{"role": "user", "content": render_prompt(question.record)}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=False,
    )
    ids = tokenizer.encode(prompt, add_special_tokens=False)
    if not ids or len(ids) > max_tokens:
        invalid("Compiled prompt exceeds token budget; no truncation is allowed")
    special = set(tokenizer.all_special_ids)
    if special.intersection(
        tokenizer.encode(render_prompt(question.record), add_special_tokens=False)
    ):
        invalid("Input contains reserved model control tokens")
    tokens = []
    for label in LABELS[: len(question.record["criteria"])]:
        encoded = tokenizer.encode(label, add_special_tokens=False)
        if (
            len(encoded) != 1
            or tokenizer.encode(prompt + label, add_special_tokens=False)
            != ids + encoded
        ):
            invalid("Candidate is not one token at the answer boundary")
        if encoded[0] in special or encoded[0] in tokens:
            invalid("Invalid candidate token mapping")
        tokens.append(encoded[0])
    return Compiled(question, ids, tokens)


def probabilities(scores):
    if len(scores) < 2 or any(
        type(v) not in (int, float) or not math.isfinite(v) for v in scores
    ):
        raise ServingError(
            502, "invalid_distribution", "Backend returned invalid candidate scores"
        )
    shift = max(scores)
    values = [math.exp(v - shift) for v in scores]
    total = sum(values)
    return [v / total for v in values]


def confidence(probs):
    if len(probs) == 1:
        return 1.0
    entropy = -sum(p * math.log(p) for p in probs if p > 0)
    return min(1.0, max(0.0, 1.0 - entropy / math.log(len(probs))))


def answer(question, probs):
    keys = [c["id"] for c in question.record["criteria"]]
    if (
        len(keys) != len(probs)
        or any(
            type(p) not in (int, float) or not math.isfinite(p) or not 0 <= p <= 1
            for p in probs
        )
        or not math.isclose(sum(probs), 1, abs_tol=1e-6)
    ):
        raise ServingError(
            502,
            "invalid_distribution",
            "Backend distribution does not match candidates",
        )
    if question.kind == "noul":
        return {"type": "noul", "noul": probs[keys.index("yes")]}
    value = {
        "type": question.kind,
        "probabilities": dict(zip(keys, probs)),
        "confidence": confidence(probs),
    }
    if question.kind == "choice":
        value["choice"] = keys[max(range(len(probs)), key=probs.__getitem__)]
    else:
        value.update(
            score=sum(i * p for i, p in enumerate(probs)), legend=question.legend
        )
    return value