arch v2

witgym/config.py

@@ -19,22 +19,36 @@ DTYPE = torch.bfloat16
 # Sampling
 EXTRACT_TEMP = 0.2
 EXTRACT_DO_SAMPLE = False
-GENERATE_TEMP = 1.3
+GENERATE_TEMP = 1.15  # lowered after anchor/retrieval fixes to reduce domain invention
 GENERATE_MIN_P = 0.06
-GENERATE_MAX_NEW_TOKENS = 55
+GENERATE_MAX_NEW_TOKENS = 40
 RANK_TEMP = 0.1
 EXTRACT_MAX_NEW_TOKENS = 220   # JSON never exceeds ~150 tokens; 220 = 47% safety margin
 
 # Context management
-CONTEXT_WINDOW = 32768        # Conservative for MPS memory
+CONTEXT_WINDOW = 4096         # MPS-friendly; triggers earlier compression if history grows
 COMPRESSION_THRESHOLD = 0.80  # Compress at 80% full
 KEEP_LAST_N_TURNS = 4
 
 # RAG
 TOP_K_SCENES = 2
+RETRIEVE_POOL_SIZE = 12
+ENABLE_CROSS_ENCODER_RERANK = True
+RERANK_MODEL_ID = "cross-encoder/ettin-reranker-32m-v1"
+RERANK_DEVICE = "cpu"  # keep 9B on MPS; rerank on CPU to avoid unified-memory spikes
 INDEX_PATH = "data/index.json"
 TRANSCRIPT_DIR = "data/transcripts"
 
+# Generation guards
+ENABLE_BAD_WORD_GUARD = True
+BAD_WORD_PHRASES = (
+    "PATH A",
+    "PATH B",
+    "PATHWAY SELECTION",
+)
+ENABLE_OVERLAP_GUARD = True
+OVERLAP_NGRAM_SIZE = 6  # contiguous words shared with retrieved dialogue
+
 # Cliché penalty (soft steering, not hard suppression)
 CLICHE_LOGIT_PENALTY = -5.0
 CLICHE_PENALTY_TOKENS = 6   # Penalise first N tokens of the obvious response

witgym/conversation.py

@@ -20,29 +20,33 @@ class ConversationManager:
         self.history: List[Tuple[str, str]] = []   # (user, assistant)
         self.used_archetypes: Set[ComedyArchetype] = set()
         self._summary: str = ""  # Compressed summary of old turns
-        self._mechanisms: List[Tuple[str, str, str]] = []  # (archetype, tension, distance)
+        self._mechanisms: List[Tuple[str, str, str, str, str]] = []  # (user_input, subtext, archetype, tension, distance)
 
     def add_turn(self, user_input: str, response: str, metadata: ComedyMetadata):
         self.history.append((user_input, response))
         self.used_archetypes.add(metadata.archetype)
         self._mechanisms.append((
+            user_input,
+            metadata.subtext,
             metadata.archetype.value,
             metadata.tension_type.value,
             metadata.violation_distance.value,
         ))
 
     def get_context_string(self) -> str:
-        """Return the last N turns as mechanism-only context.
+        """Return the last N turns as mechanism-only context (archetype + tension).
 
-        Structural goal: preserve callback potential (what kind of jokes were made),
-        while avoiding lexical bleed from raw user/assistant text across turns.
+        Mechanism-only context for callbacks — no prior topic text. user_input and subtext are
+        stored in _mechanisms but intentionally not exposed here — topic contamination
+        makes turn-level evaluation unreliable and anchors the current joke to prior topics.
+        Re-enable richer output when multi-turn callback quality is validated.
         """
         recent = self._mechanisms[-config.KEEP_LAST_N_TURNS:]
         lines = []
         if self._summary:
             lines.append(f"[Earlier conversation summary]: {self._summary}")
-        for i, (arch, tension, dist) in enumerate(recent, 1):
-            lines.append(f"Turn -{len(recent) - i + 1}: archetype={arch}, tension={tension}, distance={dist}")
+        for i, (user_in, subtext, arch, tension, dist) in enumerate(recent, 1):
+            lines.append(f"Turn -{len(recent) - i + 1}: archetype={arch}, tension={tension}")
         return "\n".join(lines)
 
     def needs_compression(self, tokenizer) -> bool:

witgym/engine.py

@@ -62,28 +62,31 @@ class WitGymEngine:
         scenes = retrieve_scenes(
             self.index,
             metadata,
-            self.conversation.used_archetypes,
             self.embed_model,
         )
 
-        # PASS 2 — Generate persona candidates.
-        # Medium inputs (4-6): run cynic + absurdist only.
-        # Rich inputs (> 6): run all three including conviction.
+        # PASS 2 — persona selection by twist_potential
+        # ≤6: light inputs — cynic + absurdist only (conviction needs richer tension)
+        # 7-8: full standard set — cynic, conviction, absurdist
+        # >8: peak inputs — bisociate replaces conviction (domain-pivot vs. within-frame logic)
         context_str = self.conversation.get_context_string()
-        personas_to_run = None  # None = all three
+        personas_to_run = ["cynic", "conviction", "absurdist"]  # default: standard 3
         if metadata.twist_potential <= 6:
-            personas_to_run = ["cynic", "absurdist"]  # Conviction needs richer tension to land cleanly
-            logger.info(f"twist_potential={metadata.twist_potential} ≤ 6 — skipping conviction")
+            personas_to_run = ["cynic", "absurdist"]
+            logger.info(f"twist_potential={metadata.twist_potential} ≤ 6 — cynic + absurdist only")
+        elif metadata.twist_potential > 8:
+            personas_to_run = ["cynic", "absurdist", "bisociate"]
+            logger.info(f"twist_potential={metadata.twist_potential} > 8 — bisociate replaces conviction")
 
         candidates = generate_candidates(
             user_input, metadata, scenes,
             self.model, self.tokenizer,
-            context_str, self.conversation.used_archetypes,
+            context_str,
             personas_to_run=personas_to_run,
         )
 
         # RANK — Pick best candidate
-        selected = rank_candidates(user_input, candidates, self.model, self.tokenizer)
+        selected = rank_candidates(user_input, metadata, candidates, self.model, self.tokenizer)
 
         # COMPRESS — Swartzwelder pass: generate loose, cut ruthless (skips if ≤12 words)
         selected = compress_winner(selected, self.model, self.tokenizer)

witgym/extractor.py

@@ -15,11 +15,14 @@ Return ONLY a JSON object with these exact fields (no explanation, no markdown,
   "surface": "what was literally said in one sentence",
   "subtext": "what the speaker actually means or feels",
   "archetype": one of ["status_assertion", "self_delusion", "power_inversion", "anxiety_escalation", "social_fail", "misplaced_conf"],
+  "archetype_confidence": an integer from 1 to 10 (how confident you are in the archetype choice),
   "tension_type": one of ["social_embarrass", "existential", "status_threat", "identity_expose", "logic_collapse"],
   "power_dynamic": "who has power and who doesn't, one sentence",
+  "speaker_strategy": "one short phrase describing how the speaker is trying to be perceived (e.g. competent, unbothered, in-control), or null if unclear",
   "obvious_response": "the most boring, expected response to this input",
   "violation_distance": one of ["mild", "moderate", "sharp"],
-  "twist_potential": an integer from 1 to 10 rating how much hidden comedy tension is in this input (1=completely flat, 10=extremely rich setup for wit)
+  "twist_potential": an integer from 1 to 10 rating how much hidden comedy tension is in this input (1=completely flat, 10=extremely rich setup for wit),
+  "connector": "the specific word or phrase in the input that could mean two different things simultaneously, or null if no such word exists"
 }}
 
 Think carefully about the ARCHETYPE — pick the one that most accurately describes the comedy mechanism hiding in this input.
@@ -32,6 +35,7 @@ Archetype selection guidance (avoid overusing self_delusion):
 - self_delusion: specifically a self-image story ("I'm fine / I'm great / I'm the best") contradicted by behavior/evidence in the same moment
 If unsure between self_delusion vs social_fail/anxiety_escalation, prefer the more specific one (social_fail/anxiety_escalation) unless there is an explicit self-image contradiction.
 For twist_potential: score high if the input has self-delusion, status gap, or absurd logic. Score low if it is a neutral factual statement with no tension.
+For connector: look for a single word or short phrase that carries an expected meaning in context AND a second meaning that reframes the situation. Most inputs will have null. Return null unless a genuine dual-reading exists (e.g. "manage" can mean control people or barely cope; "balance" can mean financial or emotional equilibrium).
 Return ONLY the JSON. Nothing else."""
 
 

witgym/generator.py

@@ -6,19 +6,27 @@ from typing import List, Set
 from collections import Counter
 import re as _re
 from loguru import logger
-from transformers import LogitsProcessorList
+from transformers import LogitsProcessorList, NoBadWordsLogitsProcessor
 from witgym.model import generate_text, ClichePenaltyProcessor
+from witgym import config
 from witgym.schemas import (
-    ComedyMetadata, TranscriptScene, CandidateResponse, ComedyArchetype
+    ComedyMetadata, TranscriptScene, CandidateResponse,
 )
 
 PERSONA_INSTRUCTIONS = {
     "cynic": (
+        "PATHWAY SELECTION — pick one silently before generating:\n"
+        "  PATH A (subtext contains avoidance, spiral, or self-protection): "
+        "Name the specific function the behavior is performing. "
+        "What is it protecting? What does it cost? State it precisely.\n"
+        "  PATH B (subtext contains status claim or competence assertion): "
+        "State the actual outcome as if it has already become concrete. Past tense. "
+        "One specific image or action anchored in the user's situation.\n"
+        "You pick the path silently. Output only the wit line — never the path label.\n\n"
         "You've seen this exact rationalization a hundred times. "
-        "You're not angry about it — you're just tired. "
-        "State what's actually happening with the weariness of someone who has been right about this for years. "
-        "Open with the situation, the consequence, or the outcome as your subject — "
-        "lead with something concrete and specific, not a generic observation about the person."
+        "You're not angry — you're just tired. "
+        "Lead with something concrete: the situation, the consequence, the outcome. "
+        "Not a generic observation about the person."
     ),
     "conviction": (
         "You have a firm, specific belief about how this situation works. "
@@ -28,24 +36,49 @@ PERSONA_INSTRUCTIONS = {
         "No irony. No wink. Absolute conviction."
     ),
     "absurdist": (
+        "PATHWAY SELECTION — pick one silently before generating:\n"
+        "  PATH A (subtext contains fear, avoidance, or escalation logic — anxiety_escalation or existential tension): "
+        "Follow the anxiety's own internal logic to its physical inevitable endpoint. "
+        "Name the specific form: a concrete image, measurable action, or physical object "
+        "from the user's situation.\n"
+        "  PATH B (connector field is non-null OR input has subject\u2192action\u2192object structure): "
+        "Keep the verb/action. Find the domain where that exact action applied to a different object "
+        "reveals the identical human truth. Deliver from that domain without naming the parallel.\n"
+        "You pick the path silently. Output only the wit line — never the path label.\n\n"
         "You're the only person in the room who sees where this logically ends. "
         "You're not being weird — you're just following the math. "
-        "State the inevitable conclusion with the calm of someone reading from a manual. "
-        "End with a single, specific concrete image that makes the conclusion visible — "
-        "not an abstraction, not a restatement. A physical object, a measurable action, a named institution."
+        "End with a single concrete image anchored in the user's situation. "
+        "Not an abstraction. A physical object or measurable action the input already implies."
+    ),
+    "bisociate": (
+        "PATHWAY SELECTION — pick one silently before generating:\n"
+        "  PATH A (connector field is non-null): "
+        "Your punchline must land on the second meaning of the connector word. "
+        "The setup's expected reading of that word is the straight path. You take the other one.\n"
+        "  PATH B (connector is null): "
+        "Apply the dominant verb from the subtext to the most structurally incongruous object from daily life. "
+        "Deliver from inside that object's world, as if you've been there all along. "
+        "Do not name the parallel. Do not explain the connection.\n"
+        "You pick the path silently. Output only the wit line — never the path label.\n\n"
+        "You notice that what they're describing is not unique to their situation. "
+        "The same need, avoidance, or craving exists somewhere completely different. "
+        "State the parallel as if it is obvious and everyone already knows this. "
+        "Do not reference any topic already mentioned in the conversation context."
     ),
 }
 
 GENERATION_PROMPT = """\
-You are a sharp, brief conversational wit engine.
-You are responding to: "{user_input}"
+You're in the writer's room. Someone just said: "{user_input}" — what's the line?
 
 SITUATION ANALYSIS:
 - What's happening: {surface}
 - What they really mean: {subtext}
 - The comedy mechanism: {archetype}
+- Archetype confidence: {archetype_confidence}/10
 - The tension: {tension_type}
 - Power dynamic: {power_dynamic}
+- Speaker strategy: {speaker_strategy}
+- Connector word (two readings, if present): {connector}
 
 HUMAN COMEDY PRECEDENT (structurally similar situations — same violation type, NOT same words):
 {scenes_block}
@@ -61,8 +94,11 @@ CONSTRAINTS (ALL must be satisfied):
 5. Lead with the punchline. Do not build up to it.
 6. No preamble. No "Here's a response:". No hedging. Start with the wit directly.
 7. Stay benign. The violation must be recognizable, not offensive.
-8. Do NOT use these violation types already used in this conversation: {used_archetypes_str}
-9. Suppress this boring response style: "{obvious_response}"
+8. Suppress this boring response style: "{obvious_response}"
+9. NEVER output "PATH A", "PATH B", "PATHWAY SELECTION", or any routing label. Output ONLY the comedy line.
+10. If any phrase from the precedent block appears in your line, rewrite it — use the mechanism only, never the precedent's wording.
+11. DOMAIN ANCHOR — Do not introduce institutions, departments, legal process, or workplace nouns unless already implied by surface, subtext, or power_dynamic. Prefer nouns from the user input.
+12. Treat unexpected domain shifts (HR, litigation, supply chain, middle-manager hierarchy) as defects unless the input already implies them.
 
 CONVERSATION CONTEXT (last turns, for callbacks):
 {context_str}
@@ -70,14 +106,46 @@ CONVERSATION CONTEXT (last turns, for callbacks):
 Respond now with ONE or TWO sentences. Nothing else."""
 
 
+def _normalize_for_overlap(text: str) -> str:
+    """Lowercase, strip punctuation, collapse whitespace for n-gram checks."""
+    text = text.lower()
+    text = _re.sub(r"[^\w\s]", " ", text)
+    return " ".join(text.split())
+
+
+def _ngrams(words: List[str], n: int) -> Set[str]:
+    if len(words) < n:
+        return set()
+    return {" ".join(words[i : i + n]) for i in range(len(words) - n + 1)}
+
+
+def _precedent_prompt_texts(scenes: List[TranscriptScene]) -> List[str]:
+    """Text actually injected into the generation prompt (overlap guard surface)."""
+    return [s.why_it_works for s in scenes]
+
+
+def _copies_retrieved_phrase(candidate: str, scenes: List[TranscriptScene], n: int) -> bool:
+    """True if candidate shares an n-word contiguous phrase with precedent prompt text."""
+    cand_words = _normalize_for_overlap(candidate).split()
+    if len(cand_words) < n:
+        return False
+    cand_grams = _ngrams(cand_words, n)
+    if not cand_grams:
+        return False
+    for source in _precedent_prompt_texts(scenes):
+        src_grams = _ngrams(_normalize_for_overlap(source).split(), n)
+        if cand_grams & src_grams:
+            return True
+    return False
+
 
 def _build_scenes_block(scenes: List[TranscriptScene]) -> str:
+    """Mechanism-only precedent — tags + why_it_works, no raw setup or dialogue."""
     blocks = []
     for s in scenes:
         blocks.append(
-            f"  Show: {s.show} | Character: {s.character}\n"
-            f"  Situation: {s.setup}\n"
-            f"  What they said: {s.response}\n"
+            f"  Archetype: {s.archetype.value} | Tension: {s.tension_type.value} | "
+            f"Distance: {s.violation_distance.value}\n"
             f"  Why it worked: {s.why_it_works}"
         )
     return "\n\n".join(blocks)
@@ -90,7 +158,6 @@ def generate_candidates(
     model,
     tokenizer,
     context_str: str,
-    used_archetypes: set,
     personas_to_run: List[str] = None,
 ) -> List[CandidateResponse]:
     """Generate persona candidates (1-3) with ClichePenalty applied.
@@ -99,7 +166,6 @@ def generate_candidates(
     None = all three (default). Engine gates this based on twist_potential.
     """
     scenes_block = _build_scenes_block(scenes)
-    used_str = ", ".join(a.value for a in used_archetypes) if used_archetypes else "none yet"
 
     active_personas = {
         name: instr for name, instr in PERSONA_INSTRUCTIONS.items()
@@ -108,25 +174,41 @@ def generate_candidates(
 
     cliche_processor = ClichePenaltyProcessor(metadata.obvious_response, tokenizer)
     processors = LogitsProcessorList([cliche_processor])
+    if config.ENABLE_BAD_WORD_GUARD:
+        bad_words_ids = [
+            tokenizer.encode(phrase, add_special_tokens=False)
+            for phrase in config.BAD_WORD_PHRASES
+        ]
+        processors.append(NoBadWordsLogitsProcessor(bad_words_ids, eos_token_id=tokenizer.eos_token_id))
 
     candidates = []
+    ngram_n = config.OVERLAP_NGRAM_SIZE
+    last_raw = None
+    last_persona = None
     for persona_name, persona_instruction in active_personas.items():
         prompt = GENERATION_PROMPT.format(
             user_input=user_input,
             surface=metadata.surface,
             subtext=metadata.subtext,
             archetype=metadata.archetype.value,
+            archetype_confidence=metadata.archetype_confidence,
             tension_type=metadata.tension_type.value,
             power_dynamic=metadata.power_dynamic,
+            speaker_strategy=metadata.speaker_strategy or "none",
+            connector=metadata.connector or "none",
             scenes_block=scenes_block,
             persona_name=persona_name.upper(),
             persona_instruction=persona_instruction,
-            used_archetypes_str=used_str,
             obvious_response=metadata.obvious_response,
             context_str=context_str or "(no prior context)",
         )
 
         raw = generate_text(prompt, model, tokenizer, config_type="generate", logits_processors=processors)
+        last_raw, last_persona = raw, persona_name
+        if config.ENABLE_OVERLAP_GUARD and _copies_retrieved_phrase(raw, scenes, ngram_n):
+            logger.warning(f"[{persona_name}] copies precedent phrasing — dropping candidate")
+            continue
+
         logger.info(f"[{persona_name}] → {raw[:80]}...")
 
         candidates.append(CandidateResponse(
@@ -153,22 +235,38 @@ def generate_candidates(
                 extra_ids = tokenizer.encode(word, add_special_tokens=False)
                 cliche_processor.penalty_ids.update(extra_ids[:2])  # first 2 tokens of each word
 
+    if not candidates and last_raw:
+        logger.warning(f"All candidates dropped — keeping last line from {last_persona}")
+        candidates.append(CandidateResponse(
+            persona=last_persona,
+            text=last_raw.strip(),
+            violation_type=f"{metadata.archetype.value} via {last_persona} lens",
+        ))
+
     return candidates
 
 
 RANK_PROMPT = """\
 You are judging {n} comedy responses to: "{user_input}"
+Connector word (has two simultaneous readings in the input): "{connector}"
+Extracted context:
+- subtext: "{subtext}"
+- archetype: "{archetype}"
+- tension_type: "{tension_type}"
+- speaker_strategy: "{speaker_strategy}"
 
 {candidates_block}
 
 Pick the funniest one using this exact priority order:
 
-1. CONCRETE IMAGE — The best response ends with a single specific, unexpected image or action that makes the human truth visible. Responses using only abstract concepts or bureaucratic jargon with no specific image always lose. "The drill is merely a gentle hug" beats "the patient submitted a fitness statement."
-2. SHARPNESS — Does the punchline land on first read without unpacking?
-3. TRUTH — Does it name something recognizable that nobody said out loud?
-4. BREVITY — If sharpness and image quality are equal, pick the shorter one.
+1. DOMAIN ANCHOR — Prefer lines that stay inside the user's world from surface/subtext. Down-rank candidates that import HR, legal process, supply chain, new departments, or institutional bureaucracy unless the user input already implies them. A grounded but slightly softer punchline beats a sharper line that drifts into an unsupported external domain.
+2. FINAL CLAUSE — The punchline is always the last clause. Judge the quality of the ENDING, not the setup. A sharp ending on a flat setup beats a flat ending on a sharp setup. The best ending is a single specific, unexpected image or action that makes the human truth visible. Responses where the final clause is abstract, bureaucratic, or jargon always lose.
+3. CONNECTOR — If a response lands on the second meaning of the connector word in its punchline, this is a strong quality signal. It means the wit is structurally grounded in the input, not floating free.
+4. SHARPNESS — Does the punchline land on first read without unpacking?
+5. TRUTH — Does it name something recognizable that nobody said out loud?
+6. BREVITY — If sharpness and image quality are equal, pick the shorter one.
 
-A sharp 20-word line with a specific concrete image beats a flat 10-word line of jargon.
+A sharp 20-word line with a specific concrete ending beats a flat 10-word line of jargon.
 Responses that are purely bureaucratic or purely abstract always lose, regardless of length.
 
 Reply ONLY with a single digit ({valid_digits}). Nothing else."""
@@ -176,6 +274,7 @@ Reply ONLY with a single digit ({valid_digits}). Nothing else."""
 
 def rank_candidates(
     user_input: str,
+    metadata: ComedyMetadata,
     candidates: List[CandidateResponse],
     model,
     tokenizer,
@@ -228,6 +327,11 @@ def rank_candidates(
     prompt = RANK_PROMPT.format(
         n=len(shuffled),
         user_input=user_input,
+        connector=(metadata.connector or "none"),
+        subtext=metadata.subtext,
+        archetype=metadata.archetype.value,
+        tension_type=metadata.tension_type.value,
+        speaker_strategy=metadata.speaker_strategy or "none",
         candidates_block="\n".join(lines),
         valid_digits="/".join(digits),
     )

witgym/model.py

@@ -1,6 +1,7 @@
 """Model loading and ClichePenaltyProcessor."""
 import re
 import torch
+import gc
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -153,6 +154,7 @@ def generate_text(
     # Free KV cache + intermediate buffers from unified memory after each call
     if config.DEVICE == "mps":
         torch.mps.empty_cache()
+        gc.collect()
 
     new_tokens = output_ids[0][inputs["input_ids"].shape[-1]:]
     raw = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

witgym/retriever.py

@@ -4,13 +4,14 @@ Implements Principle 2: query is the abstract comedy metadata, not raw text.
 Returns analogous situations (same violation type), not similar words.
 """
 import json
-from typing import List, Set
+from typing import List
 from loguru import logger
 import numpy as np
 from witgym import config
-from witgym.schemas import ComedyArchetype, TensionType, ViolationDistance, TranscriptScene, ComedyMetadata
+from witgym.schemas import TranscriptScene, ComedyMetadata
 
 _index_cache = None  # Loaded once, reused across calls
+_reranker_cache = None
 
 
 def load_index(index_path: str = config.INDEX_PATH) -> dict:
@@ -31,18 +32,32 @@ def load_index(index_path: str = config.INDEX_PATH) -> dict:
     return _index_cache
 
 
+def _load_reranker():
+    global _reranker_cache
+    if not config.ENABLE_CROSS_ENCODER_RERANK:
+        return None
+    if _reranker_cache is not None:
+        return _reranker_cache
+
+    try:
+        from sentence_transformers import CrossEncoder
+
+        logger.info(f"Loading reranker: {config.RERANK_MODEL_ID} on {config.RERANK_DEVICE}")
+        _reranker_cache = CrossEncoder(config.RERANK_MODEL_ID, device=config.RERANK_DEVICE)
+        return _reranker_cache
+    except Exception as e:
+        logger.warning(f"Reranker unavailable; falling back to cosine only: {e}")
+        _reranker_cache = False
+        return None
+
+
 def retrieve_scenes(
     index: dict,
     metadata: ComedyMetadata,
-    used_archetypes: Set[ComedyArchetype],
     embed_model,
     top_k: int = config.TOP_K_SCENES,
 ) -> List[TranscriptScene]:
-    """Principle 2: query on abstract comedy fields, return analogous situations.
-
-    Over-fetches (top_k * 3) then filters out already-used archetypes to
-    prevent repetition across the conversation.
-    """
+    """Principle 2: query on abstract comedy fields, return analogous situations."""
     # Mirror the indexed representation: keep the structural labels, then add
     # the extracted subtext so the query has semantic detail comparable to setup.
     query = (
@@ -66,36 +81,60 @@ def retrieve_scenes(
     # Rank all scenes by similarity
     ranked_indices = np.argsort(scores)[::-1]
 
-    # Filter: exclude scenes whose archetype is already used in this session,
-    # AND enforce intra-call diversity (no two retrieved scenes share the same archetype)
-    selected = []
-    selected_archetypes: Set[ComedyArchetype] = set()
-    for idx in ranked_indices:
-        scene = scenes[idx]
-        if scene.archetype in used_archetypes:
+    # Stage 1: pool by cosine (wide) → Stage 2: rerank (narrow)
+    pool_size = max(config.RETRIEVE_POOL_SIZE, top_k * 3)
+    pool_indices = ranked_indices[:pool_size]
+    pool_scenes = [scenes[i] for i in pool_indices]
+
+    reranker = _load_reranker()
+    if reranker is not None:
+        query_text = (
+            f"{metadata.subtext}\n"
+            f"archetype={metadata.archetype.value} "
+            f"tension={metadata.tension_type.value} "
+            f"distance={metadata.violation_distance.value}"
+        )
+        doc_texts = [
+            (
+                f"setup={s.setup}\n"
+                f"why_it_works={s.why_it_works}\n"
+                f"archetype={s.archetype.value} "
+                f"tension={s.tension_type.value} "
+                f"distance={s.violation_distance.value}"
+            )
+            for s in pool_scenes
+        ]
+        pairs = [[query_text, d] for d in doc_texts]
+        rerank_scores = np.asarray(reranker.predict(pairs), dtype=np.float32)
+        reranked_order = np.argsort(rerank_scores)[::-1]
+        ranked_scenes = [pool_scenes[i] for i in reranked_order]
+        logger.debug(
+            "Rerank top5 scores: "
+            + ", ".join(f"{rerank_scores[i]:.3f}" for i in reranked_order[:5])
+        )
+    else:
+        ranked_scenes = pool_scenes
+
+    # Same extracted archetype first; backfill from reranked pool only if thin
+    target = metadata.archetype
+    selected: List[TranscriptScene] = []
+    for scene in ranked_scenes:
+        if scene.archetype != target:
             continue
-        if scene.archetype in selected_archetypes:
-            continue  # Already have a scene of this archetype in this call
         selected.append(scene)
-        selected_archetypes.add(scene.archetype)
         if len(selected) >= top_k:
             break
 
-    # Fallback 1: relax intra-call diversity if not enough distinct archetypes found
     if len(selected) < top_k:
-        for idx in ranked_indices:
-            scene = scenes[idx]
-            if scene.archetype in used_archetypes:
+        for scene in ranked_scenes:
+            if scene in selected:
                 continue
-            if scene not in selected:
-                selected.append(scene)
+            selected.append(scene)
             if len(selected) >= top_k:
                 break
 
-    # Fallback 2: all session archetypes exhausted — return top matches without filter
     if not selected:
-        logger.warning("All archetypes used — returning top scenes without archetype filter")
-        selected = [scenes[i] for i in ranked_indices[:top_k]]
+        selected = ranked_scenes[:top_k]
 
     logger.info(f"Retrieved {len(selected)} scenes: {[s.archetype.value for s in selected]}")
     return selected

witgym/schemas.py

@@ -31,11 +31,14 @@ class ComedyMetadata(BaseModel):
     surface: str
     subtext: str
     archetype: ComedyArchetype
+    archetype_confidence: int = 7  # 1-10: confidence in archetype selection
     tension_type: TensionType
     power_dynamic: str
     obvious_response: str
     violation_distance: ViolationDistance
     twist_potential: int = 5  # 1-10: comedy richness of the input; drives pipeline gating
+    connector: Optional[str] = None  # Word/phrase with two simultaneous readings; null if absent
+    speaker_strategy: Optional[str] = None  # Brief guidance for how the speaker is trying to be perceived
 
 
 class TranscriptScene(BaseModel):