Update file rag_orchestrator.py

rag/rag_orchestrator.py

@@ -8,8 +8,8 @@
 Этот модуль не зависит от XMPP и может использоваться как в ботах,
 так и в отдельном RAG-ядре (API-сервисе).
 
-Добавлено управление токенами: резервирование места для ответа и промтов,
+Улучшено управление токенами: теперь история и контекст обрезаются с учётом
-обрезка истории и контекста с учётом доступного пространства.
+лимитов модели, резервирования для ответа и промптов.
 """
 
 import asyncio
@@ -78,6 +78,68 @@ class RAGOrchestrator:
         self.default_prompts = default_prompts or {}
         logger.info("RAGOrchestrator инициализирован")
 
+    # ------------------------------------------------------------------
+    # Вспомогательный метод для расчёта токенов с резервированием
+    # ------------------------------------------------------------------
+
+    def _prepare_prompt_parts(
+        self,
+        synthesis_template: str,
+        system_prompt: Optional[str],
+        query: str,
+        context: str,
+        max_total_tokens: int = 8192,  # для GigaChat-Max
+        reserved_for_answer: int = 1000,
+        reserved_for_overhead: int = 200  # буфер для форматирования
+    ) -> Dict[str, Any]:
+        """
+        Подготавливает части промта и обрезает историю/контекст по токенам.
+
+        Возвращает словарь с ключами:
+            - history: отфильтрованная история (список сообщений)
+            - context: отфильтрованный контекст (строка)
+            - prompt_tokens: количество токенов в промте (без истории и контекста)
+            - total_used: общее использованное количество токенов
+            - available_for_history_and_context: сколько токенов выделено для истории и контекста
+        """
+        # 1. Подсчёт токенов в статичных частях промта
+        system_tokens = count_tokens(system_prompt) if system_prompt else 0
+        synthesis_tokens = count_tokens(synthesis_template)
+        query_tokens = count_tokens(query)
+
+        # 2. Резервирование
+        # Вычитаем системный промпт, шаблон, запрос, буфер и место для ответа
+        prompt_tokens = system_tokens + synthesis_tokens + query_tokens
+        available_for_history_and_context = (
+            max_total_tokens
+            - prompt_tokens
+            - reserved_for_answer
+            - reserved_for_overhead
+        )
+
+        if available_for_history_and_context <= 0:
+            logger.warning(
+                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,
+            "system_tokens": system_tokens,
+            "synthesis_tokens": synthesis_tokens,
+            "query_tokens": query_tokens,
+        }
+
+    # ------------------------------------------------------------------
+    # Основной метод обработки запроса
+    # ------------------------------------------------------------------
+
     async def process_query(
         self,
         query: str,
@@ -115,46 +177,64 @@ class RAGOrchestrator:
         # ----- 1. Подготовка промптов -----
         if prompts is None:
             prompts = self.default_prompts.copy()
-        synthesis_template = prompts.get('synthesis', '{context}\n\n{query}\n\nОтвет:')
+        synthesis_template = prompts.get('synthesis', '')
+        if not synthesis_template:
+            synthesis_template = "{context}\n\n{query}\n\nОтвет:"
         system_prompt = prompts.get('system', None)
 
-        # ----- 2. Получение истории диалога из БД -----
+        # ----- 2. Получение истории диалога из БД (необрезанной) -----
         history = await self.db.get_history(user_jid, room_jid, limit=100)
 
-        # ----- 3. Управление токенами -----
+        # ----- 3. Расчёт лимитов токенов (НОВОЕ) -----
-        # 3.1. Резервирование места для ответа и промтов
+        # Параметры из конфига
-        max_context_tokens = getattr(self.config, 'max_context_tokens', 3000)
+        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_prompt = 0
+        reserved_for_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200)
-        if system_prompt:
-            reserved_for_prompt += count_tokens(system_prompt)
-        reserved_for_prompt += count_tokens(synthesis_template)
-        # Дополнительный резерв на служебные токены
-        overhead = 100
-        available_for_history_and_context = max_context_tokens - reserved_for_answer - reserved_for_prompt - overhead
 
-        if available_for_history_and_context < 100:
+        # Используем вспомогательный метод для расчёта
-            logger.warning(f"Доступное пространство для истории и контекста слишком мало: {available_for_history_and_context} токенов. Увеличьте max_context_tokens или уменьшите резервы.")
+        token_info = self._prepare_prompt_parts(
-            available_for_history_and_context = max(100, available_for_history_and_context)
+            synthesis_template=synthesis_template,
+            system_prompt=system_prompt,
+            query=query,
+            context="",  # пока пусто, позже добавим
+            max_total_tokens=max_model_tokens,
+            reserved_for_answer=reserved_for_answer,
+            reserved_for_overhead=reserved_for_overhead
+        )
+        available_for_history_and_context = token_info["available_for_history_and_context"]
 
-        logger.debug(f"Управление токенами: max={max_context_tokens}, резерв ответа={reserved_for_answer}, промт={reserved_for_prompt}, доступно={available_for_history_and_context}")
+        logger.debug(
+            f"Доступно для истории и контекста: {available_for_history_and_context} токенов, "
+            f"промпт: {token_info['prompt_tokens']} токенов"
+        )
 
-        # 3.2. Обрезка истории по токенам (сохраняем последние сообщения)
+        # ----- 4. Обрезка истории по токенам (с учётом лимита) -----
-        total_tokens = 0
+        # Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние
         truncated_history = []
+        total_history_tokens = 0
+        history_limit = min(available_for_history_and_context // 2, 2000)  # выделяем половину для истории
+
         for record in reversed(history):
             tokens = count_tokens(record['content'])
-            if total_tokens + tokens <= available_for_history_and_context:
+            if total_history_tokens + tokens <= history_limit:
                 truncated_history.append(record)
-                total_tokens += tokens
+                total_history_tokens += tokens
             else:
                 break
-        truncated_history.reverse()
+        truncated_history.reverse()  # восстанавливаем хронологический порядок
-        # Сразу обрезаем историю для использования в дальнейшем
-        formatted_history = truncated_history
-        logger.debug(f"История обрезана до {len(truncated_history)} сообщений, {total_tokens} токенов")
 
-        # ----- 4. Классификация намерений -----
+        # Преобразуем в формат, ожидаемый GigaChat
+        formatted_history = [
+            {"role": rec['role'], "content": rec['content']}
+            for rec in truncated_history
+        ]
+
+        logger.debug(
+            f"История обрезана: {len(truncated_history)} сообщений, "
+            f"{total_history_tokens} токенов (лимит {history_limit})"
+        )
+
+        # ----- 5. Классификация намерений (если не переопределена) -----
         intent = intent_override
         if intent is None and getattr(self.config, 'enable_intent_classification', True):
             intent_prompt = prompts.get('intent', '')
@@ -170,19 +250,19 @@ class RAGOrchestrator:
         else:
             intent = intent or "GENERAL"
 
-        # Принудительная установка SURGICAL
+        # ----- 6. Принудительная установка SURGICAL по ключевым словам -----
         keywords = getattr(self.config, 'surgical_keywords', [])
         if any(kw in query.lower() for kw in keywords) and last_file_path:
             intent = "SURGICAL"
             logger.info(f"Принудительный Intent: SURGICAL (есть файл и ключевые слова)")
 
-        # ----- 5. Обработка специализированных намерений -----
+        # ----- 7. Обработка специализированных намерений -----
         answer = None
         context = None
         ctx_for_critique = None
         synthesis_template_for_critique = None
 
-        # --- 5.1. METRICS ---
+        # --- 7.1. METRICS ---
         if intent == "METRICS":
             context = await self.kb.find_relevant_info(
                 query, user_jid, room_jid,
@@ -204,7 +284,7 @@ class RAGOrchestrator:
                 else:
                     answer = "Не удалось извлечь метрики."
 
-        # --- 5.2. SUMMARY ---
+        # --- 7.2. SUMMARY ---
         elif intent == "SUMMARY":
             if not last_file_text:
                 answer = "Нет документа для суммаризации."
@@ -218,7 +298,7 @@ class RAGOrchestrator:
                     bot_config=self.config
                 )
 
-        # --- 5.3. CONTRADICTION ---
+        # --- 7.3. CONTRADICTION ---
         elif intent == "CONTRADICTION":
             context = await self.kb.find_relevant_info(
                 query, user_jid, room_jid,
@@ -244,7 +324,7 @@ class RAGOrchestrator:
                     else:
                         answer = "✅ Противоречий не обнаружено."
 
-        # --- 5.4. TEMPLATE_FILL ---
+        # --- 7.4. TEMPLATE_FILL ---
         elif intent == "TEMPLATE_FILL":
             template_text = last_file_text or ""
             if not template_text and last_file_path and os.path.exists(last_file_path):
@@ -277,7 +357,7 @@ class RAGOrchestrator:
                     temperature=getattr(self.config, 'ai_temperature', 0.1)
                 )
 
-        # --- 5.5. SPELLCHECK ---
+        # --- 7.5. SPELLCHECK ---
         elif intent == "SPELLCHECK":
             if room_jid is not None:
                 answer = "⚠️ Проверка орфографии доступна только в личном чате."
@@ -290,7 +370,7 @@ class RAGOrchestrator:
                     # TODO: адаптировать check_spelling для работы через HTTP
                     answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
 
-        # --- 5.6. GREETING ---
+        # --- 7.6. GREETING ---
         elif intent == "GREETING":
             answer = await self.giga.chat(
                 history=formatted_history,
@@ -300,7 +380,7 @@ class RAGOrchestrator:
                 temperature=getattr(self.config, 'ai_temperature', 0.1)
             )
 
-        # --- 5.7. SURGICAL ---
+        # --- 7.7. SURGICAL ---
         elif intent == "SURGICAL":
             if not last_file_path:
                 answer = "❌ Нет загруженного документа для замены."
@@ -353,9 +433,9 @@ class RAGOrchestrator:
                     except ImportError:
                         answer = "❌ Библиотека mawo-pymorphy3 не установлена. Установите её для морфологической замены."
 
-        # ----- 6. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
+        # ----- 8. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
         if answer is None:
-            # 6.1. Расширение запроса
+            # 8.1. Расширение запроса
             expanded = await expand_query(
                 giga=self.giga,
                 query=query,
@@ -364,14 +444,37 @@ class RAGOrchestrator:
             )
             search_query = expanded if expanded and expanded != query else query
 
-            # 6.2. Поиск релевантного контекста в базе знаний
+            # 8.2. Поиск релевантного контекста в базе знаний
             context = await self.kb.find_relevant_info(
                 search_query, user_jid, room_jid,
                 top_k=getattr(self.config, 'rag_default_top_k', 30)
             )
             logger.info(f"Найден контекст длиной {len(context)} символов (room={room_jid})")
 
-            # 6.3. Переранжирование контекста (если включено)
+            # 8.3. Обрезка контекста по токенам (НОВОЕ)
+            # Оставшиеся токены после истории и промптов – используем для контекста
+            max_context_tokens = available_for_history_and_context - total_history_tokens
+            # Защита от отрицательных значений
+            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)
+                    if max_context_chars > 0:
+                        context = context[:max_context_chars]
+                    else:
+                        context = ""
+
+            # 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(
@@ -382,53 +485,20 @@ class RAGOrchestrator:
                     bot_config=self.config
                 )
 
-            # 6.4. Управление токенами для контекста и ответа
+            # 8.5. Синтез ответа с добавлением цепочки рассуждений (CoT)
-            # Вычисляем, сколько токенов уже занято историей
+            synthesis_template = synthesis_template
-            used_tokens = total_tokens + count_tokens(synthesis_template) + count_tokens(system_prompt or '')
-            # Вычисляем, сколько токенов осталось для контекста
-            remaining_for_context = available_for_history_and_context - used_tokens
-            if remaining_for_context < 0:
-                # Если даже без контекста превышаем – обрезаем историю жёстко
-                logger.warning(f"Перерасход токенов: история уже {used_tokens} токенов, доступно {available_for_history_and_context}")
-                # Урезаем историю до минимального размера (например, 5 сообщений)
-                if len(truncated_history) > 5:
-                    truncated_history = truncated_history[-5:]
-                    remaining_for_context = available_for_history_and_context - count_tokens(''.join([h['content'] for h in truncated_history])) - count_tokens(synthesis_template) - count_tokens(system_prompt or '')
-                    if remaining_for_context < 0:
-                        remaining_for_context = 100  # минимум
-            # Обрезаем контекст до remaining_for_context токенов
-            if remaining_for_context > 0 and context:
-                context_tokens = count_tokens(context)
-                if context_tokens > remaining_for_context:
-                    # Обрезаем контекст по границам предложений (упрощённо – по символам)
-                    # В идеале нужно использовать иерархическое резюмирование, но пока оставим простую обрезку
-                    chars_per_token = 3  # приблизительно
-                    max_chars = int(remaining_for_context * chars_per_token)
-                    if len(context) > max_chars:
-                        context = context[:max_chars]
-                        # Обрезаем до последнего полного предложения
-                        last_period = context.rfind('.')
-                        if last_period > 0:
-                            context = context[:last_period+1]
-                    logger.warning(f"Контекст обрезан до {len(context)} символов (приблизительно {remaining_for_context} токенов)")
-            elif remaining_for_context <= 0:
-                context = ""
-                logger.warning("Нет места для контекста в промте, ответ будет без контекста")
-
-            # 6.5. Синтез ответа с добавлением цепочки рассуждений (CoT)
-            synthesis_template_local = synthesis_template
             if intent in ("CALCULATION", "PROCEDURE"):
                 cot_instruction = (
                     "\n\nПожалуйста, покажи пошаговое решение перед итоговым ответом. "
                     "Опиши каждый шаг вычислений или действий в логической последовательности. "
                     "После всех шагов дай итоговый ответ."
                 )
-                synthesis_template_local += cot_instruction
+                synthesis_template += cot_instruction
 
-            full_query = synthesis_template_local.format(context=context, query=query)
+            full_query = synthesis_template.format(context=context, query=query)
             logger.debug(f"Полный запрос к GigaChat (первые 500 символов): {full_query[:500]}")
 
-            # 6.6. Генерация ответа через GigaChat
+            # 8.6. Генерация ответа
             answer = await self.giga.chat(
                 history=formatted_history,
                 query=full_query,
@@ -437,11 +507,10 @@ class RAGOrchestrator:
                 temperature=getattr(self.config, 'ai_temperature', 0.1)
             )
 
-            # Сохраняем контекст для самокритики
             ctx_for_critique = context
-            synthesis_template_for_critique = synthesis_template_local
+            synthesis_template_for_critique = synthesis_template
 
-        # ----- 7. Самокритика (если включена) -----
+        # ----- 9. Самокритика (если включена) -----
         if (getattr(self.config, 'enable_self_critique', False) and
             intent not in ("METRICS", "SUMMARY", "CONTRADICTION") and
             ctx_for_critique is not None and
@@ -480,18 +549,18 @@ class RAGOrchestrator:
                     ):
                         answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
 
-        # ----- 8. Сохранение истории диалога в БД -----
+        # ----- 10. Сохранение истории диалога в БД -----
         await self.db.add_history(user_jid, "user", query, room_jid)
         if answer:
             await self.db.add_history(user_jid, "assistant", answer, room_jid)
 
-        # ----- 9. Извлечение источников из контекста -----
+        # ----- 11. Извлечение источников из контекста -----
         sources = []
         if context:
             for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
                 sources.append(match.group(1))
 
-        # ----- 10. Формирование результата -----
+        # ----- 12. Формирование результата -----
         return {
             "answer": answer,
             "intent": intent,