340 lines
14 KiB
Python
340 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Главный оркестратор RAG-пайплайна (фасад).
|
||
Координирует работу менеджеров: HistoryManager, IntentRouter, QueryProcessor, IndexingManager.
|
||
Принимает запрос пользователя, обрабатывает его и возвращает ответ.
|
||
"""
|
||
|
||
import logging
|
||
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.history_manager import HistoryManager
|
||
from core.intent_router import IntentRouter
|
||
from core.query_processor import QueryProcessor
|
||
from core.indexing_manager import IndexingManager
|
||
from core.functions.intent_classify import classify_intent
|
||
from core.utils.text_utils import count_tokens
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class RAGOrchestrator:
|
||
"""
|
||
Оркестратор RAG-пайплайна.
|
||
Содержит ссылки на все сервисы и менеджеры.
|
||
Предоставляет два основных метода: process_query и index_document.
|
||
"""
|
||
|
||
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: сервис эмбеддингов
|
||
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 {}
|
||
|
||
# Инициализация менеджеров
|
||
self.history_manager = HistoryManager(
|
||
db=db,
|
||
giga=giga,
|
||
config=config,
|
||
default_prompts=self.default_prompts
|
||
)
|
||
|
||
self.intent_router = IntentRouter(
|
||
giga=giga,
|
||
kb=kb,
|
||
files=files,
|
||
config=config,
|
||
default_prompts=self.default_prompts
|
||
)
|
||
|
||
self.query_processor = QueryProcessor(
|
||
giga=giga,
|
||
kb=kb,
|
||
config=config,
|
||
default_prompts=self.default_prompts
|
||
)
|
||
|
||
self.indexing_manager = IndexingManager(
|
||
kb=kb,
|
||
giga=giga,
|
||
config=config,
|
||
default_prompts=self.default_prompts
|
||
)
|
||
|
||
logger.info("RAGOrchestrator инициализирован с менеджерами")
|
||
|
||
# ------------------------------------------------------------------
|
||
# Вспомогательный метод для расчёта токенов
|
||
# ------------------------------------------------------------------
|
||
|
||
def _prepare_prompt_parts(
|
||
self,
|
||
synthesis_template: str,
|
||
system_prompt: Optional[str],
|
||
query: str,
|
||
max_total_tokens: int = 8192,
|
||
reserved_for_answer: int = 1000,
|
||
reserved_for_overhead: int = 200
|
||
) -> Dict[str, Any]:
|
||
"""
|
||
Подсчитывает токены в статичных частях промта и вычисляет,
|
||
сколько токенов остаётся для истории и контекста.
|
||
|
||
Возвращает словарь с ключами:
|
||
- available_for_history_and_context: int
|
||
- prompt_tokens: int
|
||
- system_tokens: int
|
||
- synthesis_tokens: int
|
||
- query_tokens: int
|
||
"""
|
||
system_tokens = count_tokens(system_prompt) if system_prompt else 0
|
||
synthesis_tokens = count_tokens(synthesis_template)
|
||
query_tokens = count_tokens(query)
|
||
prompt_tokens = system_tokens + synthesis_tokens + query_tokens
|
||
|
||
available = max_total_tokens - prompt_tokens - reserved_for_answer - reserved_for_overhead
|
||
if available < 0:
|
||
logger.warning(
|
||
f"Недостаточно токенов для истории и контекста: {available}. "
|
||
f"Увеличьте max_total_tokens или уменьшите размер промптов."
|
||
)
|
||
available = max(available, 100)
|
||
|
||
return {
|
||
"available_for_history_and_context": available,
|
||
"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]:
|
||
"""
|
||
Основной метод обработки запроса.
|
||
|
||
ИСТОРИЯ ДИАЛОГА ПОЛУЧАЕТСЯ ИЗ БД, а не из запроса.
|
||
Это обеспечивает единый контекст для всех клиентов.
|
||
|
||
Аргументы:
|
||
query: текст запроса пользователя
|
||
user_jid: JID пользователя (без ресурса)
|
||
room_jid: JID комнаты (None для личного чата)
|
||
prompts: словарь промптов (если None, используются default_prompts)
|
||
intent_override: принудительное переопределение намерения
|
||
last_file_path: путь к последнему загруженному файлу
|
||
last_file_text: текст последнего загруженного файла
|
||
|
||
Возвращает:
|
||
Словарь с ключами: answer, intent, context, sources, confidence, error
|
||
"""
|
||
# 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. Расчёт лимитов токенов
|
||
max_model_tokens = getattr(self.config, 'max_model_tokens', 8192)
|
||
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,
|
||
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} токенов")
|
||
|
||
# 3. Получение истории из БД
|
||
raw_history = await self.history_manager.get_history(user_jid, room_jid, limit=100)
|
||
|
||
# 4. Сжатие истории, если она слишком длинная
|
||
max_history_tokens = min(available_for_history_and_context // 2, 2000)
|
||
formatted_history = await self.history_manager.compress_history_if_needed(
|
||
raw_history,
|
||
max_tokens=max_history_tokens,
|
||
prompt_template=prompts.get('hierarchical_summary', '')
|
||
)
|
||
|
||
# 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. Маршрутизация специализированных намерений
|
||
router_result = await self.intent_router.route(
|
||
intent=intent,
|
||
query=query,
|
||
user_jid=user_jid,
|
||
room_jid=room_jid,
|
||
prompts=prompts,
|
||
last_file_path=last_file_path,
|
||
last_file_text=last_file_text,
|
||
history=formatted_history,
|
||
system_prompt=system_prompt
|
||
)
|
||
|
||
answer = None
|
||
context = ""
|
||
sources = []
|
||
|
||
if router_result is not None:
|
||
# Намерение обработано маршрутизатором
|
||
answer = router_result.get("answer")
|
||
context = router_result.get("context", "")
|
||
sources = router_result.get("sources", [])
|
||
else:
|
||
# Обычный RAG-пайплайн (GENERAL, FACT, PROCEDURE, COMPARISON, CALCULATION)
|
||
# Вычисляем, сколько токенов осталось для контекста после истории
|
||
history_tokens = sum(count_tokens(msg['content']) for msg in formatted_history)
|
||
available_for_context = available_for_history_and_context - history_tokens
|
||
available_for_context = max(available_for_context, 0)
|
||
|
||
processor_result = await self.query_processor.process(
|
||
query=query,
|
||
user_jid=user_jid,
|
||
room_jid=room_jid,
|
||
prompts=prompts,
|
||
intent=intent,
|
||
history=formatted_history,
|
||
system_prompt=system_prompt,
|
||
available_tokens_for_context=available_for_context
|
||
)
|
||
answer = processor_result.get("answer")
|
||
context = processor_result.get("context", "")
|
||
sources = processor_result.get("sources", [])
|
||
|
||
# 8. Сохранение истории диалога в БД
|
||
await self.history_manager.save_message(user_jid, "user", query, room_jid)
|
||
if answer:
|
||
await self.history_manager.save_message(user_jid, "assistant", answer, room_jid)
|
||
|
||
# 9. Формирование результата
|
||
return {
|
||
"answer": answer or "⚠️ Не удалось сгенерировать ответ.",
|
||
"intent": intent,
|
||
"context": context,
|
||
"sources": sources,
|
||
"confidence": None,
|
||
"error": 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]:
|
||
"""
|
||
Индексирует документ в базу знаний.
|
||
Делегирует работу IndexingManager.
|
||
|
||
Аргументы:
|
||
file_name: исходное имя файла
|
||
file_text: извлечённый текст
|
||
user_jid: JID владельца
|
||
room_jid: JID комнаты (None для личного)
|
||
is_global: глобальный ли документ
|
||
title: отображаемое название
|
||
metadata: дополнительные метаданные
|
||
file_hash: SHA-256 хеш содержимого
|
||
update_if_exists: заменять ли существующий документ в комнате
|
||
|
||
Возвращает:
|
||
Словарь с ключами doc_id, chunk_count, error
|
||
"""
|
||
try:
|
||
doc_id, chunk_count = await self.indexing_manager.index_document(
|
||
file_name=file_name,
|
||
file_text=file_text,
|
||
user_jid=user_jid,
|
||
room_jid=room_jid,
|
||
is_global=is_global,
|
||
title=title,
|
||
metadata=metadata,
|
||
file_hash=file_hash,
|
||
update_if_exists=update_if_exists
|
||
)
|
||
return {"doc_id": doc_id, "chunk_count": chunk_count, "error": None}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка индексации документа {file_name}: {e}", exc_info=True)
|
||
return {"doc_id": None, "chunk_count": 0, "error": str(e)} |