src/omnisub/backends/base.py

"""Interface backend LLM + tiện ích parse JSON bền bỉ từ output model.

Mọi backend (transformers Qwen3-Omni, llama.cpp fallback) phải hiện thực
`LLMBackend`. Pipeline chỉ phụ thuộc vào interface này, không vào model cụ thể.
"""

from __future__ import annotations

import abc
import json
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence


@dataclass
class ChatMessage:
    """Một message đa phương thức gửi tới model.

    - `text`: nội dung văn bản.
    - `images`: đường dẫn ảnh/frame (đưa TRƯỚC text khi build prompt).
    - `audio`: đường dẫn audio (đưa SAU text — theo best practice Qwen multimodal).
    """

    role: str = "user"
    text: str = ""
    images: List[str] = field(default_factory=list)
    audio: List[str] = field(default_factory=list)


class LLMBackend(abc.ABC):
    """Giao diện tối thiểu cho mọi backend."""

    @abc.abstractmethod
    def chat(
        self,
        messages: Sequence[ChatMessage],
        *,
        system: Optional[str] = None,
        max_new_tokens: int = 1024,
        temperature: float = 0.7,
        top_p: float = 0.95,
        top_k: int = 64,
        **kwargs: Any,
    ) -> str:
        """Sinh văn bản từ chuỗi message. Trả về raw text."""
        raise NotImplementedError

    def chat_json(
        self,
        messages: Sequence[ChatMessage],
        *,
        system: Optional[str] = None,
        max_new_tokens: int = 1024,
        **kwargs: Any,
    ) -> Any:
        """Gọi `chat` rồi cố trích JSON từ output. Trả về dict/list đã parse."""
        raw = self.chat(
            messages, system=system, max_new_tokens=max_new_tokens, **kwargs
        )
        return extract_json(raw)

    # Tiện ích cho subclass: gom text từ message (cho backend text-only).
    @staticmethod
    def _join_text(messages: Sequence[ChatMessage]) -> str:
        return "\n".join(m.text for m in messages if m.text)


_FENCE_RE = re.compile(r"```(?:json)?\s*(.*?)```", re.DOTALL | re.IGNORECASE)


def extract_json(raw: str) -> Any:
    """Trích object/array JSON đầu tiên từ output model.

    Thử theo thứ tự: parse trực tiếp → bóc khối ```json``` → quét cặp ngoặc
    cân bằng đầu tiên. Ném `ValueError` nếu không tìm được JSON hợp lệ.
    """
    if raw is None:
        raise ValueError("Output rỗng, không có JSON.")
    text = raw.strip()

    # 1) parse thẳng
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # 2) bóc khối ```json ... ```
    m = _FENCE_RE.search(text)
    if m:
        candidate = m.group(1).strip()
        try:
            return json.loads(candidate)
        except json.JSONDecodeError:
            text = candidate  # tiếp tục quét trong nội dung khối

    # 3) quét cặp ngoặc cân bằng đầu tiên ({} hoặc [])
    for opener, closer in (("{", "}"), ("[", "]")):
        snippet = _extract_balanced(text, opener, closer)
        if snippet:
            try:
                return json.loads(snippet)
            except json.JSONDecodeError:
                continue

    raise ValueError(f"Không trích được JSON từ output:\n{raw[:500]}")


def _extract_balanced(text: str, opener: str, closer: str) -> Optional[str]:
    start = text.find(opener)
    if start == -1:
        return None
    depth = 0
    in_str = False
    escape = False
    for i in range(start, len(text)):
        ch = text[i]
        if in_str:
            if escape:
                escape = False
            elif ch == "\\":
                escape = True
            elif ch == '"':
                in_str = False
            continue
        if ch == '"':
            in_str = True
        elif ch == opener:
            depth += 1
        elif ch == closer:
            depth -= 1
            if depth == 0:
                return text[start : i + 1]
    return None