# -*- coding: utf-8 -*- """ Главный оркестратор RAG-пайплайна (фасад). Координирует работу менеджеров: HistoryManager, IntentRouter, QueryProcessor, IndexingManager. Принимает запрос пользователя, обрабатывает его и возвращает ответ. """ import logging from typing import Optional, Dict, List, Any from .services.postgres_service import PostgresService from .services.qdrant_service import QdrantService from .services.embedding_service import EmbeddingService from .services.kb_service import KBService from .services.giga_client import GigaClient from .services.file_service import FileService from .history_manager import HistoryManager from .intent_router import IntentRouter from .query_processor import QueryProcessor from .indexing_manager import IndexingManager from .functions.intent_classify import classify_intent from .utils.text_utils import count_tokens from .config_models import AppConfig 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: AppConfig, default_prompts: Optional[Dict[str, str]] = None ): """ Инициализация оркестратора. Аргументы: db: сервис PostgreSQL qdrant: сервис Qdrant embedding: сервис эмбеддингов kb: сервис базы знаний giga: клиент GigaChat files: сервис файлов config: объект конфигурации (AppConfig) 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 = self.config.max_model_tokens reserved_for_answer = self.config.reserved_for_answer_tokens reserved_for_overhead = self.config.reserved_for_overhead_tokens 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 self.config.enable_intent_classification: 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 = 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)}