530 lines
28 KiB
Python
530 lines
28 KiB
Python
# -*- 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', '{context}\n\n{query}\n\nОтвет:')
|
||
system_prompt = prompts.get('system', None)
|
||
|
||
# ----- 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)
|
||
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 <= available_for_history_and_context:
|
||
truncated_history.append(record)
|
||
total_tokens += tokens
|
||
else:
|
||
break
|
||
truncated_history.reverse()
|
||
# Сразу обрезаем историю для использования в дальнейшем
|
||
formatted_history = truncated_history
|
||
logger.debug(f"История обрезана до {len(truncated_history)} сообщений, {total_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"
|
||
|
||
# Принудительная установка 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. Обработка специализированных намерений -----
|
||
answer = None
|
||
context = None
|
||
ctx_for_critique = None
|
||
synthesis_template_for_critique = None
|
||
|
||
# --- 5.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 = "Не удалось извлечь метрики."
|
||
|
||
# --- 5.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
|
||
)
|
||
|
||
# --- 5.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 = "✅ Противоречий не обнаружено."
|
||
|
||
# --- 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]
|
||
else:
|
||
template_text = str(result)
|
||
if not template_text:
|
||
answer = "❌ Нет шаблона документа для заполнения. Загрузите файл .docx."
|
||
else:
|
||
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)
|
||
)
|
||
|
||
# --- 5.5. SPELLCHECK ---
|
||
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 = "⚠️ Проверка орфографии требует доработки (передача промптов)."
|
||
|
||
# --- 5.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)
|
||
)
|
||
|
||
# --- 5.7. SURGICAL ---
|
||
elif intent == "SURGICAL":
|
||
if not last_file_path:
|
||
answer = "❌ Нет загруженного документа для замены."
|
||
else:
|
||
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()
|
||
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 не установлена. Установите её для морфологической замены."
|
||
|
||
# ----- 6. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
|
||
if answer is None:
|
||
# 6.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
|
||
|
||
# 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})")
|
||
|
||
# 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,
|
||
query=query,
|
||
context=context,
|
||
prompt_text=None,
|
||
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
|
||
if intent in ("CALCULATION", "PROCEDURE"):
|
||
cot_instruction = (
|
||
"\n\nПожалуйста, покажи пошаговое решение перед итоговым ответом. "
|
||
"Опиши каждый шаг вычислений или действий в логической последовательности. "
|
||
"После всех шагов дай итоговый ответ."
|
||
)
|
||
synthesis_template_local += cot_instruction
|
||
|
||
full_query = synthesis_template_local.format(context=context, query=query)
|
||
logger.debug(f"Полный запрос к GigaChat (первые 500 символов): {full_query[:500]}")
|
||
|
||
# 6.6. Генерация ответа через 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_local
|
||
|
||
# ----- 7. Самокритика (если включена) -----
|
||
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 = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
|
||
|
||
# ----- 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)
|
||
|
||
# ----- 9. Извлечение источников из контекста -----
|
||
sources = []
|
||
if context:
|
||
for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
|
||
sources.append(match.group(1))
|
||
|
||
# ----- 10. Формирование результата -----
|
||
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} |