Add new file
This commit is contained in:
Markov Andrey
186
rag/query_processor.py
Normal file
186
rag/query_processor.py
Normal file
@@ -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
|
||||
}
|
||||
Reference in New Issue
Block a user