Update file rag_orchestrator.py

rag/rag_orchestrator.py

@@ -10,6 +10,10 @@
 
 Улучшено управление токенами: теперь история и контекст обрезаются с учётом
 лимитов модели, резервирования для ответа и промптов.
+
+ДОБАВЛЕНО:
+- Иерархическое резюмирование истории диалога.
+- Иерархическое резюмирование больших документов при индексации.
 """
 
 import asyncio
@@ -89,9 +93,9 @@ class RAGOrchestrator:
         system_prompt: Optional[str],
         query: str,
         context: str,
-        max_total_tokens: int = 8192,  # для GigaChat-Max
+        max_total_tokens: int = 8192,
         reserved_for_answer: int = 1000,
-        reserved_for_overhead: int = 200  # буфер для форматирования
+        reserved_for_overhead: int = 200
     ) -> Dict[str, Any]:
         """
         Подготавливает части промта и обрезает историю/контекст по токенам.
@@ -109,7 +113,6 @@ class RAGOrchestrator:
         query_tokens = count_tokens(query)
 
         # 2. Резервирование
-        # Вычитаем системный промпт, шаблон, запрос, буфер и место для ответа
         prompt_tokens = system_tokens + synthesis_tokens + query_tokens
         available_for_history_and_context = (
             max_total_tokens
@@ -123,12 +126,8 @@ class RAGOrchestrator:
                 f"Недостаточно токенов для истории и контекста: {available_for_history_and_context}. "
                 f"Увеличьте max_total_tokens или уменьшите размер промптов."
             )
-            # Если отрицательное – устанавливаем минимум 100 токенов для истории и контекста
             available_for_history_and_context = max(available_for_history_and_context, 100)
 
-        # 3. Обрезка истории (уже есть в process_query, но мы её переместим сюда)
-        # Для простоты пока оставим обрезку в process_query, но используем этот метод для расчёта лимита
-
         return {
             "available_for_history_and_context": available_for_history_and_context,
             "prompt_tokens": prompt_tokens,
@@ -137,6 +136,67 @@ class RAGOrchestrator:
             "query_tokens": query_tokens,
         }
 
+    # ------------------------------------------------------------------
+    # Метод для сжатия больших документов при индексации
+    # ------------------------------------------------------------------
+
+    async def _compress_document_if_needed(self, text: str, file_name: str) -> str:
+        """
+        Если текст превышает порог (max_tokens_for_document), сжимает его
+        с помощью иерархического резюмирования.
+        Возвращает либо исходный текст, либо сжатый.
+        """
+        # Читаем настройки из конфига
+        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}, применяем иерархическое резюмирование")
+
+        # Загружаем промпт для резюмирования
+        summary_prompt = self.default_prompts.get('hierarchical_summary', '')
+        if not summary_prompt:
+            try:
+                prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt'
+                with open(prompt_path, 'r', encoding='utf-8') as f:
+                    summary_prompt = f.read()
+            except Exception as e:
+                logger.warning(f"Не удалось загрузить промпт hierarchical_summary: {e}, используем стандартный")
+                summary_prompt = (
+                    "Ты — профессиональный реферант. Кратко изложи суть текста, сохранив ключевые факты, цифры и логические связи.\n"
+                    "Объём: не более 30% от исходного, но не менее 2 предложений.\n"
+                    "Стиль: деловой, нейтральный, без оценок.\n\n"
+                    "ТЕКСТ:\n{text}\n\nРЕЗЮМЕ:"
+                )
+
+        try:
+            compressed = await hierarchical_summarize(
+                text=text,
+                giga=self.giga,
+                prompt_template=summary_prompt,
+                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
+
     # ------------------------------------------------------------------
     # Основной метод обработки запроса
     # ------------------------------------------------------------------
@@ -186,18 +246,16 @@ class RAGOrchestrator:
         # ----- 2. Получение истории диалога из БД (необрезанной) -----
         history = await self.db.get_history(user_jid, room_jid, limit=100)
 
-        # ----- 3. Расчёт лимитов токенов (НОВОЕ) -----
+        # ----- 3. Расчёт лимитов токенов -----
-        # Параметры из конфига
+        max_model_tokens = getattr(self.config, 'max_model_tokens', 8192)
-        max_model_tokens = getattr(self.config, 'max_model_tokens', 8192)  # GigaChat-Max default
         reserved_for_answer = getattr(self.config, 'reserved_for_answer_tokens', 1000)
         reserved_for_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200)
 
-        # Используем вспомогательный метод для расчёта
         token_info = self._prepare_prompt_parts(
             synthesis_template=synthesis_template,
             system_prompt=system_prompt,
             query=query,
-            context="",  # пока пусто, позже добавим
+            context="",
             max_total_tokens=max_model_tokens,
             reserved_for_answer=reserved_for_answer,
             reserved_for_overhead=reserved_for_overhead
@@ -205,25 +263,30 @@ class RAGOrchestrator:
         available_for_history_and_context = token_info["available_for_history_and_context"]
 
         logger.debug(
-            f"Доступно для истории и контекста: {available_for_history_and_context} токенов, "
+            f"Доступно для истории и контекста: {available_for_history_and_context} токенов"
-            f"промпт: {token_info['prompt_tokens']} токенов"
         )
 
-        # ----- 4. Обрезка истории (с использованием иерархического резюмирования) -----
+        # ----- 4. Сжатие истории (иерархическое резюмирование) -----
         max_history_tokens = min(available_for_history_and_context // 2, 2000)
 
         # Если история слишком длинная, сжимаем её иерархически
         history_text = "\n".join([f"{rec['role']}: {rec['content']}" for rec in history])
         if count_tokens(history_text) > max_history_tokens:
-            logger.info(f"История слишком длинная ({count_tokens(history_text)} токенов), применяем иерархическое резюмирование")
+            logger.info(
+                f"История слишком длинная ({count_tokens(history_text)} токенов), "
+                f"применяем иерархическое резюмирование"
+            )
             summary_prompt = prompts.get('hierarchical_summary', '')
             if not summary_prompt:
-                # Загружаем из файла, если не передано
                 try:
-                    with open(self.config.prompts_dir / 'hierarchical_summary.txt', 'r', encoding='utf-8') as f:
+                    prompt_path = self.config.prompts_dir / 'hierarchical_summary.txt'
+                    with open(prompt_path, 'r', encoding='utf-8') as f:
                         summary_prompt = f.read()
                 except Exception:
-                    summary_prompt = "Кратко изложи суть диалога, сохранив ключевые факты и вопросы.\n\nТЕКСТ:\n{text}"
+                    summary_prompt = (
+                        "Кратко изложи суть диалога, сохранив ключевые факты и вопросы.\n\n"
+                        "ТЕКСТ:\n{text}"
+                    )
             try:
                 compressed_history = await hierarchical_summarize(
                     text=history_text,
@@ -234,9 +297,12 @@ class RAGOrchestrator:
                     max_depth=2,
                     temperature=getattr(self.config, 'ai_temperature', 0.1)
                 )
-                # Заменяем историю одним сжатым сообщением
+                formatted_history = [
-                formatted_history = [{"role": "system", "content": "Сжатая история диалога:\n" + compressed_history}]
+                    {"role": "system", "content": "Сжатая история диалога:\n" + compressed_history}
-                logger.info(f"История сжата с {count_tokens(history_text)} до {count_tokens(compressed_history)} токенов")
+                ]
+                logger.info(
+                    f"История сжата с {count_tokens(history_text)} до {count_tokens(compressed_history)} токенов"
+                )
             except Exception as e:
                 logger.error(f"Ошибка при иерархическом резюмировании истории: {e}")
                 # Fallback: простая обрезка
@@ -255,14 +321,15 @@ class RAGOrchestrator:
                     for rec in truncated_history
                 ]
         else:
-            # Если история укладывается, просто форматируем
+            # Если история укладывается, форматируем последние сообщения
-            truncated_history = history[-10:]  # ограничим последними 10 сообщениями для простоты, но можно взять больше
+            # Берём максимум 20 сообщений, чтобы не перегружать
+            truncated_history = history[-20:]
             formatted_history = [
                 {"role": rec['role'], "content": rec['content']}
                 for rec in truncated_history
             ]
 
-        # ----- 5. Классификация намерений (если не переопределена) -----
+        # ----- 5. Классификация намерений -----
         intent = intent_override
         if intent is None and getattr(self.config, 'enable_intent_classification', True):
             intent_prompt = prompts.get('intent', '')
@@ -395,7 +462,6 @@ class RAGOrchestrator:
                 elif not last_file_path.lower().endswith('.docx'):
                     answer = "❌ Проверка орфографии поддерживается только для файлов .docx."
                 else:
-                    # TODO: адаптировать check_spelling для работы через HTTP
                     answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
 
         # --- 7.6. GREETING ---
@@ -479,30 +545,24 @@ class RAGOrchestrator:
             )
             logger.info(f"Найден контекст длиной {len(context)} символов (room={room_jid})")
 
-            # 8.3. Обрезка контекста по токенам (НОВОЕ)
+            # 8.3. Обрезка контекста по токенам
             # Оставшиеся токены после истории и промптов – используем для контекста
-            max_context_tokens = available_for_history_and_context - total_history_tokens
+            max_context_tokens = available_for_history_and_context - sum(count_tokens(msg['content']) for msg in formatted_history)
-            # Защита от отрицательных значений
             max_context_tokens = max(max_context_tokens, 0)
 
             if max_context_tokens > 0 and context:
-                # Если контекст слишком длинный, обрезаем по токенам
-                # Используем функцию count_tokens для оценки
                 context_tokens = count_tokens(context)
                 if context_tokens > max_context_tokens:
-                    # Простая обрезка: берём первые max_context_tokens токенов
-                    # В будущем можно использовать иерархическое резюмирование
                     logger.warning(
                         f"Контекст слишком длинный ({context_tokens} токенов), обрезаем до {max_context_tokens}"
                     )
-                    # Обрезаем по символам (грубая оценка: 1 токен ≈ 3 символа)
+                    max_context_chars = int(max_context_tokens * 3.5)
-                    max_context_chars = int(max_context_tokens * 3)
                     if max_context_chars > 0:
                         context = context[:max_context_chars]
                     else:
                         context = ""
 
-            # 8.4. Переранжирование контекста (если включено)
+            # 8.4. Переранжирование контекста
             rerank_min_length = getattr(self.config, 'rerank_min_length', 5000)
             if intent != "FACT" and len(context) > rerank_min_length:
                 context = await rerank_context(
@@ -597,6 +657,10 @@ class RAGOrchestrator:
             "confidence": None
         }
 
+    # ------------------------------------------------------------------
+    # Индексация документа (с возможным сжатием)
+    # ------------------------------------------------------------------
+
     async def index_document(
         self,
         file_name: str,
@@ -611,11 +675,16 @@ class RAGOrchestrator:
     ) -> Dict[str, Any]:
         """
         Индексация документа в базу знаний.
-        Вызывается из HTTP-эндпоинта /rag/index.
+        Перед индексацией применяет иерархическое резюмирование,
+        если документ слишком большой.
         """
+        # Сжимаем документ, если он слишком большой
+        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=file_text,
+            file_text=compressed_text,
             user_jid=user_jid,
             is_global=is_global,
             title=title,