Files
fckbot/rag/rag_orchestrator.py
2026-06-30 12:21:07 +00:00

599 lines
30 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- 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 инициализирован")
# ------------------------------------------------------------------
# Вспомогательный метод для расчёта токенов с резервированием
# ------------------------------------------------------------------
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,
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. Получение истории диалога из БД (необрезанной) -----
history = await self.db.get_history(user_jid, room_jid, limit=100)
# ----- 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_overhead = getattr(self.config, 'reserved_for_overhead_tokens', 200)
# Используем вспомогательный метод для расчёта
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"Доступно для истории и контекста: {available_for_history_and_context} токенов, "
f"промпт: {token_info['prompt_tokens']} токенов"
)
# ----- 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_history_tokens + tokens <= history_limit:
truncated_history.append(record)
total_history_tokens += tokens
else:
break
truncated_history.reverse() # восстанавливаем хронологический порядок
# Преобразуем в формат, ожидаемый 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', '')
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"
# ----- 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 (есть файл и ключевые слова)")
# ----- 7. Обработка специализированных намерений -----
answer = None
context = None
ctx_for_critique = None
synthesis_template_for_critique = None
# --- 7.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 = "Не удалось извлечь метрики."
# --- 7.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
)
# --- 7.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 = "✅ Противоречий не обнаружено."
# --- 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):
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)
)
# --- 7.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 = "⚠️ Проверка орфографии требует доработки (передача промптов)."
# --- 7.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)
)
# --- 7.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 не установлена. Установите её для морфологической замены."
# ----- 8. Обычный RAG (для GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION) -----
if answer is None:
# 8.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
# 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})")
# 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(
bot=None,
query=query,
context=context,
prompt_text=None,
bot_config=self.config
)
# 8.5. Синтез ответа с добавлением цепочки рассуждений (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]}")
# 8.6. Генерация ответа
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
# ----- 9. Самокритика (если включена) -----
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 = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
# ----- 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)
# ----- 11. Извлечение источников из контекста -----
sources = []
if context:
for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context):
sources.append(match.group(1))
# ----- 12. Формирование результата -----
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}