Blame for questions.py

"""Map Gemma letter logits onto TypeSafe's three question primitives.

The letter distribution is the model's raw output. Each primitive renormalizes
the distribution over its legal answer letters, so the probabilities sum to 1
across the caller's options.
"""

from __future__ import annotations

import math
from dataclasses import dataclass, field

from .letters import LETTER_INDEX, LETTERS, _softmax

# Fixed bank of question-agnostic demonstrations. Their predicates differ from the
# target question on purpose: they teach the letter-slot format, not the task.
GENERIC_EXAMPLES: tuple[tuple[str, str, str], ...] = (
    ("I need this fixed immediately, it is blocking everything.", "Does the text express urgency?", "Y"),
    ("Here is the documentation link for reference.", "Does the text express urgency?", "N"),
    ("The app crashes when I open the settings page.", "Does the text contain a bug report?", "Y"),
    ("Just wanted to share the meeting notes from today.", "Does the text contain a bug report?", "N"),
    ("Please send this to a human agent right now.", "Does the text ask for a human agent?", "Y"),
    ("No rush on this, whenever you get a chance.", "Does the text ask for a human agent?", "N"),
    ("URGENT: the payment system is failing right now.", "Does the text express urgency?", "Y"),
    ("The sky is blue and the grass is green.", "Does the text express urgency?", "N"),
)

DEFAULT_EXAMPLE_COUNT = 6


@dataclass(frozen=True)
class QuestionOption:
    letter: str
    label: str
    description: str | None = None


@dataclass
class Question:
    id: str
    type: str  # noul | choice | score
    instructions: str
    options: list[QuestionOption] = field(default_factory=list)
    criteria: object = None


@dataclass(frozen=True)
class QuestionAnswer:
    type: str
    choice: str | None = None
    probabilities: dict[str, float] = field(default_factory=dict)
    confidence: float | None = None
    noul: float | None = None
    score: float | None = None
    legend: dict[str, str] | None = None

    def to_dict(self) -> dict:
        payload: dict = {"type": self.type}
        if self.type == "noul":
            payload["noul"] = round(self.noul or 0.0, 6)
        elif self.type == "choice":
            payload["choice"] = self.choice
            payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()}
            payload["confidence"] = round(self.confidence or 0.0, 6)
        elif self.type == "score":
            payload["score"] = round(self.score or 0.0, 6)
            payload["legend"] = self.legend or {}
            payload["probabilities"] = {k: round(v, 6) for k, v in self.probabilities.items()}
            payload["confidence"] = round(self.confidence or 0.0, 6)
        return payload


def confidence_from_probabilities(probabilities: list[float]) -> float:
    """1.0 when all mass is on one option, 0.0 when uniform."""
    n = len(probabilities)
    if n <= 1:
        return 1.0
    entropy = -sum(p * math.log(p) for p in probabilities if p > 0)
    return max(0.0, 1.0 - entropy / math.log(n))


def _predicate_forms(instructions: str) -> tuple[str, str]:
    """Turn an instruction into affirmative/negative option text.

    "Does the text request a refund?" -> ("Yes, the text requests a refund.",
                                          "No, the text does not request a refund.")
    """
    text = instructions.strip().rstrip(".").strip()
    lowered = text.lower()

    if lowered.startswith("does the text "):
        body = text[len("Does the text ") :]
        affirmative = f"the text {body}"
    elif lowered.startswith("is the text "):
        body = text[len("Is the text ") :]
        affirmative = f"the text {body}"
    elif lowered.startswith("is this "):
        body = text[len("Is this ") :]
        affirmative = f"this {body}"
    elif lowered.startswith("does this "):
        body = text[len("Does this ") :]
        affirmative = f"this {body}"
    elif lowered.startswith("does the "):
        body = text[len("Does the ") :]
        affirmative = f"the {body}"
    elif lowered.startswith("is the "):
        body = text[len("Is the ") :]
        affirmative = f"the {body}"
    elif lowered.startswith("did the "):
        body = text[len("Did the ") :]
        affirmative = f"the {body}"
    elif lowered.startswith("can the "):
        body = text[len("Can the ") :]
        affirmative = f"the {body}"
    else:
        affirmative = text

    affirmative = affirmative.rstrip("?").strip()
    if affirmative.lower().startswith("the "):
        negative = "the " + affirmative[4:].split(" ", 1)[0] + " does not" + (
            " " + affirmative[4:].split(" ", 1)[1] if " " in affirmative[4:] else ""
        )
    else:
        negative = f"it is not the case that {affirmative}"
    return f"Yes, {affirmative}.", f"No, {negative}."


def build_question(qid: str, spec: dict) -> Question:
    """Parse one TypeSafe question spec into a letter-addressable Question."""
    qtype = spec.get("type")
    if qtype not in {"noul", "choice", "score"}:
        raise ValueError(f"question {qid!r} has unsupported type {qtype!r}")

    instructions = spec.get("instructions")
    if not isinstance(instructions, str) or not instructions.strip():
        raise ValueError(f"question {qid!r} requires non-empty string instructions")

    criteria = spec.get("criteria")

    if qtype == "noul":
        if isinstance(criteria, dict):
            yes_desc = criteria.get("true")
            no_desc = criteria.get("false")
        else:
            yes_desc = no_desc = None
        yes_text, no_text = _predicate_forms(instructions)
        return Question(
            id=qid,
            type=qtype,
            instructions=instructions,
            options=[
                QuestionOption("Y", "yes", yes_desc or yes_text),
                QuestionOption("N", "no", no_desc or no_text),
            ],
            criteria=criteria,
        )

    if qtype == "choice":
        if not isinstance(criteria, dict) or not criteria:
            raise ValueError(f"choice question {qid!r} requires a non-empty criteria map")
        if len(criteria) > len(LETTERS):
            raise ValueError(f"choice question {qid!r} supports at most 26 options")
        entries = list(criteria.items())
    else:
        if not isinstance(criteria, list) or len(criteria) < 2:
            raise ValueError(f"score question {qid!r} requires an ordered list of at least two levels")
        if len(criteria) > len(LETTERS):
            raise ValueError(f"score question {qid!r} supports at most 26 levels")
        entries = [(f"level {i}", level) for i, level in enumerate(criteria)]

    options = [
        QuestionOption(
            LETTERS[i],
            label,
            description if isinstance(description, str) else str(description),
        )
        for i, (label, description) in enumerate(entries)
    ]
    return Question(id=qid, type=qtype, instructions=instructions, options=options, criteria=criteria)


def score_question(question: Question, letter_logits: dict[str, float], temperature: float = 1.0) -> QuestionAnswer:
    """Turn raw letter logits into a typed answer, renormalized over legal options."""
    option_letters = [option.letter for option in question.options]
    missing = [letter for letter in option_letters if letter not in letter_logits]
    if missing:
        raise ValueError(f"question {question.id!r} is missing logits for {missing}")

    logits = [letter_logits[letter] for letter in option_letters]
    probs = _softmax(logits, temperature)

    if question.type == "noul":
        noul = probs[0]
        return QuestionAnswer(
            type="noul",
            noul=noul,
        )

    if question.type == "choice":
        pairs = sorted(zip(question.options, probs), key=lambda op: op[1], reverse=True)
        top = pairs[0]
        return QuestionAnswer(
            type="choice",
            choice=top[0].label,
            probabilities={option.label: p for option, p in pairs},
            confidence=confidence_from_probabilities(probs),
        )

    score = sum(i * p for i, p in enumerate(probs))
    return QuestionAnswer(
        type="score",
        score=score,
        probabilities={str(i): p for i, p in enumerate(probs)},
        legend={str(i): option.description or "" for i, option in enumerate(question.options)},
        confidence=confidence_from_probabilities(probs),
    )


def score_question_tensor(question: Question, letter_logits, temperature: float = 1.0) -> QuestionAnswer:
    """Score a ``[26]`` device tensor, transferring only final answer values.

    The logits, legal-option gather, softmax, entropy, ranking, and score
    reduction all remain on the tensor's device.  Constructing ``QuestionAnswer``
    necessarily creates Python/JSON values at the API boundary.
    """
    if temperature <= 0:
        raise ValueError("temperature must be > 0")
    if letter_logits.dim() != 1 or letter_logits.shape[0] != len(LETTERS):
        raise ValueError("expected one [26] letter-logit tensor")

    import torch

    indices = torch.tensor(
        [LETTER_INDEX[option.letter] for option in question.options],
        device=letter_logits.device,
        dtype=torch.long,
    )
    probs = torch.softmax(letter_logits.index_select(0, indices).float() / temperature, dim=0)

    if question.type == "noul":
        return QuestionAnswer(type="noul", noul=float(probs[0].item()))

    if len(question.options) <= 1:
        confidence_value = 1.0
    else:
        entropy = -(probs * torch.where(probs > 0, probs.log(), torch.zeros_like(probs))).sum()
        confidence = 1.0 - entropy / torch.log(torch.tensor(float(len(question.options)), device=probs.device))
        confidence_value = float(torch.clamp(confidence, min=0.0).item())

    if question.type == "choice":
        # Stable GPU argsort preserves the request's option order for exact ties.
        order = torch.argsort(probs, descending=True, stable=True).tolist()
        probability_values = probs.tolist()
        return QuestionAnswer(
            type="choice",
            choice=question.options[order[0]].label,
            probabilities={question.options[index].label: probability_values[index] for index in order},
            confidence=confidence_value,
        )

    probability_values = probs.tolist()
    levels = torch.arange(len(question.options), device=probs.device, dtype=probs.dtype)
    score = float((levels * probs).sum().item())
    return QuestionAnswer(
        type="score",
        score=score,
        probabilities={str(index): value for index, value in enumerate(probability_values)},
        legend={str(index): option.description or "" for index, option in enumerate(question.options)},
        confidence=confidence_value,
    )