From 2008bb12839355dedc4286a7b919552261b20550 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 12:21:07 +0000 Subject: [PATCH] Update file rag_orchestrator.py --- rag/rag_orchestrator.py | 237 ++++++++++++++++++++++++++-------------- 1 file changed, 153 insertions(+), 84 deletions(-) diff --git a/rag/rag_orchestrator.py b/rag/rag_orchestrator.py index 64fe329..8522495 100644 --- a/rag/rag_orchestrator.py +++ b/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.1. Резервирование места для ответа и промтов - max_context_tokens = getattr(self.config, 'max_context_tokens', 3000) + # ----- 3. Расчёт лимитов токенов (НОВОЕ) ----- + # Параметры из конфига + 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 - 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 + reserved_for_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200) - 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) + # Используем вспомогательный метод для расчёта + token_info = self._prepare_prompt_parts( + 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. Обрезка истории по токенам (сохраняем последние сообщения) - total_tokens = 0 + # ----- 4. Обрезка истории по токенам (с учётом лимита) ----- + # Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние 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() - # Сразу обрезаем историю для использования в дальнейшем - formatted_history = truncated_history - logger.debug(f"История обрезана до {len(truncated_history)} сообщений, {total_tokens} токенов") + truncated_history.reverse() # восстанавливаем хронологический порядок - # ----- 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. Управление токенами для контекста и ответа - # Вычисляем, сколько токенов уже занято историей - 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 + # 8.5. Синтез ответа с добавлением цепочки рассуждений (CoT) + synthesis_template = 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,