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

rag/rag_server.py

@@ -1,16 +1,20 @@
+# ============================================================
+# rag/rag_server.py
+# Полная версия с интеграцией всех новых модулей
+# ============================================================
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 """
 RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов.
 ДОБАВЛЕНО:
-- Новые эндпоинты для команд бота (clear, list, delete_global, template_*, reset_history, room_*, stats, generate_document)
+- Новые эндпоинты для команд бота
 - Graceful shutdown (сигналы SIGTERM/SIGINT)
 - Расширенный /health с проверкой PostgreSQL и Qdrant
 - Ограничение CORS через переменную окружения
 - Метрики Prometheus (/metrics)
 - Структурированное JSON-логирование
 - Гарантированное удаление временных файлов
-- Интеграция ReindexWorker, SOMA, ReAct, планирования, агентов
+- Интеграция ReindexWorker, SOMA, ReAct, планирования, агентов, ToolHandler
 """
 
 import asyncio
@@ -54,6 +58,7 @@ from .functions.react import react_loop
 from .functions.plan_generation import generate_plan
 from .agents.coordinator import AgentCoordinator
 from .agents.roles import MethodologistAgent, BoxSolutionAgent, PersonalAssistantAgent
+from .functions.tools import ToolHandler
 
 logger = logging.getLogger(__name__)
 
@@ -113,6 +118,7 @@ _orchestrator: Optional[RAGOrchestrator] = None
 _config = None
 _reindex_worker: Optional[ReindexWorker] = None
 _agent_coordinator: Optional[AgentCoordinator] = None
+_tool_handler: Optional[ToolHandler] = None
 
 # ============================================================
 # ФУНКЦИИ ИНИЦИАЛИЗАЦИИ
@@ -124,7 +130,7 @@ def get_orchestrator() -> RAGOrchestrator:
     return _orchestrator
 
 def init_orchestrator(config):
-    global _orchestrator, _config, _reindex_worker, _agent_coordinator
+    global _orchestrator, _config, _reindex_worker, _agent_coordinator, _tool_handler
     _config = config
     logger.info("Инициализация RAG-сервера...")
 
@@ -172,6 +178,12 @@ def init_orchestrator(config):
     )
     files = FileService(config)
 
+    # ---- ToolHandler ----
+    _tool_handler = None
+    if getattr(config, 'enable_react', False):
+        _tool_handler = ToolHandler(config)
+        logger.info("ToolHandler инициализирован")
+
     # ---- Оркестратор ----
     default_prompts = config.prompts_content or {}
     orchestrator = RAGOrchestrator(
@@ -182,7 +194,8 @@ def init_orchestrator(config):
         giga=giga,
         files=files,
         config=config,
-        default_prompts=default_prompts
+        default_prompts=default_prompts,
+        tool_handler=_tool_handler
     )
     _orchestrator = orchestrator
 
@@ -305,6 +318,9 @@ async def shutdown_event():
     if _reindex_worker:
         await _reindex_worker.stop()
     if _orchestrator:
+        # Закрываем ToolHandler через QueryProcessor
+        if hasattr(_orchestrator, 'query_processor'):
+            await _orchestrator.query_processor.close()
         try:
             await _orchestrator.db.close()
         except Exception as e:
@@ -472,15 +488,13 @@ async def generate_document(request: Request, req: GenerateDocumentRequest, _: s
         if not template_info:
             raise HTTPException(status_code=404, detail="Шаблон не найден")
         template_path = template_info['file_path']
-        # Здесь должна быть логика генерации документа с использованием KB и Giga
-        # Для демонстрации возвращаем ссылку (в реальности нужно генерировать файл)
         result = await orch.generate_document_from_template(
             template_path=template_path,
             template_name=req.template_name,
             user_jid=req.user_jid,
             room_jid=req.room_jid,
         )
-        return {"document": result}  # result может быть URL или base64
+        return {"document": result}
     except HTTPException:
         raise
     except Exception as e:
@@ -579,8 +593,9 @@ async def shutdown():
     if _reindex_worker:
         await _reindex_worker.stop()
     if _orchestrator:
+        if hasattr(_orchestrator, 'query_processor'):
+            await _orchestrator.query_processor.close()
         await _orchestrator.db.close()
-    # здесь можно добавить остановку uvicorn
 
 # ============================================================
 # ТОЧКА ВХОДА