Редактировать rag_orchestrator.py

rag/rag_orchestrator.py

@@ -7,6 +7,9 @@
 
 Этот модуль не зависит от XMPP и может использоваться как в ботах,
 так и в отдельном RAG-ядре (API-сервисе).
+
+Добавлено управление токенами: резервирование места для ответа и промтов,
+обрезка истории и контекста с учётом доступного пространства.
 """
 
 import asyncio
@@ -110,42 +113,48 @@ class RAGOrchestrator:
                 - confidence (float): оценка уверенности (если есть)
         """
         # ----- 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. Получение истории диалога из БД -----
-        # Запрашиваем последние 100 сообщений для данного пользователя/комнаты
         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)
+        reserved_for_answer = getattr(self.config, 'reserved_for_answer_tokens', 1000)
+        reserved_for_prompt = 0
+        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 или уменьшите резервы.")
+            available_for_history_and_context = max(100, available_for_history_and_context)
+
+        logger.debug(f"Управление токенами: max={max_context_tokens}, резерв ответа={reserved_for_answer}, промт={reserved_for_prompt}, доступно={available_for_history_and_context}")
+
+        # 3.2. Обрезка истории по токенам (сохраняем последние сообщения)
         total_tokens = 0
         truncated_history = []
-        # Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние
         for record in reversed(history):
             tokens = count_tokens(record['content'])
-            if total_tokens + tokens <= max_context_tokens:
+            if total_tokens + tokens <= available_for_history_and_context:
                 truncated_history.append(record)
                 total_tokens += tokens
             else:
                 break
-        # Восстанавливаем хронологический порядок (от старого к новому)
         truncated_history.reverse()
-        # Преобразуем в формат, ожидаемый GigaChat: [{"role": "user/assistant", "content": "..."}]
+        # Сразу обрезаем историю для использования в дальнейшем
         formatted_history = truncated_history
+        logger.debug(f"История обрезана до {len(truncated_history)} сообщений, {total_tokens} токенов")
 
-        logger.debug(
+        # ----- 4. Классификация намерений -----
-            f"История из БД: {len(truncated_history)} сообщений, "
-            f"{total_tokens} токенов из {max_context_tokens} допустимых"
-        )
-
-        # ----- 4. Классификация намерений (если не переопределена) -----
         intent = intent_override
         if intent is None and getattr(self.config, 'enable_intent_classification', True):
             intent_prompt = prompts.get('intent', '')
@@ -161,20 +170,19 @@ class RAGOrchestrator:
         else:
             intent = intent or "GENERAL"
 
-        # ----- 5. Принудительная установка SURGICAL по ключевым словам -----
+        # Принудительная установка 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 (есть файл и ключевые слова)")
 
-        # ----- 6. Обработка специализированных намерений -----
+        # ----- 5. Обработка специализированных намерений -----
         answer = None
         context = None
         ctx_for_critique = None
         synthesis_template_for_critique = None
 
-        # --- 6.1. METRICS: извлечение числовых показателей ---
+        # --- 5.1. METRICS ---
         if intent == "METRICS":
             context = await self.kb.find_relevant_info(
                 query, user_jid, room_jid,
@@ -196,7 +204,7 @@ class RAGOrchestrator:
                 else:
                     answer = "Не удалось извлечь метрики."
 
-        # --- 6.2. SUMMARY: суммаризация документа ---
+        # --- 5.2. SUMMARY ---
         elif intent == "SUMMARY":
             if not last_file_text:
                 answer = "Нет документа для суммаризации."
@@ -210,7 +218,7 @@ class RAGOrchestrator:
                     bot_config=self.config
                 )
 
-        # --- 6.3. CONTRADICTION: проверка противоречий ---
+        # --- 5.3. CONTRADICTION ---
         elif intent == "CONTRADICTION":
             context = await self.kb.find_relevant_info(
                 query, user_jid, room_jid,
@@ -236,11 +244,10 @@ class RAGOrchestrator:
                     else:
                         answer = "✅ Противоречий не обнаружено."
 
-        # --- 6.4. TEMPLATE_FILL: заполнение шаблона документа ---
+        # --- 5.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):
-                # Если текст не передан, извлекаем его из файла
                 result = await asyncio.to_thread(self.files.process_any_file, last_file_path)
                 if isinstance(result, tuple) and len(result) == 2:
                     template_text = result[0]
@@ -249,7 +256,6 @@ class RAGOrchestrator:
             if not template_text:
                 answer = "❌ Нет шаблона документа для заполнения. Загрузите файл .docx."
             else:
-                # Обрезаем шаблон до 5000 символов, чтобы не перегружать контекст
                 truncated_template = template_text[:5000]
                 search_query = f"{query}\n{truncated_template}"
                 context = await self.kb.find_relevant_info(search_query, user_jid, room_jid, top_k=15)
@@ -271,10 +277,7 @@ class RAGOrchestrator:
                     temperature=getattr(self.config, 'ai_temperature', 0.1)
                 )
 
-        # --- 6.5. SPELLCHECK: проверка орфографии ---
+        # --- 5.5. SPELLCHECK ---
-        # Внимание: эта функция требует доступа к промпту spellcheck и локальному file_service
-        # Пока оставлена заглушкой, так как требует отдельного рефакторинга.
-        # В будущем check_spelling должен быть адаптирован для работы через сервер.
         elif intent == "SPELLCHECK":
             if room_jid is not None:
                 answer = "⚠️ Проверка орфографии доступна только в личном чате."
@@ -287,7 +290,7 @@ class RAGOrchestrator:
                     # TODO: адаптировать check_spelling для работы через HTTP
                     answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
 
-        # --- 6.6. GREETING: приветствие (без контекста) ---
+        # --- 5.6. GREETING ---
         elif intent == "GREETING":
             answer = await self.giga.chat(
                 history=formatted_history,
@@ -297,12 +300,11 @@ class RAGOrchestrator:
                 temperature=getattr(self.config, 'ai_temperature', 0.1)
             )
 
-        # --- 6.7. SURGICAL: хирургическая замена в документе ---
+        # --- 5.7. SURGICAL ---
         elif intent == "SURGICAL":
             if not last_file_path:
                 answer = "❌ Нет загруженного документа для замены."
             else:
-                # Парсим запрос вида "замени X на Y"
                 import re
                 q = re.sub(r'\s+', ' ', query.lower()).strip()
                 m = re.search(r'замени\s+(.+?)\s+на\s+(.+)', q)
@@ -313,11 +315,9 @@ class RAGOrchestrator:
                 else:
                     old_word = m.group(1).strip()
                     new_word = m.group(2).strip()
-                    # Морфологическая замена с помощью pymorphy3
                     try:
                         from mawo_pymorphy3 import create_analyzer
                         morph = create_analyzer()
-                        # Генерируем все формы старого слова
                         old_forms = set()
                         parsed_old = morph.parse(old_word)[0]
                         old_forms.add(old_word)
@@ -329,7 +329,6 @@ class RAGOrchestrator:
                                 inflected = parsed_old.inflect({case, number})
                                 if inflected:
                                     old_forms.add(inflected.word)
-                        # Сопоставляем с формами нового слова
                         replacements = {}
                         for old_form in old_forms:
                             parsed_old_form = morph.parse(old_form)[0]
@@ -343,23 +342,20 @@ class RAGOrchestrator:
                             parsed_new = morph.parse(new_word)[0]
                             new_form = parsed_new.inflect(tags)
                             replacements[old_form] = new_form.word if new_form else new_word
-                        # Добавляем исходные формы
                         replacements[old_word] = new_word
                         replacements[parsed_old.normal_form] = new_word
 
-                        # Выполняем замену (синхронно в потоке)
                         new_path = await asyncio.to_thread(self.files.surgical_replace, last_file_path, replacements)
                         if new_path:
                             answer = f"Замена '{old_word}' → '{new_word}' выполнена. Файл сохранён: {new_path}"
-                            # Можно также вернуть файл клиенту, но пока только сообщаем
                         else:
                             answer = f"❌ Ошибка при замене '{old_word}' → '{new_word}'."
                     except ImportError:
                         answer = "❌ Библиотека mawo-pymorphy3 не установлена. Установите её для морфологической замены."
 
-        # ----- 7. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
+        # ----- 6. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
         if answer is None:
-            # 7.1. Расширение запроса (добавляем синонимы и аббревиатуры)
+            # 6.1. Расширение запроса
             expanded = await expand_query(
                 giga=self.giga,
                 query=query,
@@ -368,39 +364,71 @@ class RAGOrchestrator:
             )
             search_query = expanded if expanded and expanded != query else query
 
-            # 7.2. Поиск релевантного контекста в базе знаний
+            # 6.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})")
 
-            # 7.3. Переранжирование контекста (отбрасываем нерелевантные чанки)
+            # 6.3. Переранжирование контекста (если включено)
             rerank_min_length = getattr(self.config, 'rerank_min_length', 5000)
             if intent != "FACT" and len(context) > rerank_min_length:
                 context = await rerank_context(
-                    bot=None,  # передаём None, так как используем кросс-энкодер
+                    bot=None,
                     query=query,
                     context=context,
                     prompt_text=None,
                     bot_config=self.config
                 )
 
-            # 7.4. Синтез ответа с добавлением цепочки рассуждений (CoT) для расчётов и инструкций
+            # 6.4. Управление токенами для контекста и ответа
-            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 += cot_instruction
+                synthesis_template_local += 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]}")
 
-            # 7.5. Генерация ответа через GigaChat
+            # 6.6. Генерация ответа через GigaChat
             answer = await self.giga.chat(
                 history=formatted_history,
                 query=full_query,
@@ -411,10 +439,9 @@ class RAGOrchestrator:
 
             # Сохраняем контекст для самокритики
             ctx_for_critique = context
-            synthesis_template_for_critique = synthesis_template
+            synthesis_template_for_critique = synthesis_template_local
 
-        # ----- 8. Самокритика (если включена) -----
+        # ----- 7. Самокритика (если включена) -----
-        # Проверяем ответ на галлюцинации и соответствие контексту
         if (getattr(self.config, 'enable_self_critique', False) and
             intent not in ("METRICS", "SUMMARY", "CONTRADICTION") and
             ctx_for_critique is not None and
@@ -432,7 +459,6 @@ class RAGOrchestrator:
                 )
                 if not is_ok:
                     logger.warning("Ответ не прошёл самокритику, перегенерация")
-                    # Перегенерируем с тем же контекстом
                     full_query_retry = synthesis_template_for_critique.format(
                         context=ctx_for_critique,
                         query=query
@@ -444,7 +470,6 @@ class RAGOrchestrator:
                         file_id=None,
                         temperature=getattr(self.config, 'ai_temperature', 0.1)
                     )
-                    # Повторная проверка после перегенерации
                     if not await critique_answer(
                         giga=self.giga,
                         query=query,
@@ -455,24 +480,23 @@ class RAGOrchestrator:
                     ):
                         answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
 
-        # ----- 9. Сохранение истории диалога в БД -----
+        # ----- 8. Сохранение истории диалога в БД -----
-        # Сохраняем запрос пользователя и ответ бота
         await self.db.add_history(user_jid, "user", query, room_jid)
         if answer:
             await self.db.add_history(user_jid, "assistant", answer, room_jid)
 
-        # ----- 10. Извлечение источников из контекста для ответа -----
+        # ----- 9. Извлечение источников из контекста -----
         sources = []
         if context:
             for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
                 sources.append(match.group(1))
 
-        # ----- 11. Формирование результата -----
+        # ----- 10. Формирование результата -----
         return {
             "answer": answer,
             "intent": intent,
             "context": context,
-            "sources": list(set(sources)),  # удаляем дубликаты
+            "sources": list(set(sources)),
             "confidence": None
         }