Add new file

2026-06-30 12:44:44 +00:00
parent 8a2d5ba2ff
commit 8804bf736d
1 changed files with 239 additions and 0 deletions
--- a/rag/functions/hierarchical_summarize.py
+++ b/rag/functions/hierarchical_summarize.py
@@ -0,0 +1,239 @@
+# -*- coding: utf-8 -*-
+"""
+Модуль иерархического резюмирования текста.
+Разбивает текст на блоки, резюмирует каждый блок, затем объединяет резюме
+и при необходимости повторяет процесс (рекурсивно), пока не достигнет целевой длины.
+
+Используется для сжатия длинных диалогов и документов перед отправкой в LLM.
+
+Алгоритм:
+1. Разбить текст на предложения (с помощью razdel или простого split).
+2. Группировать предложения в блоки по заданному количеству токенов (chunk_size).
+3. Для каждого блока сгенерировать краткое резюме (через GigaChat).
+4. Объединить резюме в один текст.
+5. Если общее количество токенов резюме превышает target_tokens, повторить рекурсивно.
+6. Вернуть итоговое резюме.
+"""
+
+import logging
+import asyncio
+from typing import List, Optional, Tuple
+
+from core.services.giga_client import GigaClient
+from core.utils.text_utils import count_tokens, split_into_chunks
+
+logger = logging.getLogger(__name__)
+
+
+async def hierarchical_summarize(
+    text: str,
+    giga: GigaClient,
+    prompt_template: str,
+    target_tokens: int = 2000,
+    chunk_size_tokens: int = 500,
+    max_depth: int = 3,
+    temperature: float = 0.1,
+    min_chunk_tokens: int = 50,
+) -> str:
+    """
+    Иерархическое резюмирование текста.
+
+    Аргументы:
+        text (str): исходный текст.
+        giga (GigaClient): клиент GigaChat для генерации резюме.
+        prompt_template (str): шаблон промпта для резюмирования блока.
+                               Должен содержать плейсхолдер {text}.
+        target_tokens (int): целевой размер итогового резюме в токенах.
+        chunk_size_tokens (int): размер блока для резюмирования в токенах.
+        max_depth (int): максимальная глубина рекурсии (предотвращает бесконечный цикл).
+        temperature (float): температура генерации.
+        min_chunk_tokens (int): минимальный размер блока в токенах (меньшие блоки не резюмируются).
+
+    Возвращает:
+        str: итоговое резюме.
+    """
+    if not text or not text.strip():
+        return ""
+
+    # Если текст уже помещается в целевой лимит, возвращаем как есть
+    if count_tokens(text) <= target_tokens:
+        logger.debug(f"Текст уже укладывается в {target_tokens} токенов, резюмирование не требуется")
+        return text
+
+    # Если глубина достигла максимума, просто обрезаем текст
+    if max_depth <= 0:
+        logger.warning(f"Достигнута максимальная глубина резюмирования, обрезаем текст до {target_tokens} токенов")
+        return _truncate_by_tokens(text, target_tokens)
+
+    # Разбиваем текст на блоки (используем существующую функцию split_into_chunks)
+    # Но нам нужно разбиение по предложениям, а не просто по токенам.
+    # Используем собственный метод, который сохраняет целостность предложений.
+    blocks = _split_into_sentence_blocks(text, chunk_size_tokens)
+
+    if len(blocks) == 1 and count_tokens(blocks[0]) > target_tokens:
+        # Если один большой блок, рекурсивно резюмируем его
+        logger.debug(f"Один блок размером {count_tokens(blocks[0])} токенов, рекурсивное резюмирование")
+        return await hierarchical_summarize(
+            text=blocks[0],
+            giga=giga,
+            prompt_template=prompt_template,
+            target_tokens=target_tokens,
+            chunk_size_tokens=chunk_size_tokens,
+            max_depth=max_depth - 1,
+            temperature=temperature,
+            min_chunk_tokens=min_chunk_tokens
+        )
+
+    # Резюмируем каждый блок (параллельно)
+    logger.info(f"Резюмирование {len(blocks)} блоков (глубина {max_depth})")
+    summaries = []
+    for idx, block in enumerate(blocks):
+        if count_tokens(block) <= min_chunk_tokens:
+            # Маленький блок – оставляем как есть
+            summaries.append(block)
+            logger.debug(f"Блок {idx+1} слишком мал ({count_tokens(block)} токенов), пропускаем резюмирование")
+        else:
+            prompt = prompt_template.format(text=block)
+            try:
+                summary = await giga.chat(
+                    history=[],
+                    query=prompt,
+                    system_prompt=None,
+                    file_id=None,
+                    temperature=temperature
+                )
+                summaries.append(summary.strip())
+                logger.debug(f"Блок {idx+1} зарезюмирован ({count_tokens(block)} -> {count_tokens(summary)} токенов)")
+                # Небольшая задержка между запросами, чтобы не перегружать API
+                await asyncio.sleep(0.2)
+            except Exception as e:
+                logger.error(f"Ошибка резюмирования блока {idx+1}: {e}")
+                # В случае ошибки используем исходный блок
+                summaries.append(block)
+
+    # Объединяем резюме
+    combined = "\n\n".join(summaries)
+
+    # Если объединённое резюме всё ещё слишком большое, рекурсивно повторяем
+    if count_tokens(combined) > target_tokens:
+        logger.debug(f"Объединённое резюме ({count_tokens(combined)} токенов) превышает {target_tokens}, рекурсивный вызов")
+        return await hierarchical_summarize(
+            text=combined,
+            giga=giga,
+            prompt_template=prompt_template,
+            target_tokens=target_tokens,
+            chunk_size_tokens=chunk_size_tokens,
+            max_depth=max_depth - 1,
+            temperature=temperature,
+            min_chunk_tokens=min_chunk_tokens
+        )
+
+    return combined
+
+
+def _split_into_sentence_blocks(text: str, chunk_size_tokens: int) -> List[str]:
+    """
+    Разбивает текст на блоки по предложениям, стараясь не превышать chunk_size_tokens.
+    Использует razdel, если доступен, иначе регулярное выражение.
+    """
+    # Попытка использовать razdel для русского языка
+    try:
+        from razdel import sentenize
+        sentences = [sent.text for sent in sentenize(text)]
+    except ImportError:
+        import re
+        # Простое разбиение по точкам, вопросительным и восклицательным знакам
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+
+    blocks = []
+    current_block = []
+    current_tokens = 0
+
+    for sent in sentences:
+        sent_tokens = count_tokens(sent)
+        # Если одно предложение уже превышает chunk_size, разбиваем его принудительно
+        if sent_tokens > chunk_size_tokens:
+            # Если есть текущий блок, сохраняем его
+            if current_block:
+                blocks.append(" ".join(current_block))
+                current_block = []
+                current_tokens = 0
+            # Разбиваем длинное предложение на части по chunk_size_tokens (грубо)
+            parts = _split_long_sentence(sent, chunk_size_tokens)
+            blocks.extend(parts)
+            continue
+
+        if current_tokens + sent_tokens <= chunk_size_tokens:
+            current_block.append(sent)
+            current_tokens += sent_tokens
+        else:
+            if current_block:
+                blocks.append(" ".join(current_block))
+            current_block = [sent]
+            current_tokens = sent_tokens
+
+    if current_block:
+        blocks.append(" ".join(current_block))
+
+    return blocks
+
+
+def _split_long_sentence(sentence: str, chunk_size_tokens: int) -> List[str]:
+    """
+    Разбивает длинное предложение на части по chunk_size_tokens, пытаясь сохранить слова целиком.
+    """
+    # Разбиваем по пробелам
+    words = sentence.split()
+    if not words:
+        return []
+
+    parts = []
+    current_part = []
+    current_tokens = 0
+
+    for word in words:
+        word_tokens = count_tokens(word)
+        if current_tokens + word_tokens <= chunk_size_tokens:
+            current_part.append(word)
+            current_tokens += word_tokens
+        else:
+            if current_part:
+                parts.append(" ".join(current_part))
+            current_part = [word]
+            current_tokens = word_tokens
+
+    if current_part:
+        parts.append(" ".join(current_part))
+
+    return parts
+
+
+def _truncate_by_tokens(text: str, target_tokens: int) -> str:
+    """
+    Грубо обрезает текст до target_tokens токенов (по символам).
+    Используется как fallback при достижении максимальной глубины.
+    """
+    if count_tokens(text) <= target_tokens:
+        return text
+
+    # Оценка: 1 токен ≈ 3 символа для русского текста
+    max_chars = int(target_tokens * 3.5)  # небольшой запас
+    # Обрезаем по границе предложения
+    try:
+        from razdel import sentenize
+        sentences = [sent.text for sent in sentenize(text)]
+    except ImportError:
+        import re
+        sentences = re.split(r'(?<=[.!?])\s+', text)
+
+    result = []
+    current_len = 0
+    for sent in sentences:
+        sent_len = len(sent)
+        if current_len + sent_len <= max_chars:
+            result.append(sent)
+            current_len += sent_len
+        else:
+            break
+
+    return " ".join(result) if result else text[:max_chars]