display.py

"""Helpers that turn raw model output into Markdown for the chatbot widget."""

from __future__ import annotations

import json
import re
from typing import Any

from i18n import t

# KaTeX renders negation macros (\neq, \notin, …) via \not + SVG slash
# overlays.  The chat bubble's inflated line-height stretches those SVGs
# into solid-black rectangles.  Replacing with Unicode avoids the SVG
# path entirely — KaTeX renders the single glyph from its math fonts
# instead.
#
# The dict covers both single-macro forms (\neq) and \not compound forms
# (\not=, \not\equiv).  Sorted longest-first when building the regex so
# longer patterns match before shorter prefixes.
_NEGATION_UNICODE: dict[str, str] = {
    # \not + symbol compounds
    r"\not\subseteq": "⊈",
    r"\not\supseteq": "⊉",
    r"\not\parallel": "∦",
    r"\not\approx":   "≉",
    r"\not\subset":   "⊄",
    r"\not\supset":   "⊅",
    r"\not\simeq":    "≄",
    r"\not\exists":   "∄",
    r"\not\equiv":    "≢",
    r"\not\cong":     "≇",
    r"\not\leq":      "≰",
    r"\not\geq":      "≱",
    r"\not\sim":      "≁",
    r"\not\mid":      "∤",
    r"\not\le":       "≰",
    r"\not\ge":       "≱",
    r"\not\in":       "∉",
    r"\not\ni":       "∌",
    r"\not=":         "≠",
    r"\not<":         "≮",
    r"\not>":         "≯",
    # Standard single-macro negations
    r"\nsubseteq":    "⊈",
    r"\nsupseteq":    "⊉",
    r"\nparallel":    "∦",
    r"\nexists":      "∄",
    r"\nequiv":       "≢",
    r"\notin":        "∉",
    r"\nless":        "≮",
    r"\nleq":         "≰",
    r"\ngeq":         "≱",
    r"\ngtr":         "≯",
    r"\nmid":         "∤",
    r"\ncong":        "≇",
    r"\nsim":         "≁",
    r"\neq":          "≠",
    r"\ne":           "≠",
}
_NEGATION_RE = re.compile(
    "("
    + "|".join(
        re.escape(m)
        for m in sorted(_NEGATION_UNICODE, key=len, reverse=True)
    )
    + r")(?![a-zA-Z])",
)


_NOT_SPACE_RE = re.compile(r"\\not\s+(?=[\[=<>]|\\[a-zA-Z])")


def _replace_negation_macros(text: str) -> str:
    """Swap negation macros for Unicode glyphs to bypass SVG overlays."""
    # Collapse optional whitespace after \not so \not \equiv → \not\equiv
    text = _NOT_SPACE_RE.sub(r"\\not", text)
    return _NEGATION_RE.sub(lambda m: _NEGATION_UNICODE[m.group(1)], text)


# LaTeX environments that should be treated as display math.
_MATH_ENVS = (
    "equation", "align", "aligned", "gather", "gathered",
    "multline", "cases", "array", "split",
    "matrix", "bmatrix", "pmatrix", "vmatrix", "Vmatrix",
    "eqnarray",
)
_ENV_NAMES = "|".join(_MATH_ENVS)

# Matches \begin{env}...\end{env} blocks NOT already inside $$ delimiters.
# Uses a backreference (\2) to ensure begin/end names match.
_BARE_ENV_RE = re.compile(
    r"(?<!\$)"
    r"(\\begin\{(" + _ENV_NAMES + r")\*?\}"
    r"[\s\S]*?"
    r"\\end\{\2\*?\})"
    r"(?!\$)",
)


def _wrap_bare_latex_envs(text: str) -> str:
    r"""Wrap bare \begin{env}...\end{env} blocks in $$ for KaTeX.

    Models sometimes output LaTeX environments without wrapping them in
    display-math delimiters.  KaTeX only renders them when they appear
    inside ``$$...$$`` or ``\[...\]``.
    """
    return _BARE_ENV_RE.sub(r"\n\n$$\1$$\n\n", text)


def _escape_latex_delimiters(text: str) -> str:
    r"""Replace LaTeX delimiters with HTML entities so a downstream
    Markdown / KaTeX pass does not match them across the <details>
    boundary.

    Covers ``\(...\)``, ``\[...\]``, ``$...$``, and ``$$...$$``.
    """
    return (
        text
        .replace("\\(", "&#92;(")
        .replace("\\)", "&#92;)")
        .replace("\\[", "&#92;[")
        .replace("\\]", "&#92;]")
        .replace("$", "&#36;")
    )


def preprocess_latex(text: str) -> str:
    """Normalize model output so KaTeX renders math correctly.

    Applied to the assistant *content* (not thinking) before display.
    """
    text = _replace_negation_macros(text)
    text = _wrap_bare_latex_envs(text)
    return text


def build_display_content(reasoning: str, content: str) -> str:
    """Wrap reasoning + content into a collapsible Markdown block."""
    parts: list[str] = []
    if reasoning:
        is_complete = bool(content)
        summary = t("display.thinking_done") if is_complete else t("display.thinking")
        open_attr = "" if is_complete else " open"
        safe_reasoning = _escape_latex_delimiters(reasoning.strip())
        parts.append(
            f'<details class="thinking-block"{open_attr}>'
            f'<summary>{summary}</summary>\n\n{safe_reasoning}\n\n</details>\n\n'
        )
    if content:
        parts.append(preprocess_latex(content))
    return "".join(parts)


def format_tool_calls_for_display(tool_calls: list[Any]) -> str:
    return "\n\n".join(_format_single_tool_call(tc) for tc in tool_calls)


def format_tool_call_prompt(tc: Any, current_idx: int, total: int) -> str:
    header = f"**{t('tool.call_header', i=current_idx, n=total)}**\n\n"
    return header + _format_single_tool_call(tc)


def _format_single_tool_call(tc: Any) -> str:
    if isinstance(tc, dict):
        name = tc["function"]["name"]
        args_raw = tc["function"]["arguments"]
    else:
        name = tc.function.name
        args_raw = tc.function.arguments
    try:
        args_formatted = json.dumps(json.loads(args_raw), indent=2, ensure_ascii=False)
    except (json.JSONDecodeError, TypeError):
        args_formatted = args_raw
    return f"{t('tool.call_label')}: **{name}**\n```json\n{args_formatted}\n```"