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

This commit is contained in:
Markov Andrey
2026-06-30 12:03:11 +00:00
parent 4b57abf269
commit 9f23c4ed08

View File

@@ -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
}