LLM-infa/api/lib/summarize.py

"""Summarization with two backends:

* `summarize_llm` — uses any external OpenAI-compatible API (configured via
  `LLMConfig`). Long transcripts are map-reduced: per-chunk summaries are
  concatenated and re-summarized.
* `summarize_extractive` — pure-Python TF-IDF-ish fallback (RU + EN aware).

`summarize(text, llm_cfg=...)` picks the LLM path when a config is supplied,
otherwise falls back to extractive. The pipeline calls this entry point.
"""
from __future__ import annotations

import logging
import re
from collections import Counter
from typing import Optional

from api.lib.llm import LLMConfig, LLMError, chat_complete


logger = logging.getLogger("summarize")


_RU_STOPWORDS = {
    "и", "в", "во", "не", "на", "я", "что", "тот", "быть", "с", "со", "а", "весь",
    "как", "это", "но", "он", "она", "оно", "мы", "вы", "они", "к", "у", "из", "за",
    "по", "до", "от", "для", "же", "или", "бы", "если", "так", "там", "тут", "когда",
    "где", "кто", "какой", "этот", "эта", "эти", "мой", "твой", "наш", "ваш",
    "свой", "его", "её", "их", "был", "была", "было", "были", "есть", "нет", "да",
    "ну", "вот", "ещё", "еще", "уже", "только", "очень", "может", "можно", "надо",
    "будет", "будут", "всё", "все", "этого", "этом", "этой", "тоже", "также", "чтобы",
    "кого", "чего", "чем", "тем", "том", "нем", "ней", "себя", "себе", "мне", "тебе",
    "нам", "вам", "нас", "вас", "про", "без", "над", "под", "через", "при", "об",
    "о", "обо", "ли", "вон", "сюда", "туда", "оттуда", "потом", "теперь", "тогда",
    "сейчас", "иногда", "всегда", "никогда", "просто", "именно", "ведь", "значит",
    "хотя", "что-то", "кто-то",
}

_EN_STOPWORDS = {
    "the", "a", "an", "and", "or", "but", "is", "are", "was", "were", "be", "been",
    "being", "have", "has", "had", "do", "does", "did", "will", "would", "should",
    "could", "may", "might", "must", "shall", "can", "i", "you", "he", "she", "it",
    "we", "they", "them", "my", "your", "his", "her", "its", "our", "their", "this",
    "that", "these", "those", "in", "on", "at", "by", "for", "with", "about", "as",
    "of", "to", "from", "into", "so", "not", "no", "yes", "if", "then", "than",
    "when", "where", "why", "how", "what", "which", "who", "whom", "whose", "me",
    "us", "him", "just", "only", "also", "too", "very", "there", "here", "out",
    "up", "down", "over", "under", "between", "through", "again", "more", "most",
    "some", "any", "all", "each", "every", "such", "while", "because", "until",
    "during", "above", "below", "now", "ever", "never",
}

_STOPWORDS = _RU_STOPWORDS | _EN_STOPWORDS

_WORD_RE = re.compile(r"[\w']+", re.UNICODE)
_SENT_SPLIT_RE = re.compile(r"(?<=[.!?…])\s+(?=[А-ЯA-Z\"«„])")


DEFAULT_SYSTEM_PROMPT = (
    "Ты ассистент, который делает краткую и точную выжимку текста. "
    "Сохраняй ключевые факты, имена, цифры, причинно-следственные связи. "
    "Если оригинал на русском — пиши по-русски, иначе — на языке оригинала. "
    "Не добавляй информацию, которой нет в исходном тексте."
)

DEFAULT_USER_TEMPLATE = (
    "Сделай связную краткую выжимку следующего текста. "
    "Выдели основные тезисы и логику изложения. Объём — 5–12 предложений.\n\n"
    "---\n{text}\n---"
)

DEFAULT_REDUCE_TEMPLATE = (
    "Ниже — последовательность кратких выжимок частей одного длинного текста. "
    "Объедини их в единую связную выжимку, сохранив главные факты и логику.\n\n"
    "---\n{text}\n---"
)


def _split_sentences(text: str) -> list[str]:
    text = re.sub(r"\s+", " ", text).strip()
    if not text:
        return []
    parts = _SENT_SPLIT_RE.split(text)
    if len(parts) <= 1:
        parts = re.split(r"(?<=[.!?…])\s+", text)
    return [p.strip() for p in parts if p.strip()]


def _chunk_by_chars(text: str, max_chars: int) -> list[str]:
    if len(text) <= max_chars:
        return [text]
    chunks: list[str] = []
    cur: list[str] = []
    cur_len = 0
    for s in _split_sentences(text):
        if cur_len + len(s) + 1 > max_chars and cur:
            chunks.append(" ".join(cur))
            cur, cur_len = [], 0
        cur.append(s)
        cur_len += len(s) + 1
    if cur:
        chunks.append(" ".join(cur))
    return chunks


def summarize_extractive(
    text: str,
    *,
    max_sentences: int = 15,
    min_sentences: int = 3,
    ratio: float = 0.15,
) -> str:
    """Frequency-based extractive summary. Picks top sentences by mean word weight."""
    text = (text or "").strip()
    if not text:
        return ""
    sentences = _split_sentences(text)
    if not sentences:
        return ""

    target = max(min_sentences, min(max_sentences, int(len(sentences) * ratio)))
    target = max(1, min(target, len(sentences)))
    if len(sentences) <= target:
        return " ".join(sentences)

    words = _WORD_RE.findall(text.lower())
    freq = Counter(w for w in words if w not in _STOPWORDS and len(w) > 2)
    if not freq:
        return " ".join(sentences[:target])

    max_f = max(freq.values())
    norm = {w: c / max_f for w, c in freq.items()}

    scored: list[tuple[int, float]] = []
    for i, s in enumerate(sentences):
        sw = _WORD_RE.findall(s.lower())
        if not sw:
            continue
        score = sum(norm.get(w, 0.0) for w in sw) / len(sw)
        scored.append((i, score))
    if not scored:
        return " ".join(sentences[:target])

    top = sorted(scored, key=lambda x: -x[1])[:target]
    top.sort(key=lambda x: x[0])
    return " ".join(sentences[i] for i, _ in top)


def summarize_llm(
    text: str,
    cfg: LLMConfig,
    *,
    system_prompt: Optional[str] = None,
    user_template: Optional[str] = None,
    reduce_template: Optional[str] = None,
    chunk_chars: int = 10000,
) -> str:
    """Summarize via an external OpenAI-compatible API. Map-reduce for long texts."""
    text = (text or "").strip()
    if not text:
        return ""

    sys_p = system_prompt or DEFAULT_SYSTEM_PROMPT
    user_t = user_template or DEFAULT_USER_TEMPLATE
    reduce_t = reduce_template or DEFAULT_REDUCE_TEMPLATE

    def _one(chunk: str, template: str) -> str:
        return chat_complete(
            cfg,
            [
                {"role": "system", "content": sys_p},
                {"role": "user", "content": template.format(text=chunk)},
            ],
        )

    chunks = _chunk_by_chars(text, chunk_chars)
    logger.info("LLM summarize: %d chunk(s), total chars=%d", len(chunks), len(text))

    if len(chunks) == 1:
        return _one(chunks[0], user_t)

    partials: list[str] = []
    for i, ch in enumerate(chunks, start=1):
        logger.info("LLM summarize chunk %d/%d", i, len(chunks))
        partials.append(_one(ch, user_t))

    combined = "\n\n".join(partials)
    if len(combined) > chunk_chars:
        # If even concatenated partials are too long, recurse: summarize partials in batches.
        return summarize_llm(
            combined,
            cfg,
            system_prompt=sys_p,
            user_template=reduce_t,
            reduce_template=reduce_t,
            chunk_chars=chunk_chars,
        )
    logger.info("LLM reduce step over %d partials", len(partials))
    return _one(combined, reduce_t)


def summarize(
    text: str,
    *,
    llm_cfg: Optional[LLMConfig] = None,
    system_prompt: Optional[str] = None,
    user_template: Optional[str] = None,
    max_sentences: int = 15,
) -> str:
    """Primary entry point. Uses LLM if configured, else extractive fallback."""
    if llm_cfg is not None:
        try:
            return summarize_llm(
                text,
                llm_cfg,
                system_prompt=system_prompt,
                user_template=user_template,
            )
        except LLMError as e:
            logger.warning("LLM summarization failed, falling back to extractive: %s", e)
    return summarize_extractive(text, max_sentences=max_sentences)