diff --git a/rag/functions/hierarchical_summarize.py b/rag/functions/hierarchical_summarize.py new file mode 100644 index 0000000..21070a1 --- /dev/null +++ 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] \ No newline at end of file