Редактировать rag_orchestrator.py
This commit is contained in:
Markov Andrey
506
rag/rag_orchestrator.py
Normal file
506
rag/rag_orchestrator.py
Normal file
@@ -0,0 +1,506 @@
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
Главный оркестратор RAG-пайплайна.
|
||||
Принимает запрос пользователя, получает историю из БД, выполняет все этапы:
|
||||
классификацию, расширение, поиск, переранжирование, синтез, самокритику.
|
||||
Сохраняет историю в БД.
|
||||
|
||||
Этот модуль не зависит от XMPP и может использоваться как в ботах,
|
||||
так и в отдельном RAG-ядре (API-сервисе).
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
from typing import Optional, Dict, List, Any
|
||||
|
||||
# Импорт сервисов и функций
|
||||
from core.services.postgres_service import PostgresService
|
||||
from core.services.qdrant_service import QdrantService
|
||||
from core.services.embedding_service import EmbeddingService
|
||||
from core.services.kb_service import KBService
|
||||
from core.services.giga_client import GigaClient
|
||||
from core.services.file_service import FileService
|
||||
from core.functions.intent_classify import classify_intent
|
||||
from core.functions.expand_query import expand_query
|
||||
from core.functions.extract_metrics import extract_metrics
|
||||
from core.functions.summarize_document import summarize_document
|
||||
from core.functions.check_consistency import check_consistency
|
||||
from core.functions.critique_answer import critique_answer
|
||||
from core.functions.rerank_context import rerank_context
|
||||
from core.functions.check_spelling import check_spelling
|
||||
from core.utils.text_utils import count_tokens
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RAGOrchestrator:
|
||||
"""
|
||||
Оркестратор RAG-пайплайна.
|
||||
Содержит ссылки на все сервисы и выполняет полный цикл обработки запроса.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db: PostgresService,
|
||||
qdrant: QdrantService,
|
||||
embedding: EmbeddingService,
|
||||
kb: KBService,
|
||||
giga: GigaClient,
|
||||
files: FileService,
|
||||
config,
|
||||
default_prompts: Optional[Dict[str, str]] = None
|
||||
):
|
||||
"""
|
||||
Инициализация оркестратора.
|
||||
|
||||
Аргументы:
|
||||
db: сервис PostgreSQL (для истории и метаданных)
|
||||
qdrant: сервис Qdrant (векторный поиск)
|
||||
embedding: сервис эмбеддингов (GigaChat)
|
||||
kb: сервис базы знаний (индексация, поиск)
|
||||
giga: клиент GigaChat (генерация)
|
||||
files: сервис файлов (извлечение текста)
|
||||
config: объект конфигурации (BotConfig)
|
||||
default_prompts: словарь промптов по умолчанию
|
||||
"""
|
||||
self.db = db
|
||||
self.qdrant = qdrant
|
||||
self.embedding = embedding
|
||||
self.kb = kb
|
||||
self.giga = giga
|
||||
self.files = files
|
||||
self.config = config
|
||||
self.default_prompts = default_prompts or {}
|
||||
logger.info("RAGOrchestrator инициализирован")
|
||||
|
||||
async def process_query(
|
||||
self,
|
||||
query: str,
|
||||
user_jid: str,
|
||||
room_jid: Optional[str],
|
||||
prompts: Optional[Dict[str, str]] = None,
|
||||
intent_override: Optional[str] = None,
|
||||
last_file_path: Optional[str] = None,
|
||||
last_file_text: Optional[str] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Основной метод обработки запроса.
|
||||
|
||||
ИСТОРИЯ ДИАЛОГА ПОЛУЧАЕТСЯ ИЗ БД, а не из запроса.
|
||||
Это обеспечивает единый контекст для всех клиентов (XMPP, Telegram и т.д.).
|
||||
|
||||
Аргументы:
|
||||
query (str): текст запроса пользователя
|
||||
user_jid (str): JID пользователя (без ресурса)
|
||||
room_jid (Optional[str]): JID комнаты (None для личного чата)
|
||||
prompts (Optional[Dict[str, str]]): словарь промптов для текущего запроса.
|
||||
Если не передан, используются default_prompts.
|
||||
intent_override (Optional[str]): принудительное переопределение намерения
|
||||
last_file_path (Optional[str]): путь к последнему загруженному файлу
|
||||
last_file_text (Optional[str]): текст последнего загруженного файла
|
||||
|
||||
Возвращает:
|
||||
Dict[str, Any]: словарь с ключами:
|
||||
- answer (str): итоговый ответ
|
||||
- intent (str): распознанное намерение
|
||||
- context (str): использованный контекст (для отладки)
|
||||
- sources (List[str]): список источников
|
||||
- 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Ответ:"
|
||||
system_prompt = prompts.get('system', None)
|
||||
|
||||
# ----- 2. Получение истории диалога из БД -----
|
||||
# Запрашиваем последние 100 сообщений для данного пользователя/комнаты
|
||||
history = await self.db.get_history(user_jid, room_jid, limit=100)
|
||||
|
||||
# ----- 3. Обрезка истории по токенам (чтобы не превысить лимит модели) -----
|
||||
max_context_tokens = getattr(self.config, 'max_context_tokens', 3000)
|
||||
total_tokens = 0
|
||||
truncated_history = []
|
||||
# Идём с конца (свежие сообщения) к началу, чтобы сохранить самые последние
|
||||
for record in reversed(history):
|
||||
tokens = count_tokens(record['content'])
|
||||
if total_tokens + tokens <= max_context_tokens:
|
||||
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)} сообщений, "
|
||||
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', '')
|
||||
if intent_prompt:
|
||||
intent = await classify_intent(
|
||||
giga=self.giga,
|
||||
query=query,
|
||||
prompt_text=intent_prompt,
|
||||
bot_config=self.config
|
||||
)
|
||||
else:
|
||||
intent = "GENERAL"
|
||||
else:
|
||||
intent = intent or "GENERAL"
|
||||
|
||||
# ----- 5. Принудительная установка 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. Обработка специализированных намерений -----
|
||||
answer = None
|
||||
context = None
|
||||
ctx_for_critique = None
|
||||
synthesis_template_for_critique = None
|
||||
|
||||
# --- 6.1. METRICS: извлечение числовых показателей ---
|
||||
if intent == "METRICS":
|
||||
context = await self.kb.find_relevant_info(
|
||||
query, user_jid, room_jid,
|
||||
top_k=getattr(self.config, 'rag_metrics_top_k', 30)
|
||||
)
|
||||
if not context:
|
||||
answer = "Не найдено данных для извлечения метрик."
|
||||
else:
|
||||
metrics_prompt = prompts.get('metrics', '')
|
||||
metrics = await extract_metrics(
|
||||
giga=self.giga,
|
||||
context=context,
|
||||
prompt_text=metrics_prompt,
|
||||
bot_config=self.config
|
||||
)
|
||||
if metrics:
|
||||
lines = [f"- {m.get('metric_name')}: {m.get('value')} {m.get('unit', '')}" for m in metrics[:10]]
|
||||
answer = "📊 **Извлечённые метрики:**\n" + "\n".join(lines)
|
||||
else:
|
||||
answer = "Не удалось извлечь метрики."
|
||||
|
||||
# --- 6.2. SUMMARY: суммаризация документа ---
|
||||
elif intent == "SUMMARY":
|
||||
if not last_file_text:
|
||||
answer = "Нет документа для суммаризации."
|
||||
else:
|
||||
summary_prompt = prompts.get('summary', '')
|
||||
answer = await summarize_document(
|
||||
giga=self.giga,
|
||||
text=last_file_text,
|
||||
title="Ваш документ",
|
||||
prompt_text=summary_prompt,
|
||||
bot_config=self.config
|
||||
)
|
||||
|
||||
# --- 6.3. CONTRADICTION: проверка противоречий ---
|
||||
elif intent == "CONTRADICTION":
|
||||
context = await self.kb.find_relevant_info(
|
||||
query, user_jid, room_jid,
|
||||
top_k=getattr(self.config, 'rag_contradiction_top_k', 10)
|
||||
)
|
||||
if not context:
|
||||
answer = "Недостаточно данных для проверки противоречий."
|
||||
else:
|
||||
chunks = [c.strip() for c in context.split("\n\n") if c.strip()]
|
||||
if len(chunks) < 2:
|
||||
answer = "Недостаточно фрагментов."
|
||||
else:
|
||||
consistency_prompt = prompts.get('consistency', '')
|
||||
consistency = await check_consistency(
|
||||
giga=self.giga,
|
||||
chunks=chunks,
|
||||
query=query,
|
||||
prompt_text=consistency_prompt,
|
||||
bot_config=self.config
|
||||
)
|
||||
if "[CONFLICT]" in consistency:
|
||||
answer = f"⚠️ **Обнаружены противоречия:**\n{consistency}"
|
||||
else:
|
||||
answer = "✅ Противоречий не обнаружено."
|
||||
|
||||
# --- 6.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]
|
||||
else:
|
||||
template_text = str(result)
|
||||
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)
|
||||
fill_prompt = prompts.get('generate_document', '')
|
||||
if not fill_prompt:
|
||||
fill_prompt = "Заполни плейсхолдеры, используя базу знаний."
|
||||
prompt = (
|
||||
f"Перед тобой шаблон документа. Заполни плейсхолдеры, используя ТОЛЬКО базу знаний.\n\n"
|
||||
f"[ТЕКСТ ШАБЛОНА]:\n{truncated_template}\n\n"
|
||||
f"[ДАННЫЕ ИЗ БАЗЫ ЗНАНИЙ]:\n{context}\n\n"
|
||||
f"Инструкция по заполнению:\n{fill_prompt}\n\n"
|
||||
f"Формат ответа: [SURGICAL_REPLACE]\nинструкция_из_шаблона ||| текст_из_БЗ\n[/SURGICAL_REPLACE]"
|
||||
)
|
||||
answer = await self.giga.chat(
|
||||
history=formatted_history,
|
||||
query=prompt,
|
||||
system_prompt=system_prompt,
|
||||
file_id=None,
|
||||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||||
)
|
||||
|
||||
# --- 6.5. SPELLCHECK: проверка орфографии ---
|
||||
# Внимание: эта функция требует доступа к промпту spellcheck и локальному file_service
|
||||
# Пока оставлена заглушкой, так как требует отдельного рефакторинга.
|
||||
# В будущем check_spelling должен быть адаптирован для работы через сервер.
|
||||
elif intent == "SPELLCHECK":
|
||||
if room_jid is not None:
|
||||
answer = "⚠️ Проверка орфографии доступна только в личном чате."
|
||||
else:
|
||||
if not last_file_text or not last_file_path:
|
||||
answer = "❌ Нет документа для проверки. Сначала отправьте файл .docx."
|
||||
elif not last_file_path.lower().endswith('.docx'):
|
||||
answer = "❌ Проверка орфографии поддерживается только для файлов .docx."
|
||||
else:
|
||||
# TODO: адаптировать check_spelling для работы через HTTP
|
||||
answer = "⚠️ Проверка орфографии требует доработки (передача промптов)."
|
||||
|
||||
# --- 6.6. GREETING: приветствие (без контекста) ---
|
||||
elif intent == "GREETING":
|
||||
answer = await self.giga.chat(
|
||||
history=formatted_history,
|
||||
query=query,
|
||||
system_prompt=system_prompt,
|
||||
file_id=None,
|
||||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||||
)
|
||||
|
||||
# --- 6.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)
|
||||
if not m:
|
||||
m = re.search(r'заменить\s+(.+?)\s+на\s+(.+)', q)
|
||||
if not m:
|
||||
answer = "❌ Не удалось распознать, что на что заменять. Используйте формат: замени слово на слово"
|
||||
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)
|
||||
old_forms.add(parsed_old.normal_form)
|
||||
cases = ['nomn', 'gent', 'datv', 'accs', 'ablt', 'loct']
|
||||
numbers = ['sing', 'plur']
|
||||
for number in numbers:
|
||||
for case in cases:
|
||||
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]
|
||||
tags = set()
|
||||
if parsed_old_form.tag.case:
|
||||
tags.add(parsed_old_form.tag.case)
|
||||
if parsed_old_form.tag.number:
|
||||
tags.add(parsed_old_form.tag.number)
|
||||
if parsed_old_form.tag.gender:
|
||||
tags.add(parsed_old_form.tag.gender)
|
||||
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) -----
|
||||
if answer is None:
|
||||
# 7.1. Расширение запроса (добавляем синонимы и аббревиатуры)
|
||||
expanded = await expand_query(
|
||||
giga=self.giga,
|
||||
query=query,
|
||||
prompt_text=prompts.get('expand', ''),
|
||||
bot_config=self.config
|
||||
)
|
||||
search_query = expanded if expanded and expanded != query else query
|
||||
|
||||
# 7.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. Переранжирование контекста (отбрасываем нерелевантные чанки)
|
||||
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, так как используем кросс-энкодер
|
||||
query=query,
|
||||
context=context,
|
||||
prompt_text=None,
|
||||
bot_config=self.config
|
||||
)
|
||||
|
||||
# 7.4. Синтез ответа с добавлением цепочки рассуждений (CoT) для расчётов и инструкций
|
||||
synthesis_template = synthesis_template
|
||||
if intent in ("CALCULATION", "PROCEDURE"):
|
||||
cot_instruction = (
|
||||
"\n\nПожалуйста, покажи пошаговое решение перед итоговым ответом. "
|
||||
"Опиши каждый шаг вычислений или действий в логической последовательности. "
|
||||
"После всех шагов дай итоговый ответ."
|
||||
)
|
||||
synthesis_template += cot_instruction
|
||||
|
||||
# Подставляем контекст и запрос в шаблон
|
||||
full_query = synthesis_template.format(context=context, query=query)
|
||||
logger.debug(f"Полный запрос к GigaChat (первые 500 символов): {full_query[:500]}")
|
||||
|
||||
# 7.5. Генерация ответа через GigaChat
|
||||
answer = await self.giga.chat(
|
||||
history=formatted_history,
|
||||
query=full_query,
|
||||
system_prompt=system_prompt,
|
||||
file_id=None,
|
||||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||||
)
|
||||
|
||||
# Сохраняем контекст для самокритики
|
||||
ctx_for_critique = context
|
||||
synthesis_template_for_critique = synthesis_template
|
||||
|
||||
# ----- 8. Самокритика (если включена) -----
|
||||
# Проверяем ответ на галлюцинации и соответствие контексту
|
||||
if (getattr(self.config, 'enable_self_critique', False) and
|
||||
intent not in ("METRICS", "SUMMARY", "CONTRADICTION") and
|
||||
ctx_for_critique is not None and
|
||||
answer is not None):
|
||||
critique_prompt = prompts.get('critique', '')
|
||||
if critique_prompt:
|
||||
logger.debug("Запуск самокритики")
|
||||
is_ok = await critique_answer(
|
||||
giga=self.giga,
|
||||
query=query,
|
||||
context=ctx_for_critique,
|
||||
answer=answer,
|
||||
prompt_text=critique_prompt,
|
||||
bot_config=self.config
|
||||
)
|
||||
if not is_ok:
|
||||
logger.warning("Ответ не прошёл самокритику, перегенерация")
|
||||
# Перегенерируем с тем же контекстом
|
||||
full_query_retry = synthesis_template_for_critique.format(
|
||||
context=ctx_for_critique,
|
||||
query=query
|
||||
)
|
||||
answer = await self.giga.chat(
|
||||
history=formatted_history,
|
||||
query=full_query_retry,
|
||||
system_prompt=system_prompt,
|
||||
file_id=None,
|
||||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||||
)
|
||||
# Повторная проверка после перегенерации
|
||||
if not await critique_answer(
|
||||
giga=self.giga,
|
||||
query=query,
|
||||
context=ctx_for_critique,
|
||||
answer=answer,
|
||||
prompt_text=critique_prompt,
|
||||
bot_config=self.config
|
||||
):
|
||||
answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
|
||||
|
||||
# ----- 9. Сохранение истории диалога в БД -----
|
||||
# Сохраняем запрос пользователя и ответ бота
|
||||
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. Извлечение источников из контекста для ответа -----
|
||||
sources = []
|
||||
if context:
|
||||
for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
|
||||
sources.append(match.group(1))
|
||||
|
||||
# ----- 11. Формирование результата -----
|
||||
return {
|
||||
"answer": answer,
|
||||
"intent": intent,
|
||||
"context": context,
|
||||
"sources": list(set(sources)), # удаляем дубликаты
|
||||
"confidence": None
|
||||
}
|
||||
|
||||
async def index_document(
|
||||
self,
|
||||
file_name: str,
|
||||
file_text: str,
|
||||
user_jid: str,
|
||||
room_jid: Optional[str],
|
||||
is_global: bool = False,
|
||||
title: Optional[str] = None,
|
||||
metadata: Optional[Dict] = None,
|
||||
file_hash: Optional[str] = None,
|
||||
update_if_exists: bool = True
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Индексация документа в базу знаний.
|
||||
Вызывается из HTTP-эндпоинта /rag/index.
|
||||
"""
|
||||
doc_id, chunk_count = await self.kb.add_document(
|
||||
file_name=file_name,
|
||||
file_text=file_text,
|
||||
user_jid=user_jid,
|
||||
is_global=is_global,
|
||||
title=title,
|
||||
metadata=metadata,
|
||||
room_jid=room_jid,
|
||||
file_hash=file_hash,
|
||||
update_if_exists=update_if_exists
|
||||
)
|
||||
return {"doc_id": doc_id, "chunk_count": chunk_count}
|
||||
Reference in New Issue
Block a user