From 2e776beece415a3fd6dbb65e56e25267b17e8282 Mon Sep 17 00:00:00 2001 From: Markov Andrey Date: Tue, 30 Jun 2026 13:50:30 +0000 Subject: [PATCH] Add new file --- rag/query_processor.py | 186 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 186 insertions(+) create mode 100644 rag/query_processor.py diff --git a/rag/query_processor.py b/rag/query_processor.py new file mode 100644 index 0000000..3f541f3 --- /dev/null +++ b/rag/query_processor.py @@ -0,0 +1,186 @@ +# -*- coding: utf-8 -*- +""" +Обработчик обычных RAG-запросов (FACT, PROCEDURE, COMPARISON, CALCULATION, GENERAL). +Выполняет расширение запроса, поиск в БЗ, переранжирование, синтез ответа, +самокритику и перегенерацию при необходимости. +""" + +import logging +import re +from typing import Optional, Dict, List, Any + +from core.services.giga_client import GigaClient +from core.services.kb_service import KBService +from core.functions.expand_query import expand_query +from core.functions.rerank_context import rerank_context +from core.functions.critique_answer import critique_answer +from core.utils.text_utils import count_tokens + +logger = logging.getLogger(__name__) + + +class QueryProcessor: + """ + Обрабатывает обычные RAG-запросы (не специализированные намерения). + """ + + def __init__( + self, + giga: GigaClient, + kb: KBService, + config, + default_prompts: Optional[Dict[str, str]] = None + ): + """ + Инициализация процессора запросов. + + Аргументы: + giga: клиент GigaChat + kb: сервис базы знаний + config: объект конфигурации + default_prompts: словарь промптов по умолчанию + """ + self.giga = giga + self.kb = kb + self.config = config + self.default_prompts = default_prompts or {} + + async def process( + self, + query: str, + user_jid: str, + room_jid: Optional[str], + prompts: Dict[str, str], + intent: str, + history: List[Dict[str, str]], + system_prompt: Optional[str], + available_tokens_for_context: int, + ) -> Dict[str, Any]: + """ + Выполняет полный RAG-пайплайн для обычного запроса. + + Аргументы: + query: текст запроса пользователя + user_jid: JID пользователя + room_jid: JID комнаты (None для личного чата) + prompts: словарь промптов (expand, synthesis, critique, ...) + intent: код намерения (для выбора стратегии) + history: история диалога (уже сжатая, если нужно) + system_prompt: системный промпт + available_tokens_for_context: сколько токенов доступно для контекста (после вычета истории и промптов) + + Возвращает: + Словарь с ключами 'answer', 'context', 'sources', 'confidence' + """ + # 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 + + # 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})") + + # 3. Обрезка контекста по токенам + if context: + context_tokens = count_tokens(context) + if context_tokens > available_tokens_for_context: + logger.warning( + f"Контекст слишком длинный ({context_tokens} токенов), " + f"обрезаем до {available_tokens_for_context}" + ) + max_context_chars = int(available_tokens_for_context * 3.5) + if max_context_chars > 0: + context = context[:max_context_chars] + else: + context = "" + + # 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 + ) + + # 5. Синтез ответа с добавлением цепочки рассуждений (CoT) для CALCULATION и PROCEDURE + synthesis_template = prompts.get('synthesis', '') + if not synthesis_template: + synthesis_template = "{context}\n\n{query}\n\nОтвет:" + + 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]}") + + # 6. Генерация ответа + answer = await self.giga.chat( + history=history, + query=full_query, + system_prompt=system_prompt, + file_id=None, + temperature=getattr(self.config, 'ai_temperature', 0.1) + ) + + # 7. Самокритика (если включена) + if getattr(self.config, 'enable_self_critique', False) and context: + critique_prompt = prompts.get('critique', '') + if critique_prompt: + logger.debug("Запуск самокритики") + is_ok = await critique_answer( + giga=self.giga, + query=query, + context=context, + answer=answer, + prompt_text=critique_prompt, + bot_config=self.config + ) + if not is_ok: + logger.warning("Ответ не прошёл самокритику, перегенерация") + full_query_retry = synthesis_template.format(context=context, query=query) + answer = await self.giga.chat( + history=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=context, + answer=answer, + prompt_text=critique_prompt, + bot_config=self.config + ): + answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные." + + # 8. Извлечение источников + sources = [] + if context: + for match in re.finditer(r'\[источник:\s*([^\]]+)\]', context): + sources.append(match.group(1)) + + return { + "answer": answer, + "context": context, + "sources": list(set(sources)), + "confidence": None + } \ No newline at end of file