agents.py

"""The emotion agents and the orchestrator that decides who speaks.

Flow for each user message:
  1. An orchestrator picks the 2-4 emotions most relevant to the message.
  2. Each chosen emotion replies *in character*, in parallel.
  3. A gentle "reflection" reads the room and helps the user name what's
     really going on underneath.

Everything degrades gracefully: with no model token (or no network) the app falls
back to a lightweight keyword-based simulation so the UI is always usable.
"""

from __future__ import annotations

import html
import json
import re
from concurrent.futures import ThreadPoolExecutor

from emotions import EMOTIONS, EMOTION_ORDER

# Token budgets. These are generous on purpose: "reasoning"/"thinking" models
# (Nemotron, Qwen3, gemma-with-thinking, etc.) spend tokens thinking *before*
# they emit the visible answer, so a tight budget leaves content empty. The
# prompts still ask for 1-2 short sentences, so non-reasoning models stop early.
ORCH_MAX_TOKENS = 512
CHIME_MAX_TOKENS = 600
REFLECT_MAX_TOKENS = 700
TEMPERATURE = 0.7
TOP_P = 0.95
# How many of the most recent turns to feed back to the agents as context.
HISTORY_TURNS = 100


def _get(obj, key: str, default=None):
    if isinstance(obj, dict):
        return obj.get(key, default)
    return getattr(obj, key, default)


# Reasoning models wrap their thoughts in <think>...</think> tags inside the
# content (or put them in a separate field). Strip the thinking so we keep only
# the user-facing answer.
_THINK_BLOCK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)


def _strip_reasoning(text: str) -> str:
    """Remove a reasoning model's chain-of-thought, leaving the answer."""
    # Drop any complete <think>...</think> blocks.
    text = _THINK_BLOCK_RE.sub("", text)
    low = text.lower()
    if "</think>" in low:
        # Answer is whatever follows the final closing tag.
        text = text[low.rfind("</think>") + len("</think>"):]
    elif "<think>" in low:
        # Only an opening tag: the model ran out of budget mid-thought and never
        # reached an answer. Keep anything before it (usually nothing).
        text = text[: low.find("<think>")]
    return text.strip()


def _chat_text(
    client,
    messages: list[dict[str, str]],
    *,
    max_tokens: int,
) -> str:
    """Call the Hugging Face chat-completion client and return response text."""
    resp = client.chat_completion(
        messages,
        max_tokens=max_tokens,
        temperature=TEMPERATURE,
        top_p=TOP_P,
    )
    choices = _get(resp, "choices", [])
    if not choices:
        return ""

    choice = choices[0]
    message = _get(choice, "message")
    if message is not None:
        content = _get(message, "content", "")
        if content is None:
            content = ""
        if isinstance(content, list):
            content = "".join(str(_get(part, "text", part)) for part in content)
        return _strip_reasoning(str(content))

    return _strip_reasoning(str(_get(choice, "text", "")))


_TAG_RE = re.compile(r"<[^>]+>")


def _readable(turn: dict) -> str:
    """Turn one stored chat entry into a clean 'Speaker: text' line.

    User messages are already plain text. Assistant entries are the HTML
    bubbles produced in app.py, so we strip the markup and recover the
    emotion's name (the bubble's first inner line) as the speaker label.
    """
    role = turn.get("role")
    content = str(turn.get("content", ""))

    if role == "user":
        text = content.strip()
        return f"You: {text}" if text else ""

    # Assistant bubble: split on tags, then unescape, to get [label, body...].
    parts = _TAG_RE.sub("\n", content.replace("<br>", " ").replace("<br/>", " "))
    chunks = [html.unescape(p).strip() for p in parts.splitlines()]
    chunks = [c for c in chunks if c]
    if not chunks:
        return ""
    if len(chunks) == 1:
        return f"Emotions: {chunks[0]}"
    label = chunks[0]  # e.g. "😨 Fear" or "A gentle reflection"
    body = " ".join(chunks[1:])
    return f"{label}: {body}"


def _history_text(history: list[dict]) -> str:
    """Flatten recent conversation into a clean, readable transcript."""
    lines = [_readable(turn) for turn in history[-HISTORY_TURNS:]]
    return "\n".join(line for line in lines if line)


# --------------------------------------------------------------------------- #
# Orchestrator: who should chime in?
# --------------------------------------------------------------------------- #

def choose_emotions(message: str, history: list[dict], client) -> list[str]:
    if client is None:
        return _fallback_choose(message)

    roster = "\n".join(
        f"- {EMOTIONS[k].name}: {EMOTIONS[k].tagline}" for k in EMOTION_ORDER
    )
    prompt = (
        "You are the control console inside someone's mind, like in the movie "
        "Inside Out. Read the person's latest message and decide which emotions "
        "would naturally light up and want to speak.\n\n"
        f"The emotions available are:\n{roster}\n\n"
        f"Recent conversation:\n{_history_text(history)}\n\n"
        f"Latest message: \"{message}\"\n\n"
        "Pick the 2 to 4 emotions that fit best. Favor an interesting, honest "
        "mix rather than always the same crowd. Respond with ONLY a JSON array "
        "of lowercase emotion keys from this set: "
        f"{json.dumps(EMOTION_ORDER)}."
    )
    try:
        text = _chat_text(
            client,
            [{"role": "user", "content": prompt}],
            max_tokens=ORCH_MAX_TOKENS,
        )
        match = re.search(r"\[.*\]", text, re.DOTALL)
        keys = json.loads(match.group(0)) if match else []
        keys = [k for k in keys if k in EMOTIONS]
        if keys:
            return keys[:4]
    except Exception:
        pass
    return _fallback_choose(message)


def _fallback_choose(message: str) -> list[str]:
    """Keyword heuristic used when no model is available."""
    m = message.lower()
    buckets = {
        "joy": ["happy", "great", "excited", "love", "win", "good", "fun", "yay", "proud"],
        "sadness": ["sad", "lonely", "miss", "lost", "cry", "hurt", "down", "grief", "tired"],
        "anxiety": ["worried", "anxious", "nervous", "stress", "deadline", "what if", "scared of"],
        "fear": ["afraid", "fear", "danger", "risk", "unsafe", "panic"],
        "anger": ["angry", "mad", "unfair", "annoyed", "furious", "hate", "frustrat"],
        "envy": ["jealous", "envy", "wish i", "they have", "compare", "better than me"],
        "embarrassment": ["embarrass", "awkward", "cringe", "ashamed", "stupid", "regret"],
        "disgust": ["gross", "disgust", "toxic", "fake", "ew", "hate the"],
        "ennui": ["bored", "meh", "whatever", "pointless", "dull", "nothing matters"],
    }
    hits = [k for k in EMOTION_ORDER if any(w in m for w in buckets[k])]
    if not hits:
        hits = ["joy", "sadness", "anxiety"]
    return hits[:4]


# --------------------------------------------------------------------------- #
# Individual emotion replies
# --------------------------------------------------------------------------- #

def emotion_reply(key: str, message: str, history: list[dict], client) -> str:
    emo = EMOTIONS[key]
    if client is None:
        return _fallback_line(key, message)

    system = (
        f"{emo.persona}\n\n"
        "You are one of several emotions reacting to the same person inside "
        "their head. Speak in first person, directly to them, in your own "
        "distinct voice. Keep it to 1-2 short, natural sentences. Stay fully in "
        "character. Do not narrate or use stage directions. Be warm and human."
    )
    user = (
        f"Recent conversation:\n{_history_text(history)}\n\n"
        f"Their latest message: \"{message}\"\n\n"
        f"React as {emo.name}."
    )
    try:
        text = _chat_text(
            client,
            [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            max_tokens=CHIME_MAX_TOKENS,
        )
        return text or _fallback_line(key, message)
    except Exception:
        return _fallback_line(key, message)


_FALLBACK_LINES = {
    "joy": "Hey, there's something good we can hold onto here — let's find it together!",
    "sadness": "It's okay if this feels heavy. I'm right here with you while it does.",
    "fear": "Just so we're careful — what's the part of this that feels risky to you?",
    "anger": "Wait, is this actually fair to you? You're allowed to take up space.",
    "disgust": "Ugh, honestly? You deserve better than whatever this is.",
    "anxiety": "Okay okay — let's think ahead. What's the plan if things don't go smoothly?",
    "envy": "I notice a little wanting in there... what is it you really wish you had?",
    "embarrassment": "I keep wondering how this looks to other people... but maybe that's okay.",
    "ennui": "Eh. Does this even matter to you, or are we just going through the motions?",
}


def _fallback_line(key: str, message: str) -> str:
    return _FALLBACK_LINES.get(key, "I'm feeling something about this too.")


# --------------------------------------------------------------------------- #
# Gentle closing reflection
# --------------------------------------------------------------------------- #

def reflection(message: str, replies: list[tuple[str, str]], client) -> str:
    if client is None:
        names = ", ".join(EMOTIONS[k].name for k, _ in replies)
        return (
            f"It sounds like {names} all showed up for you here. "
            "Which of them feels the most true right now?"
        )
    spoken = "\n".join(f"{EMOTIONS[k].name}: {text}" for k, text in replies)
    prompt = (
        "You are a calm, caring inner guide helping someone understand their "
        "feelings, in the spirit of Inside Out. The person said:\n"
        f"\"{message}\"\n\n"
        "Their emotions responded:\n"
        f"{spoken}\n\n"
        "In 1-2 warm sentences, gently reflect back what might be going on "
        "underneath for them, and invite them to notice which feeling rings "
        "most true. Do not list the emotions mechanically. Be tender, not "
        "clinical. End with a soft, open question."
    )
    try:
        text = _chat_text(
            client,
            [{"role": "user", "content": prompt}],
            max_tokens=REFLECT_MAX_TOKENS,
        )
        if text:
            return text
    except Exception:
        pass

    names = ", ".join(EMOTIONS[k].name for k, _ in replies)
    return (
        f"It sounds like {names} all showed up for you here. "
        "Which of them feels the most true right now?"
    )


# --------------------------------------------------------------------------- #
# Top-level turn
# --------------------------------------------------------------------------- #

def run_turn(message: str, history: list[dict], client=None):
    """Run one full turn.

    Returns (replies, reflection_text) where replies is a list of
    (emotion_key, text) tuples in a natural speaking order.
    """
    keys = choose_emotions(message, history, client)

    # Respond in parallel for snappiness, then restore the chosen order.
    with ThreadPoolExecutor(max_workers=len(keys) or 1) as pool:
        texts = list(
            pool.map(lambda k: emotion_reply(k, message, history, client), keys)
        )
    replies = list(zip(keys, texts))
    reflect = reflection(message, replies, client)
    return replies, reflect