Tolan/backend/emotion_text_utils.py

"""LLM 情感标签解析与展示文本清理（字幕 / TTS 共用）"""
from __future__ import annotations

import re
from typing import Optional

VALID_EMOTIONS = frozenset(
    {
        "neutral",
        "happy",
        "sad",
        "excited",
        "angry",
        "surprised",
        "thoughtful",
        "encouraging",
        "confused",
        "worried",
        "satisfied",
        "friendly",
        "sorry",
    }
)

EMOTION_TAG_PATTERN = re.compile(r"\[EMOTION:\s*(\w+)\]", re.IGNORECASE)
_EMOTION_WORDS_ALT = "|".join(re.escape(w) for w in sorted(VALID_EMOTIONS, key=len, reverse=True))
_LEADING_PLAIN_EMOTION = re.compile(
    rf"^\s*(?P<word>{_EMOTION_WORDS_ALT})(?:\s*[\r\n]+|\s+(?=\S)|\s*$)",
    re.IGNORECASE,
)


def normalize_emotion(raw: str) -> str:
    word = (raw or "").strip().lower()
    return word if word in VALID_EMOTIONS else "neutral"


def parse_emotion_from_text(text: str) -> Optional[str]:
    """从 LLM 输出中解析情感（支持 [EMOTION:xxx] 及首行裸词）。"""
    if not text:
        return None
    tag_match = EMOTION_TAG_PATTERN.search(text)
    if tag_match:
        return normalize_emotion(tag_match.group(1))
    plain_match = _LEADING_PLAIN_EMOTION.match(text)
    if plain_match:
        return normalize_emotion(plain_match.group("word"))
    return None


def strip_emotion_from_text(text: str) -> str:
    """移除情感标签/裸词，供字幕与 TTS 使用。"""
    cleaned = EMOTION_TAG_PATTERN.sub("", text)
    cleaned = _LEADING_PLAIN_EMOTION.sub("", cleaned, count=1)
    cleaned = cleaned.strip()
    if re.fullmatch(rf"(?:{_EMOTION_WORDS_ALT})", cleaned, re.IGNORECASE):
        return ""
    return cleaned


def clean_text_for_display(text: str) -> str:
    """字幕展示文本清理（emotion 已通过独立 WebSocket 消息下发）。"""
    return strip_emotion_from_text(text)