From 9f23c4ed08a9fa14246987f81b71d7c1e64d63e4 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 12:03:11 +0000 Subject: [PATCH] =?UTF-8?q?=D0=A0=D0=B5=D0=B4=D0=B0=D0=BA=D1=82=D0=B8?= =?UTF-8?q?=D1=80=D0=BE=D0=B2=D0=B0=D1=82=D1=8C=20rag=5Forchestrator.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- rag/rag_orchestrator.py | 144 +++++++++++++++++++++++----------------- 1 file changed, 84 insertions(+), 60 deletions(-) diff --git a/rag/rag_orchestrator.py b/rag/rag_orchestrator.py index dbc1a26..64fe329 100644 --- a/rag/rag_orchestrator.py +++ b/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', '') - if not synthesis_template: - synthesis_template = "{context}\n\n{query}\n\nОтвет:" + synthesis_template = prompts.get('synthesis', '{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( - f"История из БД: {len(truncated_history)} сообщений, " - f"{total_tokens} токенов из {max_context_tokens} допустимых" - ) - - # ----- 4. Классификация намерений (если не переопределена) ----- + # ----- 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: проверка орфографии --- - # Внимание: эта функция требует доступа к промпту spellcheck и локальному file_service - # Пока оставлена заглушкой, так как требует отдельного рефакторинга. - # В будущем check_spelling должен быть адаптирован для работы через сервер. + # --- 5.5. SPELLCHECK --- 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) для расчётов и инструкций - synthesis_template = synthesis_template + # 6.4. Управление токенами для контекста и ответа + # Вычисляем, сколько токенов уже занято историей + 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.format(context=context, query=query) + full_query = synthesis_template_local.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 }