fckbot/rag/query_processor.py

# -*- coding: utf-8 -*-
"""
Обработчик обычных RAG-запросов (FACT, PROCEDURE, COMPARISON, CALCULATION, GENERAL).
Выполняет расширение запроса, поиск в БЗ, переранжирование, синтез ответа,
самокритику и перегенерацию при необходимости.
"""

import logging
import re
from typing import Optional, Dict, List, Any

from rag.services.giga_client import GigaClient
from rag.services.kb_service import KBService
from rag.functions.expand_query import expand_query
from rag.functions.rerank_context import rerank_context
from rag.functions.critique_answer import critique_answer
from rag.utils.text_utils import count_tokens
from rag.config_models import AppConfig

logger = logging.getLogger(__name__)


class QueryProcessor:
    """
    Обрабатывает обычные RAG-запросы (не специализированные намерения).
    """

    def __init__(
        self,
        giga: GigaClient,
        kb: KBService,
        config: AppConfig,
        default_prompts: Optional[Dict[str, str]] = None
    ):
        """
        Инициализация процессора запросов.

        Аргументы:
            giga: клиент GigaChat
            kb: сервис базы знаний
            config: объект конфигурации (AppConfig)
            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. Расширение запроса
        expand_prompt = prompts.get('expand', '')
        expanded = await expand_query(
            giga=self.giga,
            query=query,
            prompt_text=expand_prompt,
            bot_config=self.config
        )
        search_query = expanded if expanded and expanded != query else query

        # 2. Поиск релевантного контекста в базе знаний
        top_k = self.config.rag.default_top_k
        context = await self.kb.find_relevant_info(
            search_query, user_jid, room_jid,
            top_k=top_k
        )
        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 = self.config.rag.rerank_min_length
        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. Генерация ответа
        temperature = self.config.ai.temperature
        answer = await self.giga.chat(
            history=history,
            query=full_query,
            system_prompt=system_prompt,
            file_id=None,
            temperature=temperature
        )

        # 7. Самокритика (если включена)
        if self.config.features.enable_self_critique 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=temperature
                    )
                    # Повторная проверка после перегенерации
                    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
        }