diff --git a/core/rag_orchestrator.py b/core/rag_orchestrator.py new file mode 100644 index 0000000..dbc1a26 --- /dev/null +++ b/core/rag_orchestrator.py @@ -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} \ No newline at end of file