239 lines
11 KiB
Python
239 lines
11 KiB
Python
# -*- 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] |