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 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
from core.prompt_builder import PromptBuilder
from core.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
    ) -> None:
        self.giga: GigaClient = giga
        self.kb: KBService = kb
        self.config: AppConfig = config
        self.default_prompts: Dict[str, str] = default_prompts or {}
        self.prompt_builder: PromptBuilder = PromptBuilder(config)

    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: str = await expand_query(
            giga=self.giga,
            query=query,
            prompt_text=prompts.get('expand', ''),
            bot_config=self.config
        )
        search_query: str = expanded if expanded and expanded != query else query

        # 2. Поиск релевантного контекста в базе знаний
        context: str = 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: int = count_tokens(context)
            if context_tokens > available_tokens_for_context:
                logger.warning(
                    f"Контекст слишком длинный ({context_tokens} токенов), "
                    f"обрезаем до {available_tokens_for_context}"
                )
                max_context_chars: int = int(available_tokens_for_context * 3.5)
                if max_context_chars > 0:
                    context = context[:max_context_chars]
                else:
                    context = ""

        # 4. Переранжирование контекста (если включено и контекст достаточно длинный)
        rerank_min_length: int = 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. Формирование промта с помощью PromptBuilder
        extra_instructions: str = ""
        if intent in ("CALCULATION", "PROCEDURE"):
            extra_instructions = (
                "Пожалуйста, покажи пошаговое решение перед итоговым ответом. "
                "Опиши каждый шаг вычислений или действий в логической последовательности. "
                "После всех шагов дай итоговый ответ."
            )

        synthesis_template: str = prompts.get('synthesis', '')
        if synthesis_template and not extra_instructions:
            extra_instructions = synthesis_template.format(context=context, query=query) if '{context}' in synthesis_template else synthesis_template

        prompt: str = self.prompt_builder.build_prompt(
            query=query,
            intent=intent,
            context=context,
            history=history,
            system_prompt=system_prompt,
            extra_instructions=extra_instructions
        )
        logger.debug(f"Сформированный промт (первые 500 символов): {prompt[:500]}")

        # 6. Генерация ответа
        answer: str = await self.giga.chat(
            history=[],
            query=prompt,
            system_prompt=None,
            file_id=None,
            temperature=getattr(self.config, 'ai_temperature', 0.1)
        )

        # 7. Самокритика (если включена)
        if getattr(self.config, 'enable_self_critique', False) and context:
            critique_prompt: str = prompts.get('critique', '')
            if critique_prompt:
                logger.debug("Запуск самокритики")
                is_ok: bool = 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("Ответ не прошёл самокритику, перегенерация")
                    answer = await self.giga.chat(
                        history=[],
                        query=prompt,
                        system_prompt=None,
                        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: List[str] = []
        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
        }