diff --git a/rag/rag_server.py b/rag/rag_server.py index 7a498ca..b2d78cc 100644 --- a/rag/rag_server.py +++ b/rag/rag_server.py @@ -2,27 +2,29 @@ # -*- coding: utf-8 -*- """ RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов. -Предоставляет API для выполнения RAG-запросов и индексации документов. -ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL). - ДОБАВЛЕНО: -- Rate limiting через slowapi (защита от перегрузки) -- Централизованное логирование всех запросов (middleware) -- Аутентификация через API-ключ (X-API-Key) -- Полная OpenAPI-документация с примерами и тегами +- Новые эндпоинты для команд бота (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 asyncio +import signal import time from pathlib import Path from typing import Optional, Dict, List, Any -# Добавляем путь к rag, если запускаем из корня проекта sys.path.insert(0, str(Path(__file__).parent.parent)) from fastapi import FastAPI, HTTPException, Depends, UploadFile, File as FastAPIFile, Request, status @@ -30,13 +32,12 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.utils import get_openapi from pydantic import BaseModel, Field import uvicorn - -# Rate limiting 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 @@ -47,220 +48,87 @@ 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']) # ============================================================ -# МОДЕЛИ ДАННЫХ (Pydantic) ДЛЯ API +# МОДЕЛИ ДАННЫХ ДЛЯ НОВЫХ ЭНДПОИНТОВ # ============================================================ -class QueryRequest(BaseModel): - """ - Модель запроса к RAG-серверу на выполнение RAG-запроса. - ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД. - """ - query: str = Field( - ..., - description="Текст запроса пользователя", - example="Какое тактовое время установлено на линии сборки?" - ) - user_jid: str = Field( - ..., - description="JID пользователя (без ресурса)", - example="user@domain.ru" - ) - room_jid: Optional[str] = Field( - None, - description="JID комнаты (None для личного чата)", - example="room@conference.domain.ru" - ) - prompts: Dict[str, str] = Field( - default_factory=dict, - description="Словарь промптов (system, synthesis, intent, expand, critique и т.д.)" - ) - intent_override: Optional[str] = Field( - None, - description="Принудительное переопределение намерения", - example="FACT" - ) - last_file_path: Optional[str] = Field( - None, - description="Путь к последнему загруженному файлу на сервере" - ) - last_file_text: Optional[str] = Field( - None, - description="Текст последнего загруженного файла" - ) +class ClearRequest(BaseModel): + user_jid: str = Field(..., description="JID пользователя") + room_jid: Optional[str] = Field(None, description="JID комнаты (если очистка комнатной БЗ)") + is_global: bool = Field(False, description="Очищать глобальную БЗ (только для админов)") - class Config: - schema_extra = { - "example": { - "query": "Какое тактовое время установлено на линии сборки?", - "user_jid": "user@domain.ru", - "room_jid": None, - "prompts": {}, - "intent_override": None, - "last_file_path": None, - "last_file_text": None - } - } +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 QueryResponse(BaseModel): - """Модель ответа RAG-сервера на RAG-запрос.""" - answer: str = Field( - ..., - description="Итоговый ответ", - example="Тактовое время на линии сборки составляет 120 секунд. [Источник: Регламент_линии_2025.docx]" - ) - intent: str = Field( - default="GENERAL", - description="Распознанное намерение", - example="FACT" - ) - context: Optional[str] = Field( - None, - description="Использованный контекст (для отладки)" - ) - sources: List[str] = Field( - default_factory=list, - description="Список источников", - example=["Регламент_линии_2025.docx"] - ) - confidence: Optional[float] = Field( - None, - description="Оценка уверенности (0-1)", - example=0.95 - ) - error: Optional[str] = Field( - None, - 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 IndexRequest(BaseModel): - """Модель запроса на индексацию документа.""" - file_name: str = Field( - ..., - description="Исходное имя файла", - example="Регламент_отдела.docx" - ) - file_text: str = Field( - ..., - description="Извлечённый текст документа", - example="Полный текст документа..." - ) - user_jid: str = Field( - ..., - description="JID владельца документа", - example="user@domain.ru" - ) - room_jid: Optional[str] = Field( - None, - description="JID комнаты (если документ комнатный)", - example="room@conference.domain.ru" - ) - is_global: bool = Field( - False, - description="Глобальный ли документ" - ) - title: Optional[str] = Field( - None, - description="Отображаемое название" - ) - metadata: Optional[Dict[str, Any]] = Field( - None, - description="Дополнительные метаданные" - ) - file_hash: Optional[str] = Field( - None, - description="SHA-256 хеш содержимого" - ) - update_if_exists: bool = Field( - True, - description="Заменять ли существующий документ" - ) +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 IndexResponse(BaseModel): - """Модель ответа на индексацию.""" - doc_id: Optional[int] = Field( - None, - description="Идентификатор документа в БД", - example=123 - ) - chunk_count: int = Field( - 0, - description="Количество проиндексированных чанков", - example=15 - ) - error: Optional[str] = Field( - None, - description="Сообщение об ошибке, если есть" - ) - - -class VisionResponse(BaseModel): - """Модель ответа на распознавание изображения.""" - text: str = Field( - ..., - description="Распознанный текст", - example="Текст с изображения..." - ) - status: str = Field( - ..., - description="Статус: ok, no_text_found, error", - example="ok" - ) - error: Optional[str] = Field( - None, - description="Сообщение об ошибке" - ) - - -class TranscribeResponse(BaseModel): - """Модель ответа на транскрибацию аудио.""" - text: str = Field( - ..., - description="Распознанный текст", - example="Распознанная речь..." - ) - status: str = Field( - ..., - description="Статус: ok, no_speech_detected, error", - example="ok" - ) - error: Optional[str] = Field( - None, - description="Сообщение об ошибке" - ) +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: - """Возвращает экземпляр RAGOrchestrator (инициализирует при первом вызове).""" - global _orchestrator if _orchestrator is None: - raise RuntimeError("RAGOrchestrator не инициализирован. Сервер не запущен корректно.") + raise RuntimeError("RAGOrchestrator не инициализирован") return _orchestrator - def init_orchestrator(config): - """Инициализирует все сервисы и создаёт RAGOrchestrator.""" - global _orchestrator, _config + global _orchestrator, _config, _reindex_worker, _agent_coordinator _config = config logger.info("Инициализация RAG-сервера...") - # ---- 1. Инициализация сервисов ---- - # PostgreSQL (история, документы) + # ---- Сервисы ---- db = PostgresService( host=config.db_host, port=config.db_port, @@ -268,8 +136,6 @@ def init_orchestrator(config): password=config.db_password, db_name=config.db_name ) - - # Qdrant (векторный поиск) qdrant = QdrantService( host=config.qdrant_host, port=config.qdrant_port, @@ -279,17 +145,14 @@ def init_orchestrator(config): distance=config.qdrant_distance, prefer_grpc=False ) - - # Эмбеддинги (GigaChat) 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 + verify_ssl=config.embedding_verify_ssl, + redis_url=getattr(config, 'redis_url', None) ) - - # База знаний (индексация, поиск) kb = KBService( db=db, qdrant=qdrant, @@ -300,8 +163,6 @@ def init_orchestrator(config): approx_chunk_chars=config.chunking_approx_chunk_chars, approx_overlap_chars=config.chunking_approx_overlap_chars ) - - # GigaChat клиент (генерация) giga = GigaClient( api_key=config.gigachat_api_key, model=config.ai_model, @@ -309,21 +170,10 @@ def init_orchestrator(config): timeout=config.ai_timeout, verify_ssl=False ) - - # Файловый сервис (извлечение текста) files = FileService(config) - # ---- 2. Создание оркестратора ---- - # Загружаем дефолтные промпты (опционально) - default_prompts = {} - system_prompt_path = getattr(config, 'system_prompt_file', None) - if system_prompt_path and system_prompt_path.exists(): - try: - with open(system_prompt_path, 'r', encoding='utf-8') as f: - default_prompts['system'] = f.read() - except Exception as e: - logger.warning(f"Не удалось загрузить системный промпт: {e}") - + # ---- Оркестратор ---- + default_prompts = config.prompts_content or {} orchestrator = RAGOrchestrator( db=db, qdrant=qdrant, @@ -334,91 +184,75 @@ def init_orchestrator(config): 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 ПРИЛОЖЕНИЕ С КАСТОМНОЙ OPENAPI, RATE LIMITING И ЛОГИРОВАНИЕМ +# FASTAPI ПРИЛОЖЕНИЕ # ============================================================ -# 1. Настройка rate limiter limiter = Limiter(key_func=get_remote_address) - -# 2. Создание приложения app = FastAPI( title="Эфцекабот RAG-сервер", - description=""" - ## Единый RAG-сервис для интеллектуальной обработки запросов - - **Эфцекабот** – корпоративная платформа для гибридного RAG на базе GigaChat. - - ### Основные возможности - - **Гибридный поиск** (dense + sparse) с RRF-ранжированием - - **Самокритика и перегенерация** ответов - - **Иерархическое резюмирование** длинных документов и диалогов - - **Поддержка личных, комнатных и глобальных документов** - - **Распознавание текста на изображениях (OCR)** - - **Транскрибация аудио** через SaluteSpeech - - ### Аутентификация - Все эндпоинты (кроме `/health`) требуют передачи API-ключа в заголовке `X-API-Key`. - - ### Rate Limiting - - `/rag/query`: 10 запросов в минуту - - `/rag/index`: 5 запросов в минуту - - `/rag/vision`: 15 запросов в минуту - - `/rag/transcribe`: 10 запросов в минуту - """, + description="...", version="3.0.0", - contact={ - "name": "Команда Эфцекабот", - "email": "support@fckbot.ru", - }, - license_info={ - "name": "ISC License", - "url": "https://opensource.org/licenses/ISC", - }, + contact={"name": "Команда Эфцекабот", "email": "support@fckbot.ru"}, + license_info={"name": "ISC License"}, docs_url="/docs", redoc_url="/redoc", openapi_url="/openapi.json", ) - -# 3. Подключаем обработчик ошибок rate limiting app.state.limiter = limiter app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler) -# 4. CORS +# ---- CORS (ограничение через переменную окружения) ---- +cors_origins = os.getenv("CORS_ORIGINS", "*").split(",") app.add_middleware( CORSMiddleware, - allow_origins=["*"], + allow_origins=cors_origins if cors_origins != ["*"] else ["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) -# 5. Централизованное логирование запросов (middleware) +# ---- Middleware для логирования запросов ---- @app.middleware("http") async def log_requests(request: Request, call_next): - """Логирует все входящие запросы с временем выполнения и IP-адресом.""" - start_time = time.time() + start = time.time() response = await call_next(request) - duration = time.time() - start_time - client_host = request.client.host if request.client else "unknown" + duration = time.time() - start logger.info( - f"{request.method} {request.url.path} → {response.status_code} " - f"({duration:.3f}s) from {client_host}" + f"{request.method} {request.url.path} -> {response.status_code} " + f"({duration:.3f}s) from {request.client.host}" ) return response - -# 6. Кастомизация OpenAPI +# ---- Кастомная OpenAPI с security ---- def custom_openapi(): if app.openapi_schema: return app.openapi_schema - openapi_schema = get_openapi( title=app.title, version=app.version, @@ -427,373 +261,350 @@ def custom_openapi(): contact=app.contact, license_info=app.license_info, ) - - # Добавляем глобальные компоненты безопасности openapi_schema["components"]["securitySchemes"] = { "APIKeyHeader": { "type": "apiKey", "in": "header", "name": "X-API-Key", - "description": "API-ключ для доступа к сервису. Получите у администратора.", + "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("Оркестратор не инициализирован.") - else: - try: - await _orchestrator.db.connect() - logger.info("Подключение к PostgreSQL установлено") - except Exception as e: - logger.error(f"Ошибка подключения к PostgreSQL: {e}") + 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() - logger.info("Подключение к PostgreSQL закрыто") except Exception as e: logger.error(f"Ошибка закрытия PostgreSQL: {e}") - # ============================================================ -# ЭНДПОИНТЫ API С ТЕГАМИ, ПОДРОБНЫМИ ОПИСАНИЯМИ И RATE LIMITING +# НОВЫЕ ЭНДПОИНТЫ # ============================================================ -@app.post( - "/rag/query", - response_model=QueryResponse, - tags=["RAG"], - summary="Выполнить RAG-запрос", - description=""" - Отправляет текстовый запрос пользователя, выполняет полный RAG-пайплайн: - - 1. **Классификация намерений** – определяет тип запроса (FACT, PROCEDURE, METRICS и т.д.) - 2. **Расширение запроса** – добавляет синонимы и связанные термины - 3. **Гибридный поиск** (dense + sparse) в Qdrant - 4. **Переранжирование** через кросс-энкодер - 5. **Генерация ответа** через GigaChat - 6. **Самокритика** и перегенерация при необходимости - 7. **Сохранение истории** диалога в PostgreSQL - - **Требования**: заголовок `X-API-Key` с валидным ключом. - **История**: сервер автоматически получает историю диалога из БД. - **Rate limit**: 10 запросов в минуту. - """, - response_description="Ответ содержит сгенерированный текст, намерение, источники и уверенность" -) -@limiter.limit("10/minute") -async def rag_query( - request: Request, # обязательный параметр для slowapi - query_req: QueryRequest, - _: str = Depends(verify_api_key) -): - orchestrator = get_orchestrator() - - try: - result = await orchestrator.process_query( - query=query_req.query, - user_jid=query_req.user_jid, - room_jid=query_req.room_jid, - prompts=query_req.prompts, - intent_override=query_req.intent_override, - last_file_path=query_req.last_file_path, - last_file_text=query_req.last_file_text - ) - return QueryResponse( - answer=result.get('answer', ''), - intent=result.get('intent', 'GENERAL'), - context=result.get('context'), - sources=result.get('sources', []), - confidence=result.get('confidence'), - error=result.get('error') - ) - except Exception as e: - logger.error(f"Ошибка в /rag/query: {e}", exc_info=True) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=str(e) - ) - - -@app.post( - "/rag/index", - response_model=IndexResponse, - tags=["RAG"], - summary="Индексировать документ", - description=""" - Индексирует документ в базу знаний. - - **Процесс:** - 1. Документ разбивается на чанки (с учётом токенов) - 2. Для каждого чанка вычисляются dense- и sparse-эмбеддинги - 3. Чанки сохраняются в Qdrant с метаданными - 4. Метаданные сохраняются в PostgreSQL с флагом `indexed` - - **Поддерживаемые типы:** - - Личные (`is_global=False`, `room_jid=None`) - - Комнатные (`room_jid` указан) - - Глобальные (`is_global=True`, только для админов) - - **Обновление:** если документ с таким именем уже существует в комнате, - он будет заменён новой версией (по хешу). - - **Rate limit**: 5 запросов в минуту. - """, - response_description="Возвращает идентификатор документа и количество чанков" -) +@app.post("/rag/clear", tags=["Admin"], summary="Очистка базы знаний") @limiter.limit("5/minute") -async def index_document( - request: Request, - index_req: IndexRequest, - _: str = Depends(verify_api_key) -): - orchestrator = get_orchestrator() - +async def clear_kb(request: Request, req: ClearRequest, _: str = Depends(verify_api_key)): + orch = get_orchestrator() try: - result = await orchestrator.index_document( - file_name=index_req.file_name, - file_text=index_req.file_text, - user_jid=index_req.user_jid, - room_jid=index_req.room_jid, - is_global=index_req.is_global, - title=index_req.title, - metadata=index_req.metadata, - file_hash=index_req.file_hash, - update_if_exists=index_req.update_if_exists - ) - return IndexResponse( - doc_id=result.get('doc_id'), - chunk_count=result.get('chunk_count', 0), - error=result.get('error') - ) + 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.error(f"Ошибка в /rag/index: {e}", exc_info=True) - raise HTTPException( - status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, - detail=str(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", - response_model=VisionResponse, - tags=["Vision"], - summary="Распознать текст на изображении", - description=""" - Отправляет изображение на OCR-распознавание через Tesseract. - - **Поддерживаемые форматы:** - - JPEG, PNG, BMP, TIFF - - **Возвращает:** - - Распознанный текст - - Статус: `ok`, `no_text_found`, `error` - - **Требования:** заголовок `X-API-Key`. - **Rate limit**: 15 запросов в минуту. - """, - response_description="Текст, распознанный на изображении" -) +@app.post("/rag/vision", tags=["Vision"]) @limiter.limit("15/minute") -async def vision_endpoint( - request: Request, - file: UploadFile = FastAPIFile( - ..., - description="Файл изображения (JPEG, PNG, BMP, TIFF)" - ), - _: str = Depends(verify_api_key) -): - orchestrator = get_orchestrator() - - # Проверяем расширение +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=status.HTTP_400_BAD_REQUEST, - detail="Поддерживаются только изображения (jpg, png, bmp, tiff)" - ) - - # Сохраняем во временный файл - with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp: - tmp_path = tmp.name - shutil.copyfileobj(file.file, tmp) - - try: - loop = asyncio.get_running_loop() - text = await loop.run_in_executor( - None, - orchestrator.files.extract_text_from_image, - tmp_path - ) - if not text: - return VisionResponse(text="", status="no_text_found") - return VisionResponse(text=text, status="ok") - except Exception as e: - logger.error(f"Ошибка OCR: {e}", exc_info=True) - return VisionResponse(text="", status="error", error=str(e)) - finally: - if os.path.exists(tmp_path): - os.unlink(tmp_path) + 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", - response_model=TranscribeResponse, - tags=["Audio"], - summary="Транскрибировать аудио", - description=""" - Отправляет аудиофайл на транскрибацию через SaluteSpeech. - - **Поддерживаемые форматы:** - - OGG, WAV, MP3, AMR, M4A - - **Требования:** - - Заголовок `X-API-Key` - - Переменная окружения `SALUTE_SPEECH_AUTH` - - Установленный `ffmpeg` на сервере - - **Возвращает:** - - Распознанный текст - - Статус: `ok`, `no_speech_detected`, `error` - - **Rate limit**: 10 запросов в минуту. - """, - response_description="Текст, распознанный из аудио" -) +@app.post("/rag/transcribe", tags=["Audio"]) @limiter.limit("10/minute") -async def transcribe_endpoint( - request: Request, - file: UploadFile = FastAPIFile( - ..., - description="Аудиофайл (OGG, WAV, MP3, AMR, M4A)" - ), - _: str = Depends(verify_api_key) -): - orchestrator = get_orchestrator() - - # Проверяем расширение +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=status.HTTP_400_BAD_REQUEST, - detail="Поддерживаются только аудио (ogg, wav, mp3, amr, m4a)" - ) - - # Сохраняем во временный файл - with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp: - tmp_path = tmp.name - shutil.copyfileobj(file.file, tmp) - - try: - text = await orchestrator.files.transcribe_audio(tmp_path) - if not text: - return TranscribeResponse(text="", status="no_speech_detected") - return TranscribeResponse(text=text, status="ok") - except Exception as e: - logger.error(f"Ошибка транскрибации: {e}", exc_info=True) - return TranscribeResponse(text="", status="error", error=str(e)) - finally: - if os.path.exists(tmp_path): - os.unlink(tmp_path) + 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"], - summary="Проверка состояния сервиса", - description=""" - Проверяет, запущен ли RAG-сервер и инициализирован ли оркестратор. - - **Ответ:** - - `healthy` – сервер работает - - `unhealthy` – оркестратор не инициализирован - - **Примечание:** этот эндпоинт **не требует** аутентификации и **не имеет** rate limiting. - """, - response_description="Статус сервиса" -) +@app.get("/health", tags=["Health"]) async def health_check(): - """Проверка здоровья сервиса (без аутентификации).""" if _orchestrator is None: return {"status": "unhealthy", "message": "Orchestrator not initialized"} - return {"status": "healthy"} + 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(): - """Главная функция запуска RAG-сервера.""" import argparse - - parser = argparse.ArgumentParser(description="RAG-сервер") - parser.add_argument( - "--profile-dir", - required=True, - help="Путь к директории профиля бота (например, /usr/local/etc/bots/metabot)" - ) - parser.add_argument( - "--host", - default="0.0.0.0", - help="Хост для привязки сервера" - ) - parser.add_argument( - "--port", - type=int, - default=8080, - help="Порт для привязки сервера" - ) - parser.add_argument( - "--log-level", - default="info", - choices=["debug", "info", "warning", "error", "critical"], - help="Уровень логирования" - ) + 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() - # Загружаем конфигурацию через Pydantic - try: - config = load_config(args.profile_dir) - except Exception as e: - print(f"Ошибка загрузки конфигурации: {e}") - sys.exit(1) + 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) - # Настраиваем логирование - log_level = getattr(logging, args.log_level.upper(), logging.INFO) - setup_logging(config.name, config.log_file_path, log_level=log_level) - - logger.info(f"Запуск RAG-сервера с профилем {args.profile_dir}") - logger.info(f"Хост: {args.host}, порт: {args.port}") - - # Устанавливаем конфиг для аутентификации set_auth_config(config) + init_orchestrator(config) - # Инициализируем оркестратор - try: - init_orchestrator(config) - except Exception as e: - logger.error(f"Ошибка инициализации оркестратора: {e}", exc_info=True) - sys.exit(1) + signal.signal(signal.SIGTERM, handle_sigterm) + signal.signal(signal.SIGINT, handle_sigterm) - # Запускаем Uvicorn uvicorn.run( "rag.rag_server:app", host=args.host, @@ -802,6 +613,5 @@ def main(): reload=False ) - if __name__ == "__main__": main() \ No newline at end of file