Files
fckbot/rag/rag_server.py
2026-06-30 21:20:48 +00:00

617 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/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()