diff --git a/rag/indexing_manager.py b/rag/indexing_manager.py new file mode 100644 index 0000000..e2a209f --- /dev/null +++ b/rag/indexing_manager.py @@ -0,0 +1,158 @@ +# -*- 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 \ No newline at end of file