"""Format WitGymResponse traces for Gradio — styled HTML transcript."""
import html
import json
import re
from dataclasses import dataclass, field
from typing import List, Tuple, Any, Optional
from witgym.schemas import WitGymResponse, PipelineEvent, ComedyMetadata, CandidateResponse, TranscriptScene
from witgym.avatars import char_avatar_url
# Maps character first-name or full-name → title string
_CHAR_TITLES: dict[str, str] = {
"michael": "Regional Manager",
"dwight": "Assistant (to the) Reg. Manager",
"jim": "Sales Representative",
"pam": "Receptionist",
"kevin": "Accountant",
"andy": "Sales Representative",
"stanley": "Sales Representative",
"angela": "Head of Accounting",
"ryan": "Temp → VP → Temp",
"kelly": "Customer Service Rep",
"michael scott": "Regional Manager",
"dwight schrute": "Asst. (to the) Reg. Manager",
"jim halpert": "Sales Representative",
"pam beesly": "Receptionist",
"kevin malone": "Accountant",
"andy bernard": "Sales Representative",
"stanley hudson": "Sales Representative",
"angela martin": "Head of Accounting",
"ryan howard": "Temp → VP → Temp",
"kelly kapoor": "Customer Service Rep",
}
def _avatar_url(character: str) -> str:
return char_avatar_url(character)
def _char_title(character: str) -> str:
key = character.lower().strip()
return _CHAR_TITLES.get(key, "The Office")
def _esc(text: str) -> str:
return html.escape(text or "")
def _fmt_chip(value: str) -> str:
return value.replace("_", " ").upper()
# Succinct definitions for each comedy structural property — shown on chip click
_CHIP_DEFS: dict[str, tuple[str, str]] = {
# Archetypes
"status_assertion": ("ARCHETYPE · STATUS ASSERTION", "What it is: someone is trying to claim authority, status, or competence. Why it matters: the joke works by exposing the gap between the claim and the reality."),
"self_delusion": ("ARCHETYPE · SELF-DELUSION", "What it is: someone believes a flattering story about themselves that reality cannot support. Why it matters: comedy lands when the fantasy finally runs into the facts."),
"power_inversion": ("ARCHETYPE · POWER INVERSION", "What it is: the expected pecking order flips. Why it matters: surprise grows when the person who should control the room suddenly loses it."),
"anxiety_escalation": ("ARCHETYPE · ANXIETY ESCALATION", "What it is: a small worry spirals into a full disaster movie. Why it matters: the joke gets sharper as the reaction becomes bigger than the actual problem."),
"social_fail": ("ARCHETYPE · SOCIAL PERFORMANCE FAIL", "What it is: someone tries to look normal, smooth, or impressive and fumbles it. Why it matters: comedy comes from watching the performance collapse in public."),
"misplaced_conf": ("ARCHETYPE · MISPLACED CONFIDENCE", "What it is: someone sounds certain without having the facts. Why it matters: the joke exposes confidence that has outrun competence."),
# Tensions
"social_embarrass": ("TENSION · SOCIAL EMBARRASSMENT", "What it is: the fear of looking foolish in front of other people. Why it matters: embarrassment raises the emotional stakes, so even a short line can hit hard."),
"existential": ("TENSION · EXISTENTIAL ANXIETY", "What it is: the moment quietly touches identity, meaning, or dread. Why it matters: deeper fear gives the joke weight instead of making it feel throwaway."),
"status_threat": ("TENSION · STATUS THREAT", "What it is: someone's rank, credibility, or belonging feels under pressure. Why it matters: status danger creates friction, and friction gives the line bite."),
"identity_expose": ("TENSION · IDENTITY EXPOSURE", "What it is: the person's mask is slipping. Why it matters: jokes get stronger when they reveal the truth the speaker was trying to hide."),
"logic_collapse": ("TENSION · LOGIC COLLAPSE", "What it is: the person's explanation breaks under its own rules. Why it matters: once the logic collapses, the line can simply point at the wreckage."),
# Distances
"mild": ("SHARPNESS · MILD VIOLATION", "What it is: a safe, lightly subversive joke. Why it matters: this keeps the line playful instead of confrontational."),
"moderate": ("SHARPNESS · MODERATE VIOLATION", "What it is: a joke with some edge, but still recoverable. Why it matters: this is often the sweet spot for sounding bold without losing the room."),
"sharp": ("SHARPNESS · SHARP VIOLATION", "What it is: a line that cuts close to the bone. Why it matters: sharper jokes can win bigger laughs, but they also raise the risk of silence."),
}
_PERSONA_DEFS: dict[str, tuple[str, str]] = {
"cynic": ("PERSONA · CYNIC", "This lens says the quiet part out loud. It spots the real motive, the hidden cost, or the ugly truth underneath the situation."),
"conviction": ("PERSONA · CONVICTION", "This lens commits completely to a wrong belief. It is funny because the certainty exposes something true about the speaker."),
"absurdist": ("PERSONA · ABSURDIST", "This lens follows the situation's logic farther than a normal person would. It makes the joke by treating the spiral as perfectly reasonable."),
"bisociate": ("PERSONA · BISOCIATE", "This lens jumps to a different but structurally matching world. It works when the same comic pattern suddenly shows up somewhere unexpected."),
}
def _jstr(text: str) -> str:
"""JSON-encode a string for safe embedding in an HTML onclick attribute."""
return json.dumps(text)
_MODE_BADGES = {
"banter": '⚡ BANTER',
"quick_wit": '🎯 QUICK WIT',
"coaching": '🎓 COACHING',
"humour": '🎯 QUICK WIT',
"smalltalk": '⚡ BANTER',
}
def _mode_badge_html(route: str) -> str:
return _MODE_BADGES.get(route, _MODE_BADGES["quick_wit"])
def _coach_header_html(selected_char: str) -> str:
"""Render 'JIM SAYS:' header with tiny avatar when a character is selected."""
if not selected_char or selected_char == "AI":
return '
Your humor coach
'
av = _avatar_url(selected_char)
name = selected_char.upper()
return (
f'
'
f''
f'{_esc(name)} SAYS:'
f'
'
)
def _reply_actions_html() -> str:
return (
'
'
''
'
'
)
def _trace_payload_from_result(result: WitGymResponse) -> dict[str, Any]:
metadata = result.metadata.model_dump(mode="json")
scenes = [scene.model_dump(mode="json") for scene in result.retrieved_scenes]
candidates = [candidate.model_dump(mode="json") for candidate in result.candidates]
logs = [
{"step": "metadata", "status": "ok", "detail": f"twist={metadata.get('twist_potential')} archetype={metadata.get('archetype')}"},
{"step": "retrieval", "status": "ok", "detail": ", ".join(f"{s['character']}:{s['archetype']}" for s in scenes) or "no precedent scenes"},
{"step": "candidate_generation", "status": "ok", "detail": ", ".join(f"{c['persona']}:{len(c['text'].split())}w" for c in candidates) or "no candidates"},
{"step": "ranking", "status": "ok", "detail": result.winning_persona or "none"},
{"step": "compression", "status": "ok", "detail": "selected line finalized"},
]
return {
"route": result.route,
"winning_persona": result.winning_persona,
"selected": result.selected,
"metadata": metadata,
"retrieved_scenes": scenes,
"candidates": candidates,
"explanation": result.explanation,
"logs": logs,
}
def _trace_payload_from_stream(state: "StreamingTurnState", selected_text: str) -> dict[str, Any]:
metadata = state.metadata.model_dump(mode="json") if state.metadata else None
scenes = [scene.model_dump(mode="json") for scene in state.scenes]
candidates = [candidate.model_dump(mode="json") for candidate in state.candidates]
logs = []
if metadata:
logs.append({"step": "metadata", "status": "ok", "detail": f"twist={metadata.get('twist_potential')} archetype={metadata.get('archetype')}"})
if scenes:
logs.append({"step": "retrieval", "status": "ok", "detail": ", ".join(f"{s['character']}:{s['archetype']}" for s in scenes)})
if candidates:
logs.append({"step": "candidate_generation", "status": "ok", "detail": ", ".join(f"{c['persona']}:{len(c['text'].split())}w" for c in candidates)})
if state.winning_persona:
logs.append({"step": "ranking", "status": "ok", "detail": state.winning_persona})
if state.streaming_final or state.final_text:
logs.append({"step": "compression", "status": "running" if state.streaming_final else "ok", "detail": "polishing final line" if state.streaming_final else "selected line finalized"})
return {
"route": state.route,
"winning_persona": state.winning_persona,
"selected": selected_text or state.selected,
"final_text": state.final_text or None,
"metadata": metadata,
"retrieved_scenes": scenes,
"candidates": candidates,
"active_candidate": (
{"persona": state.active_persona, "partial_text": state.active_candidate_text}
if state.active_persona and state.active_candidate_text else None
),
"streaming_final": state.streaming_final,
"logs": logs,
}
def _compact_reply_html(route: str, selected: str, *, coaching_hint: str = "", selected_char: str = "AI") -> str:
hint = f'
{_esc(coaching_hint)}
' if coaching_hint else ""
return (
f'{_mode_badge_html(route)}'
f'
'
f'{_reply_actions_html()}'
f'{_coach_header_html(selected_char)}'
f'
'
)
_THINKING_ICON = (
''
)
_STABLE_LOADING_LINE = "Checking Dunder Mifflin playbook…"
def _sanitize_streaming_reply_text(text: str) -> str:
raw = (text or "").strip()
if not raw:
return ""
json_match = re.search(r'"compressed_line"\s*:\s*"((?:[^"\\]|\\.)*)"', raw)
if json_match:
try:
return json.loads(f'"{json_match.group(1)}"').strip()
except Exception:
return json_match.group(1).strip()
if raw.startswith("{"):
try:
obj = json.loads(raw)
value = next(iter(obj.values())) if obj else ""
return str(value).strip()
except Exception:
return ""
return raw
def thinking_turn_html(user_input: str) -> str:
return (
'
'
f'
You {_esc(user_input)}
'
f'
{_THINKING_ICON}'
f'{_esc(_STABLE_LOADING_LINE)}
'
'
'
)
def _chip_onclick(value: str) -> str:
title, defn = _CHIP_DEFS.get(value, (_fmt_chip(value), "A comedy structural property."))
return f"wgOpenChip({_jstr(title)},{_jstr(defn)})"
def _persona_onclick(value: str) -> str:
title, defn = _PERSONA_DEFS.get(value, (f"PERSONA · {_fmt_chip(value)}", "This is the comic lens the coach used to write the line."))
return f"wgOpenChip({_jstr(title)},{_jstr(defn)})"
def _insight_button(label: str, onclick: str) -> str:
return (
f''
)
def _coach_notes_html(result: WitGymResponse) -> str:
meta = result.metadata
buttons = [
_insight_button("situation pattern", _chip_onclick(meta.archetype.value)),
_insight_button("social pressure", _chip_onclick(meta.tension_type.value)),
_insight_button(
f"complexity {meta.twist_potential}/10",
f"wgOpenChip({_jstr('COMEDY COMPLEXITY')},{_jstr('What it is: how many moving parts this moment has. Why it matters: lower scores usually want one clean observation, while higher scores can support a twistier line.')})",
),
_insight_button(
f"retrieved context · {len(result.retrieved_scenes)}",
f"wgOpenChip({_jstr('RETRIEVED CONTEXT')},{_jstr('These are Office scenes with a similar comedy pattern. They are reference cases for how the joke works, not lines the coach is copying.')})",
),
]
if result.winning_persona:
buttons.append(_insight_button("persona lens", _persona_onclick(result.winning_persona)))
return (
'
'
'
How WitGym built this line
'
'
Tap any note for a quick explanation.
'
'
'
+ "".join(buttons) +
'
'
'
'
)
def _meta_pass1_html(meta) -> str:
arc_click = _chip_onclick(meta.archetype.value)
ten_click = _chip_onclick(meta.tension_type.value)
vio_click = _chip_onclick(meta.violation_distance.value)
_connector_onclick = "wgOpenChip(" + _jstr("CONNECTOR") + "," + _jstr("A word or phrase with two simultaneous readings. When it lands, both meanings hit at once — that’s the mechanism.") + ")"
connector_chip = (
f'
',
]
if state.route in ("banter", "smalltalk") and state.selected:
parts += [
_compact_reply_html("banter", state.selected),
'
',
]
return "".join(parts)
if state.route == "coaching" and state.selected and state.metadata is None:
parts += [
_compact_reply_html(
"coaching",
state.selected,
coaching_hint="coaching mode — waiting for your answer",
),
'',
]
return "".join(parts)
if state.metadata is None:
parts += [
f'
{_THINKING_ICON}{_esc(_STABLE_LOADING_LINE)}
',
]
return "".join(parts)
collapsed_cls = "" if show_debug else " wg-collapsed"
chevron = "▼" if show_debug else "▶"
display_text = _sanitize_streaming_reply_text(state.final_text or state.selected)
persona_label = (
f' · {_esc(state.winning_persona)}'
if state.winning_persona else ""
)
# Guard: don't display raw JSON tokens or mixed text+JSON tails from compression streaming
raw_display_text = (state.final_text or state.selected or "").strip()
has_json_tail = '"compressed_line"' in raw_display_text or (
"{" in raw_display_text and not raw_display_text.lstrip().startswith("{")
)
is_raw_json = bool(raw_display_text and raw_display_text.lstrip().startswith("{"))
# ── Hero reply FIRST (above fold) ──────────────────────────────────────
parts.append(_mode_badge_html(state.route))
if display_text and not is_raw_json and not has_json_tail:
polish = ' · polishing…' if state.streaming_final else ""
parts += [
'
',
f'
Your humor coach{persona_label}{polish}
',
f'
{_esc(display_text)}
',
'
',
]
elif state.streaming_final or is_raw_json or has_json_tail:
# Compression pass running — show stable placeholder
parts += [
'
',
f'
Your humor coach{persona_label} · polishing…
',
f'
Finding the sharpest version…
',
'
',
]
elif not state.active_persona:
parts += [
f'
{_THINKING_ICON}'
'drafting candidates…
',
]
parts.append('')
return "".join(parts)
def format_transcript_with_streaming(
traces: List[Tuple[str, Any]],
stream_state: Optional[StreamingTurnState],
show_debug: bool = True,
max_turns: int = 5,
selected_char: str = "AI",
) -> str:
if not traces and stream_state is None:
return format_transcript_html([], show_debug=show_debug, selected_char=selected_char)
recent = traces[-max_turns:]
body = "".join(
format_trace_html(
r if isinstance(r, WitGymResponse) else WitGymResponse.model_validate(r),
user_input,
show_debug=show_debug,
is_last=False,
selected_char=selected_char,
)
for user_input, r in recent
)
if stream_state is not None:
body += format_streaming_turn_html(stream_state, show_debug=show_debug)
if stream_state is not None and stream_state.metadata is not None:
trace_payload = _trace_payload_from_stream(stream_state, stream_state.selected)
elif recent:
last_result = recent[-1][1] if isinstance(recent[-1][1], WitGymResponse) else WitGymResponse.model_validate(recent[-1][1])
trace_payload = _trace_payload_from_result(last_result)
else:
trace_payload = None
trace_attr = _esc(json.dumps(trace_payload, ensure_ascii=True)) if trace_payload else ""
return f'
',
]
if result.route in ("banter", "smalltalk"):
parts += [
_compact_reply_html("banter", result.selected, selected_char=selected_char),
'
',
]
return "".join(parts)
if result.route == "coaching" and result.coaching_question and not result.candidates:
parts += [
_compact_reply_html(
"coaching",
result.selected,
coaching_hint="coaching mode — waiting for your answer",
selected_char=selected_char,
),
'',
]
return "".join(parts)
collapsed_cls = "" if show_debug else " wg-collapsed"
chevron = "▼" if show_debug else "▶"
# Alternate candidates (non-selected) for client-side "another take"
alts = [{"persona": c.persona, "text": c.text} for c in result.candidates if c.text != result.selected]
alts_json = html.escape(json.dumps(alts))
# Winning persona label — use pre-compression value stored in response
winning_persona = result.winning_persona
persona_label = (
f' · '
if winning_persona else ""
)
# Another take button (only if alternatives exist)
another_take_btn = (
'↻ another take'
if alts else ""
)
# New-turn reveal class for animation
new_cls = " wg-coach-reply--new" if is_last else ""
beckon_cls = " wg-debug-toggle--new" if is_last else ""
# ── Hero reply FIRST (above fold) ──────────────────────────────────────
coach_hdr = (
f'
'
+ (f''
if selected_char and selected_char != "AI" else "")
+ (f'{_esc(selected_char.upper())} SAYS:' if selected_char and selected_char != "AI" else f'Your humor coach{persona_label}')
+ f'{another_take_btn}
'
)
parts += [
_mode_badge_html(result.route),
(
f'
'
),
f'{_reply_actions_html()}',
coach_hdr,
f'
{_esc(result.selected)}
',
'
',
]
if result.explanation:
parts.append(_explanation_panel_html(result.explanation))
parts.append(_coach_notes_html(result))
# ── Drill chips (coaching follow-ups) ──────────────────────────────────
if is_last:
parts.append(
'
'
'sharpen it →'
'different angle →'
'why it works →'
'
'
)
parts.append('')
return "".join(parts)
# All JS is in _GLOBAL_JS (app.py, injected via head=) — scripts in gr.HTML value
# are not executed by Gradio 6.x (set via innerHTML). See app.py _GLOBAL_JS.
_PAGE_JS = ""
def format_transcript_html(
traces: List[Tuple[str, Any]],
max_turns: int = 5,
append_html: str = "",
show_debug: bool = True,
selected_char: str = "AI",
) -> str:
if not traces and not append_html:
body = (
'
'
'
🎤
'
'
Drop the situation.
'
'
Your coach will find the line that lands.
'
'
Awkward, delusional, or painfully relatable — all valid.
'
'
'
)
else:
recent = traces[-max_turns:]
if recent:
last_trace_payload = _trace_payload_from_result(
recent[-1][1] if isinstance(recent[-1][1], WitGymResponse) else WitGymResponse.model_validate(recent[-1][1])
)
trace_attr = _esc(json.dumps(last_trace_payload, ensure_ascii=True))
else:
trace_attr = ""
body = "".join(
format_trace_html(
r if isinstance(r, WitGymResponse) else WitGymResponse.model_validate(r),
user_input,
show_debug=show_debug,
is_last=(i == len(recent) - 1),
selected_char=selected_char,
)
for i, (user_input, r) in enumerate(recent)
) + append_html
return f'
{body}
{_PAGE_JS}'
return f'
{body}
{_PAGE_JS}'
# Legacy helpers
def format_trace(result: WitGymResponse, user_input: str) -> str:
lines = [f"**You:** {user_input}", f"**WitGym:** {result.selected}", ""]
if result.route not in ("banter", "smalltalk"):
meta = result.metadata
lines.append(f"_archetype={meta.archetype.value}, tension={meta.tension_type.value}_")
return "\n".join(lines)
def format_logs(traces: List[Tuple[str, Any]], max_turns: int = 5) -> str:
return format_transcript_html(traces, max_turns)