def argmax_label(values: list, labels: list[str]) -> str | None:
"""The label with the most probability mass, or ``None`` if any value is missing or
non-numeric -- a partially parsed distribution never yields a confident-looking pick."""
numeric = [_numeric_value(value) for value in values]
if any(value is None for value in numeric):
return None
return labels[int(np.argmax(numeric))]
def argmax_annotation(values: list, labels: list[str]) -> str:
"""The top label plus its probability, formatted for heatmap cell annotations."""
label = argmax_label(values, labels)
if label is None:
return ""
probability = _numeric_value(values[labels.index(label)])
if probability is None:
return label
probability_text = f"{probability:.2f}".removeprefix("0")
return f"{label} {probability_text}"
def _numeric_value(value: object) -> float | None:
"""A finite numeric value, or ``None`` if the model emitted something unusable."""
try:
numeric = float(value)
except (TypeError, ValueError):
return None
return numeric if np.isfinite(numeric) else None
def parse_distribution(raw: object, labels: list[str]) -> list[float]:
"""Map a model's already-parsed per-question reply to per-label probabilities, in label order
(distribution-mode answers left un-normalized).
A single-pick reply is a single label string -> all the mass on that exact label; a
distribution-mode reply is a dict read label by label. Anything that doesn't match a known label
or isn't a finite number is left NaN -- we report the gap rather than massaging the reply (e.g.
stripping an echoed description) to make it fit."""
if isinstance(raw, str): # single-pick mode: a single chosen label
if raw in labels:
return [1.0 if label == raw else 0.0 for label in labels]
return [float("nan")] * len(labels)
if not isinstance(raw, dict):
return [float("nan")] * len(labels)
return [
value if (value := _numeric_value(raw.get(label))) is not None else float("nan")
for label in labels
]
def rubric_prompt(mode: str, sample_index: int, rubric_hash: str) -> str:
"""The post + all questions (with their label sets) in one prompt; ``mode`` picks the format.
``mode="dist"`` asks for a probability distribution over each question's labels; the single-pick
mode (``mode="single"``) asks for a single label per question. The uid line combines
``rubric_hash`` (which rubric version) with ``sample_index`` and a random token, so every repeat
is a distinct, independent draw and two different rubrics never share a nonce."""
lines = []
for key, (instructions, choices) in QUESTIONS.items():
labels = "\n".join(f" {label}: {desc}" for label, desc in choices.items())
lines.append(f"- {key}: {instructions}\n labels:\n{labels}")
exclusivity = (
"\n\nEach question's labels are mutually exclusive: exactly one applies. If a post could "
"arguably fit more than one, pick the single most severe / most specific label per the "
"label descriptions."
)
if mode == "single":
answer_format = (
"\n\nFor each question, pick exactly ONE label.\nRespond with ONLY a JSON object "
"mapping each question's key to one of that question's bare labels (the label only, "
"not its description), with one entry per question."
)
else:
answer_format = (
"\n\nFor each question, give a probability distribution over that question's labels "
"(values 0.00-1.00 that sum to 1).\nRespond with ONLY a JSON object mapping each "
"question's key to an object mapping that question's bare labels (the label only, "
"not its description) to probabilities, with one entry per question."
)
return (
f"uid: {rubric_hash}:{sample_index}:{token_hex(4)}\n\n"
f"Document (a reported user post):\n{json.dumps(POST, indent=2)}\n\nQuestions:\n"
+ "\n".join(lines)
+ exclusivity
+ answer_format
)
def _cost(prices: tuple[float, float], input_tokens: int, output_tokens: int) -> float:
return input_tokens / 1e6 * prices[0] + output_tokens / 1e6 * prices[1]
def _call_llm(model: str, prompt: str, temperature: float | None):
"""One LLM call -> (text, cost_usd, latency_s), routed by model name."""
reasoning = model in REASONING_MODELS
started = perf_counter()
if model.startswith("claude"):
kwargs = {
"model": model,
"max_tokens": 4096,
"messages": [{"role": "user", "content": prompt}],
}
if reasoning:
kwargs["thinking"] = {"type": "adaptive"}
elif temperature is not None:
kwargs["temperature"] = temperature
response = anthropic_client.messages.create(**kwargs)
text = next((b.text for b in response.content if b.type == "text"), "")
usage = (response.usage.input_tokens, response.usage.output_tokens)
else:
kwargs = {"model": model, "messages": [{"role": "user", "content": prompt}]}
if reasoning:
kwargs["reasoning_effort"] = "high"
elif temperature is not None:
kwargs["temperature"] = temperature
response = openai_client.chat.completions.create(**kwargs)
text = response.choices[0].message.content
usage = (response.usage.prompt_tokens, response.usage.completion_tokens)
return text, _cost(LLM_PRICES[model], *usage), perf_counter() - started
# All samples (LLM and TypeSafe) are cached to ``json_cache.json``, which ships with the cookbook, so
# re-rendering is instant and reproduces the published numbers with no API spend. ``sample_index``
# seeds the uid buster and is part of the cache key, so each of the NUM_SAMPLES repeats is its own
# entry and its own independent draw, not one draw replayed. Delete ``json_cache.json`` to re-sample
# everything live.
json_cache = JsonCache(Path("json_cache.json"))
def _rubric_fingerprint() -> str:
"""Short digest of everything that shapes the prompt/rubric: the state and every question's
text and label set. Passed into the cached calls below so that editing the post or any question
changes the cache key and forces a fresh sample, instead of silently serving a stale answer that
was generated for the old wording."""
payload = json.dumps([POST, QUESTIONS], sort_keys=True, default=str)
return hashlib.sha256(payload.encode()).hexdigest()[:12]
RUBRIC_HASH = _rubric_fingerprint()
@json_cache
def _call_typesafe(sample_index: int, rubric_hash: str):
"""Return distributions, token usage, and latency for one TypeSafe rubric call.
``rubric_hash`` invalidates samples after rubric changes.
"""
questions = {
key: Choice(instructions=instructions, criteria=choices)
for key, (instructions, choices) in QUESTIONS.items()
}
started = perf_counter()
response = client.system_one(
state={"uid": f"{rubric_hash}:{sample_index}:{token_hex(4)}", "post": POST},
questions=questions,
model=TYPESAFE_MODEL,
)
distributions = {}
for key, (_instructions, choices) in QUESTIONS.items():
probabilities = dict(response.answers[key].probabilities)
distributions[key] = [
probabilities.get(label, float("nan")) for label in choices
]
return (
distributions,
response.usage.input_tokens,
response.usage.output_tokens,
perf_counter() - started,
)
@json_cache
def ask_llm_rubric(
model: str,
mode: str,
temperature: float | None,
sample_index: int,
rubric_hash: str,
):
"""One LLM rubric query -> (per-question label distributions keyed by question key, cost_usd,
latency_s); NaNs if the reply doesn't parse.
``mode="dist"`` parses 8 label distributions; the single-pick mode (``mode="single"``) parses 8
single labels and puts all the mass on each. ``rubric_hash`` goes into the prompt's uid nonce
(and so the cache key), so an edited document/rubric busts the cache instead of serving a stale
answer."""
prompt = rubric_prompt(mode, sample_index, rubric_hash)
text, cost, latency = _call_llm(model, prompt, temperature)
# Peel a single ```json ... ``` fence (claude-haiku-4-5 sometimes adds one despite "ONLY a JSON
# object").
stripped = text.strip()
if stripped.startswith("```"):
stripped = stripped[stripped.find("\n") + 1 :] if "\n" in stripped else ""
if stripped.rstrip().endswith("```"):
stripped = stripped.rstrip()[: -len("```")]
try:
raw = json.loads(stripped)
except (ValueError, json.JSONDecodeError):
raw = {}
if not isinstance(raw, dict):
raw = {}
distributions = {
key: parse_distribution(raw.get(key), list(choices))
for key, (_instructions, choices) in QUESTIONS.items()
}
return distributions, cost, latency
