Редактировать query_processor.py

rag/query_processor.py

@@ -1,3 +1,7 @@
+# ============================================================
+# rag/query_processor.py
+# Полная версия с интеграцией всех новых функций
+# ============================================================
 # -*- coding: utf-8 -*-
 """
 Обработчик обычных RAG-запросов (FACT, PROCEDURE, COMPARISON, CALCULATION, GENERAL).
@@ -9,6 +13,7 @@
 - Интеграция ReAct (вызов внешних инструментов)
 - Интеграция планирования действий (generate_plan)
 - Интеграция многоагентности (AgentCoordinator)
+- Реальные обработчики инструментов (ToolHandler)
 """
 
 import logging
@@ -30,6 +35,7 @@ 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__)
 
@@ -47,6 +53,7 @@ class QueryProcessor:
         config: AppConfig,
         default_prompts: Optional[Dict[str, str]] = None,
         agent_coordinator: Optional[AgentCoordinator] = None,
+        tool_handler: Optional[ToolHandler] = None,
     ) -> None:
         """
         Инициализация QueryProcessor.
@@ -57,6 +64,7 @@ class QueryProcessor:
             config: конфигурация приложения.
             default_prompts: словарь промптов по умолчанию.
             agent_coordinator: координатор агентов (если включена многоагентность).
+            tool_handler: обработчик инструментов для ReAct.
         """
         self.giga: GigaClient = giga
         self.kb: KBService = kb
@@ -64,6 +72,7 @@ class QueryProcessor:
         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:
         """
@@ -77,22 +86,34 @@ class QueryProcessor:
     def _get_available_tools(self) -> List[Dict[str, Any]]:
         """
         Возвращает список доступных инструментов для ReAct.
-        В реальности инструменты определяются в конфиге и регистрируются в приложении.
-        Здесь возвращаем заглушку, которая должна быть переопределена в рантайме.
+        Использует реальные обработчики из ToolHandler, если он доступен.
         """
-        # В будущем можно загружать из конфига: getattr(self.config, 'react_tools', [])
+        if self.tool_handler is None:
+            return []
         return [
             {
                 "name": "search",
-                "description": "Поиск информации в интернете",
+                "description": "Поиск информации в интернете по запросу",
                 "parameters": {"query": "string"},
-                "handler": None  # будет передан извне (например, через регистрацию)
+                "handler": self.tool_handler.search
             },
             {
                 "name": "get_weather",
                 "description": "Получить текущую погоду в городе",
                 "parameters": {"city": "string"},
-                "handler": None
+                "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
             }
         ]
 
@@ -129,12 +150,11 @@ class QueryProcessor:
                 self.agent_coordinator is not None):
             logger.info("Запрос передан в AgentCoordinator")
             try:
-                # Пытаемся обработать запрос через агентов
                 agent_result = await self.agent_coordinator.process(
                     query=query,
-                    context="",  # агенты сами ищут контекст (или используют свой)
+                    context="",
                     history=history,
-                    tools=None  # можно передать общие инструменты
+                    tools=None
                 )
                 if agent_result:
                     logger.info("AgentCoordinator вернул ответ")
@@ -156,7 +176,7 @@ class QueryProcessor:
         )
         search_query: str = expanded if expanded and expanded != query else query
 
-        # ---- 3. Поиск релевантного контекста в базе знаний ----
+        # ---- 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)
@@ -177,7 +197,7 @@ class QueryProcessor:
                 else:
                     context = ""
 
-        # ---- 5. Переранжирование контекста ----
+        # ---- 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(
@@ -203,7 +223,7 @@ class QueryProcessor:
                 logger.info(f"Сгенерирован план из {len(plan)} шагов")
                 extra_instructions = f"Выполни следующие шаги по порядку: {json.dumps(plan, ensure_ascii=False)}"
 
-        # ---- 7. Формирование промта (с учётом возможного плана) ----
+        # ---- 7. Формирование промта ----
         if intent in ("CALCULATION", "PROCEDURE") and not extra_instructions:
             extra_instructions = (
                 "Пожалуйста, покажи пошаговое решение перед итоговым ответом. "
@@ -223,7 +243,7 @@ class QueryProcessor:
             system_prompt=system_prompt,
             extra_instructions=extra_instructions
         )
-        logger.debug(f"Сформированный промт (первые 500 символов): {prompt[:500]}")
+        logger.debug(f"Сформирован промт (первые 500 символов): {prompt[:500]}")
 
         # ---- 8. Генерация ответа ----
         answer: str = await self.giga.chat(
@@ -240,9 +260,6 @@ class QueryProcessor:
             logger.info("Запрос требует внешних данных, запускаем ReAct цикл")
             try:
                 tools = self._get_available_tools()
-                # В реальности нужно передать реальные обработчики инструментов.
-                # Здесь они могут быть зарегистрированы в конфиге.
-                # Для демонстрации мы пытаемся вызвать react_loop.
                 react_answer = await react_loop(
                     query=query,
                     history=history,
@@ -306,7 +323,6 @@ class QueryProcessor:
                     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,
@@ -314,7 +330,6 @@ class QueryProcessor:
                             file_id=None,
                             temperature=0.3
                         )
-                        # Вторая оценка (опционально)
                         evaluation2 = await soma_evaluate(
                             query=query,
                             context=context,
@@ -342,3 +357,9 @@ class QueryProcessor:
             "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 закрыт")