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 .services.giga_client import GigaClient
from .services.kb_service import KBService
from .functions.expand_query import expand_query
from .functions.rerank_context import rerank_context
from .functions.critique_answer import critique_answer
from .utils.text_utils import count_tokens
from .prompt_builder import PromptBuilder
from .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
    ):
        self.giga = giga
        self.kb = kb
        self.config = config
        self.default_prompts = default_prompts or {}
        self.prompt_builder = 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-пайплайн для обычного запроса.
        """
        # 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. Формирование промта с помощью PromptBuilder (динамические few-shot, сэндвич)
        extra_instructions = ""
        if intent in ("CALCULATION", "PROCEDURE"):
            extra_instructions = (
                "Пожалуйста, покажи пошаговое решение перед итоговым ответом. "
                "Опиши каждый шаг вычислений или действий в логической последовательности. "
                "После всех шагов дай итоговый ответ."
            )

        # Если synthesis_template задан, добавим его в extra_instructions (но PromptBuilder уже использует стандартные инструкции)
        synthesis_template = 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 = 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 = 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 = 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("Ответ не прошёл самокритику, перегенерация")
                    # Перегенерируем с тем же промтом
                    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 = []
        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
        }