158 lines
6.7 KiB
Python
158 lines
6.7 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Менеджер индексации документов.
|
||
Отвечает за сжатие больших документов и вызов KBService для индексации.
|
||
"""
|
||
|
||
import logging
|
||
from typing import Optional, Dict, Any, Tuple
|
||
|
||
from core.services.kb_service import KBService
|
||
from core.services.giga_client import GigaClient
|
||
from core.functions.hierarchical_summarize import hierarchical_summarize
|
||
from core.utils.text_utils import count_tokens
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class IndexingManager:
|
||
"""
|
||
Управляет индексацией документов в базу знаний.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
kb: KBService,
|
||
giga: GigaClient,
|
||
config,
|
||
default_prompts: Optional[Dict[str, str]] = None
|
||
):
|
||
"""
|
||
Инициализация менеджера индексации.
|
||
|
||
Аргументы:
|
||
kb: сервис базы знаний
|
||
giga: клиент GigaChat (для сжатия документов)
|
||
config: объект конфигурации
|
||
default_prompts: словарь промптов по умолчанию
|
||
"""
|
||
self.kb = kb
|
||
self.giga = giga
|
||
self.config = config
|
||
self.default_prompts = default_prompts or {}
|
||
|
||
async def compress_document_if_needed(
|
||
self,
|
||
text: str,
|
||
file_name: str
|
||
) -> str:
|
||
"""
|
||
Сжимает документ, если он превышает порог токенов.
|
||
Использует иерархическое резюмирование.
|
||
|
||
Аргументы:
|
||
text: исходный текст документа
|
||
file_name: имя файла (для логирования)
|
||
|
||
Возвращает:
|
||
str: сжатый текст или исходный, если сжатие не требуется
|
||
"""
|
||
summarization_config = getattr(self.config, 'summarization', {})
|
||
enabled = summarization_config.get('enable_hierarchical_summarization', True)
|
||
if not enabled:
|
||
return text
|
||
|
||
max_tokens = summarization_config.get('max_tokens_for_document', 8000)
|
||
target_tokens = summarization_config.get('target_tokens_after_summary', 3000)
|
||
chunk_size = summarization_config.get('chunk_size_tokens', 500)
|
||
max_depth = summarization_config.get('max_depth', 2)
|
||
|
||
text_tokens = count_tokens(text)
|
||
if text_tokens <= max_tokens:
|
||
logger.debug(f"Документ '{file_name}' ({text_tokens} токенов) не требует сжатия")
|
||
return text
|
||
|
||
logger.info(
|
||
f"Документ '{file_name}' ({text_tokens} токенов) превышает лимит {max_tokens}, "
|
||
"применяем иерархическое резюмирование"
|
||
)
|
||
|
||
# Загружаем промпт для резюмирования
|
||
prompt_template = self.default_prompts.get('hierarchical_summary', '')
|
||
if not prompt_template:
|
||
try:
|
||
prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt'
|
||
with open(prompt_path, 'r', encoding='utf-8') as f:
|
||
prompt_template = f.read()
|
||
except Exception as e:
|
||
logger.warning(f"Не удалось загрузить промпт hierarchical_summary: {e}, используем стандартный")
|
||
prompt_template = (
|
||
"Ты — профессиональный реферант. Кратко изложи суть текста, сохранив ключевые факты, цифры и логические связи.\n"
|
||
"Объём: не более 30% от исходного, но не менее 2 предложений.\n"
|
||
"Стиль: деловой, нейтральный, без оценок.\n\n"
|
||
"ТЕКСТ:\n{text}\n\nРЕЗЮМЕ:"
|
||
)
|
||
|
||
try:
|
||
compressed = await hierarchical_summarize(
|
||
text=text,
|
||
giga=self.giga,
|
||
prompt_template=prompt_template,
|
||
target_tokens=target_tokens,
|
||
chunk_size_tokens=chunk_size,
|
||
max_depth=max_depth,
|
||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||
)
|
||
logger.info(
|
||
f"Документ '{file_name}' сжат с {text_tokens} до {count_tokens(compressed)} токенов"
|
||
)
|
||
return compressed
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при сжатии документа '{file_name}': {e}, используем исходный текст")
|
||
return text
|
||
|
||
async def index_document(
|
||
self,
|
||
file_name: str,
|
||
file_text: str,
|
||
user_jid: str,
|
||
room_jid: Optional[str],
|
||
is_global: bool = False,
|
||
title: Optional[str] = None,
|
||
metadata: Optional[Dict] = None,
|
||
file_hash: Optional[str] = None,
|
||
update_if_exists: bool = True
|
||
) -> Tuple[int, int]:
|
||
"""
|
||
Индексирует документ в базу знаний, предварительно сжимая его при необходимости.
|
||
|
||
Аргументы:
|
||
file_name: имя файла
|
||
file_text: текст документа
|
||
user_jid: JID владельца
|
||
room_jid: JID комнаты (None для личного)
|
||
is_global: глобальный ли документ
|
||
title: отображаемое название
|
||
metadata: дополнительные метаданные
|
||
file_hash: SHA-256 хеш содержимого
|
||
update_if_exists: заменять ли существующий документ в комнате
|
||
|
||
Возвращает:
|
||
(doc_id, chunk_count): идентификатор документа и количество чанков
|
||
"""
|
||
# Сжимаем документ, если он слишком большой
|
||
compressed_text = await self.compress_document_if_needed(file_text, file_name)
|
||
|
||
# Передаём сжатый текст в сервис базы знаний
|
||
doc_id, chunk_count = await self.kb.add_document(
|
||
file_name=file_name,
|
||
file_text=compressed_text,
|
||
user_jid=user_jid,
|
||
is_global=is_global,
|
||
title=title,
|
||
metadata=metadata,
|
||
room_jid=room_jid,
|
||
file_hash=file_hash,
|
||
update_if_exists=update_if_exists
|
||
)
|
||
return doc_id, chunk_count |