Add new file

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