Add new file

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
+        }