#!/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, планирования, агентов """ import asyncio import json import logging import os import sys import tempfile import shutil import signal import time from pathlib import Path from typing import Optional, Dict, List, Any sys.path.insert(0, str(Path(__file__).parent.parent)) from fastapi import FastAPI, HTTPException, Depends, UploadFile, File as FastAPIFile, Request, status from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.utils import get_openapi from pydantic import BaseModel, Field import uvicorn from slowapi import Limiter, _rate_limit_exceeded_handler from slowapi.util import get_remote_address from slowapi.errors import RateLimitExceeded from prometheus_client import Counter, Histogram, Gauge, generate_latest, REGISTRY from pythonjsonlogger import jsonlogger from .utils.config_loader import load_config from .utils.logger import setup_logging from .services.postgres_service import PostgresService from .services.qdrant_service import QdrantService from .services.embedding_service import EmbeddingService from .services.kb_service import KBService from .services.giga_client import GigaClient from .services.file_service import FileService from .rag_orchestrator import RAGOrchestrator from .auth import verify_api_key, set_auth_config from .workers.reindex_worker import ReindexWorker 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 .agents.roles import MethodologistAgent, BoxSolutionAgent, PersonalAssistantAgent logger = logging.getLogger(__name__) # ============================================================ # МЕТРИКИ PROMETHEUS # ============================================================ REQUEST_COUNT = Counter('rag_requests_total', 'Total RAG requests', ['endpoint', 'status']) REQUEST_LATENCY = Histogram('rag_request_latency_seconds', 'Request latency', ['endpoint']) ACTIVE_DOCUMENTS = Gauge('rag_active_documents', 'Number of indexed documents') ERROR_COUNT = Counter('rag_errors_total', 'Total errors', ['endpoint']) # ============================================================ # МОДЕЛИ ДАННЫХ ДЛЯ НОВЫХ ЭНДПОИНТОВ # ============================================================ class ClearRequest(BaseModel): user_jid: str = Field(..., description="JID пользователя") room_jid: Optional[str] = Field(None, description="JID комнаты (если очистка комнатной БЗ)") is_global: bool = Field(False, description="Очищать глобальную БЗ (только для админов)") class ListRequest(BaseModel): user_jid: str = Field(..., description="JID пользователя") room_jid: Optional[str] = Field(None, description="JID комнаты") class DeleteGlobalRequest(BaseModel): source_name: str = Field(..., description="Название документа для удаления") class TemplateSaveRequest(BaseModel): room_jid: str = Field(..., description="JID комнаты") name: str = Field(..., description="Имя шаблона") file_path: str = Field(..., description="Путь к файлу на сервере (должен быть загружен ранее)") file_hash: str = Field(..., description="SHA-256 хеш файла") class TemplateListRequest(BaseModel): room_jid: str = Field(..., description="JID комнаты") class TemplateDeleteRequest(BaseModel): room_jid: str = Field(..., description="JID комнаты") name: str = Field(..., description="Имя шаблона") class ResetHistoryRequest(BaseModel): user_jid: str = Field(..., description="JID пользователя") room_jid: Optional[str] = Field(None, description="JID комнаты (если сброс истории комнаты)") class AddRoomRequest(BaseModel): room_jid: str = Field(..., description="JID комнаты") class GenerateDocumentRequest(BaseModel): template_name: str = Field(..., description="Имя шаблона") user_jid: str = Field(..., description="JID пользователя") room_jid: str = Field(..., description="JID комнаты") # ============================================================ # ГЛОБАЛЬНЫЕ ПЕРЕМЕННЫЕ # ============================================================ _orchestrator: Optional[RAGOrchestrator] = None _config = None _reindex_worker: Optional[ReindexWorker] = None _agent_coordinator: Optional[AgentCoordinator] = None # ============================================================ # ФУНКЦИИ ИНИЦИАЛИЗАЦИИ # ============================================================ def get_orchestrator() -> RAGOrchestrator: if _orchestrator is None: raise RuntimeError("RAGOrchestrator не инициализирован") return _orchestrator def init_orchestrator(config): global _orchestrator, _config, _reindex_worker, _agent_coordinator _config = config logger.info("Инициализация RAG-сервера...") # ---- Сервисы ---- db = PostgresService( host=config.db_host, port=config.db_port, user=config.db_user, password=config.db_password, db_name=config.db_name ) qdrant = QdrantService( host=config.qdrant_host, port=config.qdrant_port, grpc_port=config.qdrant_grpc_port, collection_name=config.qdrant_collection, vector_size=config.qdrant_vector_size, distance=config.qdrant_distance, prefer_grpc=False ) embedding = EmbeddingService( api_key=config.gigachat_api_key, model=config.embedding_model, timeout=config.embedding_timeout, cache_size=config.embedding_cache_size, verify_ssl=config.embedding_verify_ssl, redis_url=getattr(config, 'redis_url', None) ) kb = KBService( db=db, qdrant=qdrant, embedding=embedding, collection_name=config.qdrant_collection, chunk_size_tokens=config.chunk_size_tokens, overlap_tokens=config.overlap_tokens, approx_chunk_chars=config.chunking_approx_chunk_chars, approx_overlap_chars=config.chunking_approx_overlap_chars ) giga = GigaClient( api_key=config.gigachat_api_key, model=config.ai_model, temperature=config.ai_temperature, timeout=config.ai_timeout, verify_ssl=False ) files = FileService(config) # ---- Оркестратор ---- default_prompts = config.prompts_content or {} orchestrator = RAGOrchestrator( db=db, qdrant=qdrant, embedding=embedding, kb=kb, giga=giga, files=files, config=config, default_prompts=default_prompts ) _orchestrator = orchestrator # ---- Инициализация ReindexWorker ---- if getattr(config, 'enable_reindex', True): _reindex_worker = ReindexWorker( db=db, kb=kb, config=config, interval_seconds=getattr(config, 'reindex_interval', 300), batch_size=getattr(config, 'reindex_batch_size', 10) ) # ---- Инициализация AgentCoordinator (если включено) ---- if getattr(config, 'enable_agents', False): _agent_coordinator = AgentCoordinator(config, giga) methodologist_prompt = default_prompts.get('agent_methodologist', 'Ты — методолог-аналитик...') box_prompt = default_prompts.get('agent_box_solution', 'Ты — эксперт по коробочным решениям...') personal_prompt = default_prompts.get('agent_personal_assistant', 'Ты — персональный ассистент...') _agent_coordinator.register_agent('methodologist', MethodologistAgent('methodologist', methodologist_prompt, giga, config)) _agent_coordinator.register_agent('box_solution', BoxSolutionAgent('box_solution', box_prompt, giga, config)) _agent_coordinator.register_agent('personal_assistant', PersonalAssistantAgent('personal_assistant', personal_prompt, giga, config)) logger.info("RAG-сервер успешно инициализирован") return orchestrator # ============================================================ # FASTAPI ПРИЛОЖЕНИЕ # ============================================================ limiter = Limiter(key_func=get_remote_address) app = FastAPI( title="Эфцекабот RAG-сервер", description="...", version="3.0.0", contact={"name": "Команда Эфцекабот", "email": "support@fckbot.ru"}, license_info={"name": "ISC License"}, docs_url="/docs", redoc_url="/redoc", openapi_url="/openapi.json", ) app.state.limiter = limiter app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) # ---- CORS (ограничение через переменную окружения) ---- cors_origins = os.getenv("CORS_ORIGINS", "*").split(",") app.add_middleware( CORSMiddleware, allow_origins=cors_origins if cors_origins != ["*"] else ["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ---- Middleware для логирования запросов ---- @app.middleware("http") async def log_requests(request: Request, call_next): start = time.time() response = await call_next(request) duration = time.time() - start logger.info( f"{request.method} {request.url.path} -> {response.status_code} " f"({duration:.3f}s) from {request.client.host}" ) return response # ---- Кастомная OpenAPI с security ---- def custom_openapi(): if app.openapi_schema: return app.openapi_schema openapi_schema = get_openapi( title=app.title, version=app.version, description=app.description, routes=app.routes, contact=app.contact, license_info=app.license_info, ) openapi_schema["components"]["securitySchemes"] = { "APIKeyHeader": { "type": "apiKey", "in": "header", "name": "X-API-Key", "description": "API-ключ для доступа к сервису", } } openapi_schema["security"] = [{"APIKeyHeader": []}] app.openapi_schema = openapi_schema return app.openapi_schema app.openapi = custom_openapi # ============================================================ # LIFECYCLE EVENTS # ============================================================ @app.on_event("startup") async def startup_event(): if _orchestrator is None: logger.error("Оркестратор не инициализирован") return try: await _orchestrator.db.connect() logger.info("Подключение к PostgreSQL установлено") except Exception as e: logger.error(f"Ошибка подключения к PostgreSQL: {e}") if _reindex_worker: await _reindex_worker.start() logger.info("ReindexWorker запущен") try: async with _orchestrator.db.pool.acquire() as conn: count = await conn.fetchval("SELECT COUNT(*) FROM documents WHERE indexed = TRUE") ACTIVE_DOCUMENTS.set(count) except Exception as e: logger.error(f"Ошибка получения количества документов: {e}") @app.on_event("shutdown") async def shutdown_event(): if _reindex_worker: await _reindex_worker.stop() if _orchestrator: try: await _orchestrator.db.close() except Exception as e: logger.error(f"Ошибка закрытия PostgreSQL: {e}") # ============================================================ # НОВЫЕ ЭНДПОИНТЫ # ============================================================ @app.post("/rag/clear", tags=["Admin"], summary="Очистка базы знаний") @limiter.limit("5/minute") async def clear_kb(request: Request, req: ClearRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: if req.is_global: await orch.kb.clear_global_kb() elif req.room_jid: await orch.kb.clear_room_kb(req.room_jid) else: await orch.kb.clear_user_kb(req.user_jid) return {"status": "ok"} except Exception as e: logger.exception(f"Ошибка очистки БЗ: {e}") ERROR_COUNT.labels(endpoint="clear").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/list", tags=["RAG"], summary="Список документов") @limiter.limit("10/minute") async def list_documents(request: Request, req: ListRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: sources = await orch.kb.get_sources_with_type(req.user_jid, req.room_jid) return {"documents": sources} except Exception as e: logger.exception(f"Ошибка получения списка документов: {e}") ERROR_COUNT.labels(endpoint="list").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/delete_global", tags=["Admin"], summary="Удалить глобальный документ") @limiter.limit("5/minute") async def delete_global(request: Request, req: DeleteGlobalRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: async with orch.db.pool.acquire() as conn: rows = await conn.fetch("SELECT id FROM documents WHERE is_global = TRUE AND source_name = $1", req.source_name) doc_ids = [row['id'] for row in rows] if not doc_ids: raise HTTPException(status_code=404, detail="Документ не найден") for doc_id in doc_ids: await asyncio.to_thread(orch.qdrant.delete_by_doc_id, doc_id) await orch.db.delete_documents_by_ids([doc_id]) return {"status": "ok"} except HTTPException: raise except Exception as e: logger.exception(f"Ошибка удаления глобального документа: {e}") ERROR_COUNT.labels(endpoint="delete_global").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/template/save", tags=["Templates"], summary="Сохранить шаблон") @limiter.limit("10/minute") async def template_save(request: Request, req: TemplateSaveRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: success = await orch.db.save_template(req.room_jid, req.name, req.file_path, req.file_hash) if success: return {"status": "ok"} else: raise HTTPException(status_code=500, detail="Ошибка сохранения шаблона") except Exception as e: logger.exception(f"Ошибка сохранения шаблона: {e}") ERROR_COUNT.labels(endpoint="template_save").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/template/list", tags=["Templates"], summary="Список шаблонов комнаты") @limiter.limit("10/minute") async def template_list(request: Request, req: TemplateListRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: names = await orch.db.list_templates(req.room_jid) return {"templates": names} except Exception as e: logger.exception(f"Ошибка получения списка шаблонов: {e}") ERROR_COUNT.labels(endpoint="template_list").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/template/delete", tags=["Templates"], summary="Удалить шаблон") @limiter.limit("10/minute") async def template_delete(request: Request, req: TemplateDeleteRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: deleted = await orch.db.delete_template(req.room_jid, req.name) if deleted: return {"status": "ok"} else: raise HTTPException(status_code=404, detail="Шаблон не найден") except HTTPException: raise except Exception as e: logger.exception(f"Ошибка удаления шаблона: {e}") ERROR_COUNT.labels(endpoint="template_delete").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/reset_history", tags=["History"], summary="Сброс истории диалога") @limiter.limit("5/minute") async def reset_history(request: Request, req: ResetHistoryRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: if req.room_jid: await orch.db.clear_room_history(req.room_jid) else: await orch.db.clear_user_history(req.user_jid) return {"status": "ok"} except Exception as e: logger.exception(f"Ошибка сброса истории: {e}") ERROR_COUNT.labels(endpoint="reset_history").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/room/add", tags=["Rooms"], summary="Добавить комнату") @limiter.limit("10/minute") async def add_room(request: Request, req: AddRoomRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: await orch.db.add_room(req.room_jid) return {"status": "ok"} except Exception as e: logger.exception(f"Ошибка добавления комнаты: {e}") ERROR_COUNT.labels(endpoint="add_room").inc() raise HTTPException(status_code=500, detail=str(e)) @app.get("/rag/room/list", tags=["Rooms"], summary="Список комнат") @limiter.limit("10/minute") async def get_rooms(request: Request, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: rooms = await orch.db.get_rooms() return {"rooms": rooms} except Exception as e: logger.exception(f"Ошибка получения списка комнат: {e}") ERROR_COUNT.labels(endpoint="get_rooms").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/stats", tags=["Admin"], summary="Статистика") @limiter.limit("5/minute") async def get_stats(request: Request, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: async with orch.db.pool.acquire() as conn: doc_count = await conn.fetchval("SELECT COUNT(*) FROM documents WHERE indexed = TRUE") return { "documents": doc_count, "uptime": time.time() - start_time if 'start_time' in globals() else 0, } except Exception as e: logger.exception(f"Ошибка получения статистики: {e}") ERROR_COUNT.labels(endpoint="stats").inc() raise HTTPException(status_code=500, detail=str(e)) @app.post("/rag/generate_document", tags=["Documents"], summary="Генерация документа по шаблону") @limiter.limit("5/minute") async def generate_document(request: Request, req: GenerateDocumentRequest, _: str = Depends(verify_api_key)): orch = get_orchestrator() try: template_info = await orch.db.get_template(req.room_jid, req.template_name) 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 except HTTPException: raise except Exception as e: logger.exception(f"Ошибка генерации документа: {e}") ERROR_COUNT.labels(endpoint="generate_document").inc() raise HTTPException(status_code=500, detail=str(e)) # ============================================================ # СУЩЕСТВУЮЩИЕ ЭНДПОИНТЫ (с изменениями для удаления временных файлов) # ============================================================ @app.post("/rag/vision", tags=["Vision"]) @limiter.limit("15/minute") async def vision_endpoint(request: Request, file: UploadFile = FastAPIFile(...), _: str = Depends(verify_api_key)): orch = get_orchestrator() ext = os.path.splitext(file.filename)[1].lower() if ext not in ('.jpg', '.jpeg', '.png', '.bmp', '.tiff'): raise HTTPException(status_code=400, detail="Неподдерживаемый формат изображения") with tempfile.TemporaryDirectory() as tmpdir: tmp_path = os.path.join(tmpdir, f"vision_{int(time.time())}{ext}") with open(tmp_path, 'wb') as f: shutil.copyfileobj(file.file, f) try: loop = asyncio.get_running_loop() text = await loop.run_in_executor(None, orch.files.extract_text_from_image, tmp_path) if not text: return {"text": "", "status": "no_text_found"} return {"text": text, "status": "ok"} except Exception as e: logger.exception(f"Ошибка OCR: {e}") return {"text": "", "status": "error", "error": str(e)} @app.post("/rag/transcribe", tags=["Audio"]) @limiter.limit("10/minute") async def transcribe_endpoint(request: Request, file: UploadFile = FastAPIFile(...), _: str = Depends(verify_api_key)): orch = get_orchestrator() ext = os.path.splitext(file.filename)[1].lower() if ext not in ('.ogg', '.wav', '.mp3', '.amr', '.m4a'): raise HTTPException(status_code=400, detail="Неподдерживаемый аудиоформат") with tempfile.TemporaryDirectory() as tmpdir: tmp_path = os.path.join(tmpdir, f"transcribe_{int(time.time())}{ext}") with open(tmp_path, 'wb') as f: shutil.copyfileobj(file.file, f) try: text = await orch.files.transcribe_audio(tmp_path) if not text: return {"text": "", "status": "no_speech_detected"} return {"text": text, "status": "ok"} except Exception as e: logger.exception(f"Ошибка транскрибации: {e}") return {"text": "", "status": "error", "error": str(e)} # ============================================================ # HEALTHCHECK # ============================================================ @app.get("/health", tags=["Health"]) async def health_check(): if _orchestrator is None: return {"status": "unhealthy", "message": "Orchestrator not initialized"} status = {"status": "healthy", "postgres": "unknown", "qdrant": "unknown"} try: async with _orchestrator.db.pool.acquire() as conn: await conn.execute("SELECT 1") status["postgres"] = "ok" except Exception as e: status["postgres"] = f"error: {e}" status["status"] = "unhealthy" try: _orchestrator.qdrant.client.get_collection(_orchestrator.qdrant.collection_name) status["qdrant"] = "ok" except Exception as e: status["qdrant"] = f"error: {e}" status["status"] = "unhealthy" return status # ============================================================ # PROMETHEUS METRICS # ============================================================ @app.get("/metrics", tags=["Monitoring"]) async def metrics(): from fastapi.responses import Response return Response(content=generate_latest(REGISTRY), media_type="text/plain") # ============================================================ # GRACEFUL SHUTDOWN # ============================================================ def handle_sigterm(sig, frame): logger.info(f"Получен сигнал {sig}, завершение...") asyncio.create_task(shutdown()) sys.exit(0) async def shutdown(): if _reindex_worker: await _reindex_worker.stop() if _orchestrator: await _orchestrator.db.close() # здесь можно добавить остановку uvicorn # ============================================================ # ТОЧКА ВХОДА # ============================================================ def main(): import argparse parser = argparse.ArgumentParser() parser.add_argument("--profile-dir", required=True, help="Путь к профилю") parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8080) parser.add_argument("--log-level", default="info") args = parser.parse_args() config = load_config(args.profile_dir) log_level = getattr(logging, args.log_level.upper()) setup_logging(config.name, config.log_file_path, log_level=log_level, json_format=True) set_auth_config(config) init_orchestrator(config) signal.signal(signal.SIGTERM, handle_sigterm) signal.signal(signal.SIGINT, handle_sigterm) uvicorn.run( "rag.rag_server:app", host=args.host, port=args.port, log_level=args.log_level, reload=False ) if __name__ == "__main__": main()