Tolan/backend/emotion_text_utils.py

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

import re
from typing import Optional, Tuple

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))
_EMOTION_WORD_HEAD = re.compile(rf"^\s*(?P<word>{_EMOTION_WORDS_ALT})", re.IGNORECASE)


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


def _split_plain_emotion_metadata(text: str) -> Tuple[Optional[str], Optional[str]]:
    """
    若开头为裸 emotion 元数据，返回 (emotion, 正文)；否则 (None, None)。
    元数据形态：单独一行 / 空格+大写句首 / 粘连大写句首（如 friendlyHi）。
    不将「happy to see you」等正常英文句首当作元数据。
    """
    m = _EMOTION_WORD_HEAD.match(text)
    if not m:
        return None, None
    word = m.group("word")
    if word.lower() not in VALID_EMOTIONS:
        return None, None
    rest = text[m.end() :]
    if not rest:
        return normalize_emotion(word), ""
    if rest[0] in "\r\n":
        return normalize_emotion(word), rest.lstrip()
    if rest[0] == " ":
        body = rest.lstrip()
        if body and body[0].isascii() and body[0].isupper():
            return normalize_emotion(word), body
        return None, None
    if rest[0].isascii() and rest[0].isupper():
        return normalize_emotion(word), rest
    return None, None


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))
    emotion, _ = _split_plain_emotion_metadata(text)
    return emotion


def strip_emotion_from_text(text: str) -> str:
    """移除情感标签/裸词，供字幕与 TTS 使用。"""
    cleaned = EMOTION_TAG_PATTERN.sub("", text)
    _, body = _split_plain_emotion_metadata(cleaned)
    if body is not None:
        cleaned = body
    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)