backend/utils/response_cleaner.py

"""Strip model thinking tags from assistant responses."""

from __future__ import annotations

import re

_THINK_OPEN = "<" + "think>"
_THINK_CLOSE = "</" + "think>"
_REDACTED_OPEN = "<think>"
_REDACTED_CLOSE = "</think>"

_THINKING_PATTERNS = (
    re.compile(re.escape(_THINK_OPEN) + r".*?" + re.escape(_THINK_CLOSE), re.DOTALL | re.IGNORECASE),
    re.compile(
        re.escape(_REDACTED_OPEN) + r".*?" + re.escape(_REDACTED_CLOSE),
        re.DOTALL | re.IGNORECASE,
    ),
    re.compile(re.escape(_THINK_OPEN) + r"\s*" + re.escape(_THINK_CLOSE), re.IGNORECASE),
    re.compile(re.escape(_REDACTED_OPEN) + r"\s*" + re.escape(_REDACTED_CLOSE), re.IGNORECASE),
)
_SPECIAL_TOKENS = re.compile(r"<\|[^|>]*\|>")


def clean_model_response(text: str) -> str:
    if not text:
        return ""
    cleaned = text
    for pattern in _THINKING_PATTERNS:
        cleaned = pattern.sub("", cleaned)
    cleaned = _SPECIAL_TOKENS.sub("", cleaned)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned.strip()