365 lines
17 KiB
Python
365 lines
17 KiB
Python
# ============================================================
|
||
# rag/query_processor.py
|
||
# Полная версия с интеграцией всех новых функций
|
||
# ============================================================
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
Обработчик обычных RAG-запросов (FACT, PROCEDURE, COMPARISON, CALCULATION, GENERAL).
|
||
Выполняет расширение запроса, поиск в БЗ, переранжирование, синтез ответа,
|
||
самокритику и перегенерацию при необходимости.
|
||
|
||
ДОБАВЛЕНО:
|
||
- Интеграция SOMA-анализа (LLM-as-judge)
|
||
- Интеграция ReAct (вызов внешних инструментов)
|
||
- Интеграция планирования действий (generate_plan)
|
||
- Интеграция многоагентности (AgentCoordinator)
|
||
- Реальные обработчики инструментов (ToolHandler)
|
||
"""
|
||
|
||
import logging
|
||
import re
|
||
import json
|
||
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
|
||
|
||
# Импорты новых модулей (интеграция)
|
||
from .functions.soma_evaluate import soma_evaluate
|
||
from .functions.react import react_loop
|
||
from .functions.plan_generation import generate_plan
|
||
from .agents.coordinator import AgentCoordinator
|
||
from .functions.tools import ToolHandler
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class QueryProcessor:
|
||
"""
|
||
Обрабатывает обычные RAG-запросы (не специализированные намерения).
|
||
Интегрирует SOMA, ReAct, планирование и агентов.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
giga: GigaClient,
|
||
kb: KBService,
|
||
config: AppConfig,
|
||
default_prompts: Optional[Dict[str, str]] = None,
|
||
agent_coordinator: Optional[AgentCoordinator] = None,
|
||
tool_handler: Optional[ToolHandler] = None,
|
||
) -> None:
|
||
"""
|
||
Инициализация QueryProcessor.
|
||
|
||
Аргументы:
|
||
giga: клиент GigaChat.
|
||
kb: сервис базы знаний.
|
||
config: конфигурация приложения.
|
||
default_prompts: словарь промптов по умолчанию.
|
||
agent_coordinator: координатор агентов (если включена многоагентность).
|
||
tool_handler: обработчик инструментов для ReAct.
|
||
"""
|
||
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)
|
||
self.agent_coordinator: Optional[AgentCoordinator] = agent_coordinator
|
||
self.tool_handler: Optional[ToolHandler] = tool_handler
|
||
|
||
def _needs_external_data(self, query: str) -> bool:
|
||
"""
|
||
Определяет, требует ли запрос вызова внешних инструментов.
|
||
Используется для активации ReAct.
|
||
"""
|
||
keywords = ["найди", "поищи", "забронируй", "купи", "узнай курс", "какая погода",
|
||
"погода", "билеты", "отель", "рейс", "тариф", "курс"]
|
||
return any(kw in query.lower() for kw in keywords)
|
||
|
||
def _get_available_tools(self) -> List[Dict[str, Any]]:
|
||
"""
|
||
Возвращает список доступных инструментов для ReAct.
|
||
Использует реальные обработчики из ToolHandler, если он доступен.
|
||
"""
|
||
if self.tool_handler is None:
|
||
return []
|
||
return [
|
||
{
|
||
"name": "search",
|
||
"description": "Поиск информации в интернете по запросу",
|
||
"parameters": {"query": "string"},
|
||
"handler": self.tool_handler.search
|
||
},
|
||
{
|
||
"name": "get_weather",
|
||
"description": "Получить текущую погоду в городе",
|
||
"parameters": {"city": "string"},
|
||
"handler": self.tool_handler.get_weather
|
||
},
|
||
{
|
||
"name": "book_flight",
|
||
"description": "Забронировать авиабилет (заглушка)",
|
||
"parameters": {"origin": "string", "destination": "string", "date": "string"},
|
||
"handler": self.tool_handler.book_flight
|
||
},
|
||
{
|
||
"name": "get_exchange_rate",
|
||
"description": "Получить курс валют (заглушка)",
|
||
"parameters": {"from_currency": "string", "to_currency": "string"},
|
||
"handler": self.tool_handler.get_exchange_rate
|
||
}
|
||
]
|
||
|
||
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. Использование агентов (если включено) ----
|
||
if (getattr(self.config, 'enable_agents', False) and
|
||
self.agent_coordinator is not None):
|
||
logger.info("Запрос передан в AgentCoordinator")
|
||
try:
|
||
agent_result = await self.agent_coordinator.process(
|
||
query=query,
|
||
context="",
|
||
history=history,
|
||
tools=None
|
||
)
|
||
if agent_result:
|
||
logger.info("AgentCoordinator вернул ответ")
|
||
return {
|
||
"answer": agent_result,
|
||
"context": "",
|
||
"sources": [],
|
||
"confidence": None
|
||
}
|
||
except Exception as e:
|
||
logger.error(f"Ошибка при работе AgentCoordinator: {e}, переходим к обычному RAG")
|
||
|
||
# ---- 2. Расширение запроса ----
|
||
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
|
||
|
||
# ---- 3. Поиск релевантного контекста ----
|
||
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})")
|
||
|
||
# ---- 4. Обрезка контекста по токенам ----
|
||
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 = ""
|
||
|
||
# ---- 5. Переранжирование ----
|
||
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
|
||
)
|
||
|
||
# ---- 6. Планирование (если включено) ----
|
||
plan = None
|
||
extra_instructions: str = ""
|
||
if getattr(self.config, 'enable_planning', False) and intent in ("CALCULATION", "PROCEDURE"):
|
||
plan = await generate_plan(
|
||
query=query,
|
||
history=history,
|
||
config=self.config,
|
||
giga=self.giga,
|
||
prompts=self.default_prompts
|
||
)
|
||
if plan:
|
||
logger.info(f"Сгенерирован план из {len(plan)} шагов")
|
||
extra_instructions = f"Выполни следующие шаги по порядку: {json.dumps(plan, ensure_ascii=False)}"
|
||
|
||
# ---- 7. Формирование промта ----
|
||
if intent in ("CALCULATION", "PROCEDURE") and not extra_instructions:
|
||
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]}")
|
||
|
||
# ---- 8. Генерация ответа ----
|
||
answer: str = await self.giga.chat(
|
||
history=[],
|
||
query=prompt,
|
||
system_prompt=None,
|
||
file_id=None,
|
||
temperature=getattr(self.config, 'ai_temperature', 0.1)
|
||
)
|
||
|
||
# ---- 9. ReAct (если требуется внешний вызов) ----
|
||
if (getattr(self.config, 'enable_react', False) and
|
||
self._needs_external_data(query)):
|
||
logger.info("Запрос требует внешних данных, запускаем ReAct цикл")
|
||
try:
|
||
tools = self._get_available_tools()
|
||
react_answer = await react_loop(
|
||
query=query,
|
||
history=history,
|
||
tools=tools,
|
||
config=self.config,
|
||
giga=self.giga,
|
||
prompts=self.default_prompts,
|
||
max_iterations=getattr(self.config, 'react_max_iterations', 5)
|
||
)
|
||
if react_answer:
|
||
answer = react_answer
|
||
logger.info("ReAct цикл завершён, ответ получен")
|
||
except Exception as e:
|
||
logger.error(f"ReAct цикл не удался: {e}, используем обычный ответ")
|
||
|
||
# ---- 10. Самокритика ----
|
||
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 = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
|
||
|
||
# ---- 11. SOMA-анализ (если включен) ----
|
||
if getattr(self.config, 'enable_soma', False):
|
||
logger.debug("Запуск SOMA-оценки")
|
||
try:
|
||
evaluation = await soma_evaluate(
|
||
query=query,
|
||
context=context,
|
||
answer=answer,
|
||
config=self.config,
|
||
giga=self.giga,
|
||
prompts=self.default_prompts
|
||
)
|
||
if evaluation.get('verdict') == 'fail':
|
||
threshold = getattr(self.config, 'soma_threshold', 3.5)
|
||
if evaluation.get('overall_score', 0) < threshold:
|
||
logger.warning(f"SOMA оценка низкая ({evaluation['overall_score']}), перегенерация")
|
||
answer = await self.giga.chat(
|
||
history=[],
|
||
query=prompt,
|
||
system_prompt=None,
|
||
file_id=None,
|
||
temperature=0.3
|
||
)
|
||
evaluation2 = await soma_evaluate(
|
||
query=query,
|
||
context=context,
|
||
answer=answer,
|
||
config=self.config,
|
||
giga=self.giga,
|
||
prompts=self.default_prompts
|
||
)
|
||
if evaluation2.get('verdict') == 'fail':
|
||
answer = "⚠️ Извините, я не уверен в точности ответа. Проверьте данные."
|
||
else:
|
||
logger.info("SOMA оценка после перегенерации успешна")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка в SOMA-оценке: {e}, пропускаем")
|
||
|
||
# ---- 12. Извлечение источников ----
|
||
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
|
||
}
|
||
|
||
async def close(self) -> None:
|
||
"""Закрывает ресурсы (HTTP-сессию ToolHandler)."""
|
||
if self.tool_handler:
|
||
await self.tool_handler.close()
|
||
logger.info("ToolHandler закрыт") |