witgym/engine.py

"""WitGym engine — pure function core. No I/O, no UI.

This is the single file that both CLI (main.py) and Gradio (app.py) call.
"""
import re
import time
from typing import Iterator, Optional
from loguru import logger
from sentence_transformers import SentenceTransformer
from witgym import config
from witgym.model import load_model, generate_text
from witgym.schemas import WitGymResponse, PipelineEvent, fallback_metadata
from witgym.extractor import extract_comedy_metadata
from witgym.retriever import load_index, retrieve_scenes
from witgym.generator import (
    generate_candidates_stream,
    rank_candidates,
    compress_winner_stream,
)
from witgym.conversation import ConversationManager
from witgym.router import classify_intent
from witgym.prompts import BANTER_PROMPT, COACHING_ASK_PROMPT, EXPLAIN_PROMPT, SHARPEN_PROMPT, DRILL_KEYS

_shared_resources = None

STREAM_MIN_INTERVAL = 0.08
STREAM_MIN_CHARS = 3


def _cap_two_sentences(text: str) -> str:
    """Hard cap at 2 sentences — safety net if model ignores prompt constraint."""
    sentences = re.split(r'(?<=[.!?])\s+', text.strip())
    if len(sentences) <= 2:
        return text.strip()
    return " ".join(sentences[:2]).strip()


def _maybe_emit_token(
    phase: str,
    persona: Optional[str],
    acc: str,
    last_emit: float,
    metadata,
    scenes,
    candidates,
    **extra,
) -> tuple[Optional[PipelineEvent], float, str]:
    now = time.monotonic()
    if len(acc) >= STREAM_MIN_CHARS or (now - last_emit) >= STREAM_MIN_INTERVAL:
        event = PipelineEvent(
            phase=phase,
            persona=persona,
            partial_text=acc,
            metadata=metadata,
            scenes=scenes,
            candidates=candidates,
            **extra,
        )
        return event, now, ""
    return None, last_emit, acc


class SharedResources:
    """Heavy resources loaded once per process (model/tokenizer, embedder, index)."""

    def __init__(self, index_path: str = config.INDEX_PATH):
        logger.info("Loading shared WitGym resources...")
        self.model, self.tokenizer = load_model()
        logger.info(f"Loading embedding model: {config.EMBED_MODEL_ID}")
        self.embed_model = SentenceTransformer(config.EMBED_MODEL_ID, device=config.DEVICE)
        self.index = load_index(index_path)
        logger.success("Shared resources ready.")


def get_shared_resources(index_path: str = config.INDEX_PATH) -> SharedResources:
    """Process-wide singleton for Spaces / Gradio multi-session use."""
    global _shared_resources
    if _shared_resources is None:
        _shared_resources = SharedResources(index_path=index_path)
    return _shared_resources


class WitGymEngine:
    """Stateful engine. Load once, call respond() in a loop."""

    def __init__(
        self,
        index_path: str = config.INDEX_PATH,
        resources: Optional[SharedResources] = None,
        conversation: Optional[ConversationManager] = None,
        last_wit_response: Optional[WitGymResponse] = None,
        character: str = "AI",
    ):
        if resources is None:
            logger.info("Initialising WitGymEngine...")
            self._resources = SharedResources(index_path=index_path)
        else:
            self._resources = resources
        self.conversation = conversation or ConversationManager()
        self._last_wit_response: Optional[WitGymResponse] = last_wit_response
        self.character = character
        logger.success("WitGymEngine ready.")

    @property
    def model(self):
        return self._resources.model

    @property
    def tokenizer(self):
        return self._resources.tokenizer

    @property
    def embed_model(self):
        return self._resources.embed_model

    @property
    def index(self):
        return self._resources.index

    def respond(self, user_input: str) -> WitGymResponse:
        """Full two-pass pipeline. Returns WitGymResponse with all intermediate data."""
        result = None
        for event in self.respond_stream(user_input):
            if event.phase == "done" and event.response is not None:
                result = event.response
        if result is None:
            raise RuntimeError("respond_stream did not emit a done event")
        return result

    def respond_stream(self, user_input: str) -> Iterator[PipelineEvent]:
        """Incremental pipeline for Gradio streaming UI."""
        drill_type = DRILL_KEYS.get(user_input.strip())
        if drill_type and self._last_wit_response is not None:
            yield from self._handle_drill(drill_type, self._last_wit_response)
            return

        if self.conversation.coaching_state.get("awaiting"):
            from witgym.router import _heuristic_route
            if _heuristic_route(user_input) == "banter" or len(user_input.strip()) < 5:
                self.conversation.coaching_state = {}
                yield from self._handle_banter(user_input)
                return
            original = self.conversation.coaching_state.get("original_input", user_input)
            yield from self._run_wit_pipeline(
                user_input,
                route="coaching",
                display_input=user_input,
                coaching_context=(original, user_input),
                include_explanation=True,
            )
            return

        route = classify_intent(user_input, self.model, self.tokenizer)

        if route == "banter":
            yield from self._handle_banter(user_input)
            return

        if route == "coaching":
            yield from self._handle_coaching_ask(user_input)
            return

        self.conversation.set_mode("quick_wit")
        yield from self._run_wit_pipeline(user_input, route="quick_wit")

    def _handle_banter(self, user_input: str) -> Iterator[PipelineEvent]:
        logger.info("Banter route — dynamic reply")
        self.conversation.set_mode("banter")
        self._last_wit_response = None
        reply = generate_text(
            BANTER_PROMPT.format(user_input=user_input),
            self.model,
            self.tokenizer,
            config_type="generate",
        )
        reply = _cap_two_sentences(reply.strip())
        if not reply or reply.lstrip().startswith("1.") or "analyze user" in reply.lower():
            reply = "I'm here for awkward moments and sharp comebacks — toss me a situation."
        metadata = fallback_metadata(user_input)
        self.conversation.add_turn(user_input, reply, metadata, mode="banter")
        response = WitGymResponse(
            metadata=metadata,
            retrieved_scenes=[],
            candidates=[],
            selected=reply,
            route="banter",
        )
        yield PipelineEvent(phase="banter", response=response)
        yield PipelineEvent(phase="done", response=response)

    def _handle_coaching_ask(self, user_input: str) -> Iterator[PipelineEvent]:
        logger.info("Coaching route — Turn 1 clarifying question")
        self.conversation.set_mode("coaching")
        question = generate_text(
            COACHING_ASK_PROMPT.format(user_input=user_input),
            self.model,
            self.tokenizer,
            config_type="generate",
        )
        question = question.strip().split("\n")[0].strip() or "What did you actually say back?"
        self.conversation.coaching_state = {
            "awaiting": True,
            "original_input": user_input,
        }
        metadata = fallback_metadata(user_input)
        response = WitGymResponse(
            metadata=metadata,
            retrieved_scenes=[],
            candidates=[],
            selected=question,
            route="coaching",
            coaching_question=question,
        )
        yield PipelineEvent(phase="coaching_ask", partial_text=question, response=response)
        yield PipelineEvent(phase="done", response=response)

    def _handle_drill(self, drill_type: str, last: WitGymResponse) -> Iterator[PipelineEvent]:
        """Refine the last wit response. Does NOT call add_turn() — no memory pollution."""
        joke = last.selected
        metadata = last.metadata

        if drill_type == "angle":
            yield from self._run_wit_pipeline(
                metadata.surface,
                route=last.route,
                exclude_joke=joke,
            )
            return

        if drill_type == "sharpen":
            situation = last.retrieved_scenes[0].setup if last.retrieved_scenes else metadata.surface
            reply = generate_text(
                SHARPEN_PROMPT.format(situation=situation, joke=joke, subtext=metadata.subtext),
                self.model, self.tokenizer, config_type="generate",
            )
            reply = _cap_two_sentences(reply.strip())
            self._last_wit_response = WitGymResponse(
                metadata=metadata,
                retrieved_scenes=last.retrieved_scenes,
                candidates=last.candidates,
                selected=reply,
                route=last.route,
                winning_persona=last.winning_persona,
            )
        else:  # explain
            reply = generate_text(
                EXPLAIN_PROMPT.format(
                    joke=joke,
                    archetype=metadata.archetype.value,
                    subtext=metadata.subtext,
                ),
                self.model, self.tokenizer, config_type="generate",
            )
            reply = _cap_two_sentences(reply.strip())

        response = WitGymResponse(
            metadata=metadata,
            retrieved_scenes=last.retrieved_scenes,
            candidates=last.candidates,
            selected=reply,
            route="drill",
        )
        yield PipelineEvent(phase="banter", response=response)
        yield PipelineEvent(phase="done", response=response)

    def _run_wit_pipeline(
        self,
        user_input: str,
        route: str = "quick_wit",
        *,
        display_input: Optional[str] = None,
        coaching_context: Optional[tuple[str, str]] = None,
        include_explanation: bool = False,
        exclude_joke: Optional[str] = None,
    ) -> Iterator[PipelineEvent]:
        """CBR-RAG pipeline shared by quick_wit and coaching Turn 2."""
        turn_user_input = display_input or user_input

        if self.conversation.needs_compression(self.tokenizer):
            self.conversation.compress(self.model, self.tokenizer)

        metadata = extract_comedy_metadata(
            user_input,
            self.model,
            self.tokenizer,
            coaching_context=coaching_context,
        )
        yield PipelineEvent(phase="metadata", metadata=metadata)

        if metadata.twist_potential < 4:
            logger.info(f"twist_potential={metadata.twist_potential} < 4 — returning straight reply")
            selected = "Yeah, that tracks."
            self.conversation.add_turn(turn_user_input, selected, metadata, mode=route)
            response = WitGymResponse(
                metadata=metadata,
                retrieved_scenes=[],
                candidates=[],
                selected=selected,
                route=route,
            )
            yield PipelineEvent(
                phase="ranked",
                metadata=metadata,
                scenes=[],
                candidates=[],
                selected=selected,
            )
            yield PipelineEvent(phase="done", response=response)
            return

        scenes = retrieve_scenes(self.index, metadata, self.embed_model)
        yield PipelineEvent(phase="scenes", metadata=metadata, scenes=scenes)

        context_str = self.conversation.get_context_string()
        if exclude_joke:
            context_str = f"[Previous line — do NOT repeat]: {exclude_joke}\n" + context_str
        personas_to_run = ["cynic", "conviction", "absurdist"]
        if metadata.twist_potential <= 6:
            personas_to_run = ["cynic", "absurdist"]
            logger.info(f"twist_potential={metadata.twist_potential} ≤ 6 — cynic + absurdist")
        elif metadata.twist_potential > 8:
            personas_to_run = ["cynic", "absurdist", "bisociate"]
            logger.info(f"twist_potential={metadata.twist_potential} > 8 — bisociate replaces conviction")

        candidates = []
        candidate_acc = ""
        candidate_persona: Optional[str] = None
        last_emit = 0.0
        pending_chars = ""

        for item in generate_candidates_stream(
            user_input, metadata, scenes,
            self.model, self.tokenizer,
            context_str,
            personas_to_run=personas_to_run,
            character=self.character,
        ):
            event_type = item[0]
            payload = item[1] if len(item) == 2 else item[1:]
            if event_type == "candidate_start":
                candidate_acc = ""
                candidate_persona = payload
                yield PipelineEvent(
                    phase="candidate_start",
                    persona=payload,
                    metadata=metadata,
                    scenes=scenes,
                    candidates=list(candidates),
                )
            elif event_type == "candidate_token":
                persona_name, token_piece = payload
                candidate_acc += token_piece
                pending_chars += token_piece
                evt, last_emit, pending_chars = _maybe_emit_token(
                    "candidate_token",
                    persona_name,
                    pending_chars,
                    last_emit,
                    metadata,
                    scenes,
                    list(candidates),
                )
                if evt:
                    evt.partial_text = candidate_acc
                    yield evt
            elif event_type == "candidate_done":
                if pending_chars and candidate_persona:
                    yield PipelineEvent(
                        phase="candidate_token",
                        persona=candidate_persona,
                        partial_text=candidate_acc,
                        metadata=metadata,
                        scenes=scenes,
                        candidates=list(candidates),
                    )
                    pending_chars = ""
                if payload is not None:
                    candidates.append(payload)
                yield PipelineEvent(
                    phase="candidate_done",
                    persona=getattr(payload, "persona", None),
                    metadata=metadata,
                    scenes=scenes,
                    candidates=list(candidates),
                )
            elif event_type == "candidates_complete":
                candidates = payload

        selected = rank_candidates(user_input, metadata, candidates, self.model, self.tokenizer)
        winning_persona = next((c.persona for c in candidates if c.text == selected), None)
        yield PipelineEvent(
            phase="ranked",
            metadata=metadata,
            scenes=scenes,
            candidates=candidates,
            selected=selected,
            winning_persona=winning_persona,
        )

        yield PipelineEvent(phase="final_start", partial_text=selected, winning_persona=winning_persona)
        compressed = selected
        final_acc = ""
        pending_chars = ""
        last_emit = 0.0
        for event_type, payload in compress_winner_stream(selected, self.model, self.tokenizer):
            if event_type == "skip":
                compressed = payload
                yield PipelineEvent(phase="final_token", partial_text=compressed, winning_persona=winning_persona)
            elif event_type == "start":
                final_acc = selected
            elif event_type == "token":
                final_acc += payload
                pending_chars += payload
                evt, last_emit, pending_chars = _maybe_emit_token(
                    "final_token",
                    winning_persona,
                    pending_chars,
                    last_emit,
                    metadata,
                    scenes,
                    candidates,
                    winning_persona=winning_persona,
                )
                if evt:
                    evt.partial_text = final_acc
                    yield evt
            elif event_type == "done":
                if pending_chars:
                    yield PipelineEvent(
                        phase="final_token",
                        partial_text=final_acc,
                        winning_persona=winning_persona,
                    )
                compressed = payload

        selected = _cap_two_sentences(compressed)
        explanation = None
        if include_explanation:
            explanation = generate_text(
                EXPLAIN_PROMPT.format(
                    joke=selected,
                    archetype=metadata.archetype.value,
                    subtext=metadata.subtext,
                ),
                self.model,
                self.tokenizer,
                config_type="generate",
            )
            explanation = _cap_two_sentences(explanation.strip())

        self.conversation.add_turn(turn_user_input, selected, metadata, mode=route)

        response = WitGymResponse(
            metadata=metadata,
            retrieved_scenes=scenes,
            candidates=candidates,
            selected=selected,
            route=route,
            winning_persona=winning_persona,
            explanation=explanation,
        )
        self._last_wit_response = response
        yield PipelineEvent(phase="done", response=response)