Files
fckbot/rag/rag_server.py
2026-06-30 14:31:50 +00:00

614 lines
25 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-запросов.
Предоставляет API для выполнения RAG-запросов и индексации документов.
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).
ДОБАВЛЕНО:
- Эндпоинт /rag/vision для распознавания текста на изображениях (OCR)
- Эндпоинт /rag/transcribe для транскрибации аудио (SaluteSpeech)
- Использование Pydantic-модели AppConfig для конфигурации.
- Аутентификация через API-ключ (X-API-Key)
- Документирование API через OpenAPI (Swagger UI и ReDoc)
"""
import logging
import os
import sys
import tempfile
import shutil
import asyncio
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
from fastapi.middleware.cors import CORSMiddleware
from fastapi.openapi.utils import get_openapi
from pydantic import BaseModel, Field
import uvicorn
# Импорты наших модулей
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
logger = logging.getLogger(__name__)
# ============================================================
# МОДЕЛИ ДАННЫХ (Pydantic) ДЛЯ API
# ============================================================
class QueryRequest(BaseModel):
"""
Модель запроса к RAG-серверу на выполнение RAG-запроса.
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ сервер получает её из БД.
"""
query: str = Field(..., description="Текст запроса пользователя", example="Какой OEE на линии?")
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="METRICS")
last_file_path: Optional[str] = Field(None, description="Путь к последнему загруженному файлу")
last_file_text: Optional[str] = Field(None, description="Текст последнего загруженного файла")
class QueryResponse(BaseModel):
"""Модель ответа RAG-сервера на RAG-запрос."""
answer: str = Field(..., description="Итоговый ответ")
intent: str = Field(default="GENERAL", description="Распознанное намерение")
context: Optional[str] = Field(None, description="Использованный контекст (для отладки)")
sources: List[str] = Field(default_factory=list, description="Список источников")
confidence: Optional[float] = Field(None, description="Оценка уверенности (0-1)")
error: Optional[str] = Field(None, description="Сообщение об ошибке, если есть")
class IndexRequest(BaseModel):
"""Модель запроса на индексацию документа."""
file_name: str = Field(..., description="Исходное имя файла", example="report.docx")
file_text: str = Field(..., description="Извлечённый текст документа")
user_jid: str = Field(..., description="JID владельца документа", example="user@domain.ru")
room_jid: Optional[str] = Field(None, description="JID комнаты (если документ комнатный)")
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 IndexResponse(BaseModel):
"""Модель ответа на индексацию."""
doc_id: Optional[int] = Field(None, description="Идентификатор документа в БД")
chunk_count: int = Field(0, description="Количество проиндексированных чанков")
error: Optional[str] = Field(None, description="Сообщение об ошибке, если есть")
class VisionResponse(BaseModel):
"""Модель ответа на распознавание изображения."""
text: str = Field(..., description="Распознанный текст")
status: str = Field(..., description="Статус: ok, no_text_found, error")
class TranscribeResponse(BaseModel):
"""Модель ответа на транскрибацию аудио."""
text: str = Field(..., description="Распознанный текст")
status: str = Field(..., description="Статус: ok, no_speech_detected, error")
# ============================================================
# ГЛОБАЛЬНЫЙ ОРКЕСТРАТОР (инициализируется один раз при старте)
# ============================================================
_orchestrator: Optional[RAGOrchestrator] = None
_config = None
def get_orchestrator() -> RAGOrchestrator:
"""Возвращает экземпляр RAGOrchestrator (инициализирует при первом вызове)."""
global _orchestrator
if _orchestrator is None:
raise RuntimeError("RAGOrchestrator не инициализирован. Сервер не запущен корректно.")
return _orchestrator
def init_orchestrator(config):
"""Инициализирует все сервисы и создаёт RAGOrchestrator."""
global _orchestrator, _config
_config = config
logger.info("Инициализация RAG-сервера...")
# ---- 1. Инициализация сервисов ----
# PostgreSQL (история, документы)
db = PostgresService(
host=config.db_host,
port=config.db_port,
user=config.db_user,
password=config.db_password,
db_name=config.db_name
)
# Qdrant (векторный поиск)
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
)
# Эмбеддинги (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
)
# База знаний (индексация, поиск)
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
)
# GigaChat клиент (генерация)
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)
# ---- 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}")
orchestrator = RAGOrchestrator(
db=db,
qdrant=qdrant,
embedding=embedding,
kb=kb,
giga=giga,
files=files,
config=config,
default_prompts=default_prompts
)
_orchestrator = orchestrator
logger.info("RAG-сервер успешно инициализирован")
return orchestrator
# ============================================================
# FASTAPI ПРИЛОЖЕНИЕ с документацией
# ============================================================
# Описание тегов для группировки эндпоинтов в документации
tags_metadata = [
{
"name": "RAG",
"description": "Основные RAG-операции: выполнение запроса и индексация документов",
},
{
"name": "Media",
"description": "Обработка медиафайлов: OCR для изображений и транскрибация аудио",
},
{
"name": "Health",
"description": "Проверка состояния сервиса",
}
]
app = FastAPI(
title="RAG-сервер Эфцекабот",
description="""
Единый сервис для RAG-обработки запросов и индексации документов.
## Аутентификация
Для всех защищённых эндпоинтов (кроме `/health`) требуется передача API-ключа
в заголовке `X-API-Key`.
## История диалога
Сервер самостоятельно хранит историю диалогов в PostgreSQL и использует её
для контекста при ответах. Клиентам не нужно передавать историю в запросах.
## Поддерживаемые форматы файлов
- Документы: DOCX, PDF, XLSX, PPTX, TXT, CSV, JSON
- Изображения: JPEG, PNG, BMP, TIFF (OCR через Tesseract)
- Аудио: OGG, WAV, MP3, AMR, M4A (транскрибация через SaluteSpeech)
- Архивы: ZIP, 7z (рекурсивная распаковка)
""",
version="3.0.0",
openapi_tags=tags_metadata,
contact={
"name": "Андрей Марков, Вениамин Кокорин",
"email": "support@fckbot.ru",
},
license_info={
"name": "ISC",
"url": "https://opensource.org/licenses/ISC",
},
)
# Добавляем CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Настраиваем OpenAPI для отображения схемы авторизации
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,
tags=app.openapi_tags,
contact=app.contact,
license_info=app.license_info,
)
# Добавляем схему безопасности для API-ключа
openapi_schema["components"]["securitySchemes"] = {
"APIKeyHeader": {
"type": "apiKey",
"in": "header",
"name": "X-API-Key",
"description": "API-ключ для доступа к защищённым эндпоинтам",
}
}
# Применяем security ко всем эндпоинтам кроме /health
for path in openapi_schema["paths"]:
if path != "/health":
for method in openapi_schema["paths"][path]:
openapi_schema["paths"][path][method]["security"] = [{"APIKeyHeader": []}]
app.openapi_schema = openapi_schema
return app.openapi_schema
app.openapi = custom_openapi
# ============================================================
# ОБРАБОТЧИКИ СОБЫТИЙ
# ============================================================
@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}")
@app.on_event("shutdown")
async def shutdown_event():
"""Закрываем соединение с БД при завершении сервера."""
if _orchestrator:
try:
await _orchestrator.db.close()
logger.info("Подключение к PostgreSQL закрыто")
except Exception as e:
logger.error(f"Ошибка закрытия PostgreSQL: {e}")
# ============================================================
# ЭНДПОИНТЫ API (все защищены аутентификацией, кроме /health)
# ============================================================
@app.post(
"/rag/query",
response_model=QueryResponse,
tags=["RAG"],
summary="Выполнить RAG-запрос",
description="Отправляет запрос пользователя в RAG-пайплайн. "
"Сервер автоматически подгружает историю диалога из БД.",
responses={
200: {"description": "Успешный ответ", "model": QueryResponse},
403: {"description": "Неверный или отсутствующий API-ключ"},
500: {"description": "Внутренняя ошибка сервера"},
}
)
async def rag_query(request: QueryRequest, _: str = Depends(verify_api_key)):
"""
Выполняет RAG-запрос.
ИСТОРИЯ БЕРЁТСЯ СЕРВЕРОМ ИЗ БД (не передаётся в запросе).
"""
orchestrator = get_orchestrator()
try:
result = await orchestrator.process_query(
query=request.query,
user_jid=request.user_jid,
room_jid=request.room_jid,
prompts=request.prompts,
intent_override=request.intent_override,
last_file_path=request.last_file_path,
last_file_text=request.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=500, detail=str(e))
@app.post(
"/rag/index",
response_model=IndexResponse,
tags=["RAG"],
summary="Проиндексировать документ",
description="Загружает текст документа в базу знаний. "
"Документ разбивается на чанки, для каждого генерируются эмбеддинги "
"и сохраняются в Qdrant. Метаданные хранятся в PostgreSQL.",
responses={
200: {"description": "Индексация выполнена", "model": IndexResponse},
403: {"description": "Неверный или отсутствующий API-ключ"},
500: {"description": "Внутренняя ошибка сервера"},
}
)
async def index_document(request: IndexRequest, _: str = Depends(verify_api_key)):
"""Индексирует документ в базу знаний."""
orchestrator = get_orchestrator()
try:
result = await orchestrator.index_document(
file_name=request.file_name,
file_text=request.file_text,
user_jid=request.user_jid,
room_jid=request.room_jid,
is_global=request.is_global,
title=request.title,
metadata=request.metadata,
file_hash=request.file_hash,
update_if_exists=request.update_if_exists
)
return IndexResponse(
doc_id=result.get('doc_id'),
chunk_count=result.get('chunk_count', 0),
error=result.get('error')
)
except Exception as e:
logger.error(f"Ошибка в /rag/index: {e}", exc_info=True)
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/health",
tags=["Health"],
summary="Проверка состояния сервиса",
description="Возвращает статус работоспособности RAG-сервера. "
"Не требует аутентификации.",
responses={
200: {"description": "Сервис работает", "content": {"application/json": {"example": {"status": "healthy"}}}},
503: {"description": "Сервис недоступен"},
}
)
async def health_check():
"""Проверка здоровья сервиса (без аутентификации)."""
if _orchestrator is None:
return {"status": "unhealthy", "message": "Orchestrator not initialized"}
return {"status": "healthy"}
@app.post(
"/rag/vision",
response_model=VisionResponse,
tags=["Media"],
summary="Распознать текст на изображении (OCR)",
description="Отправляет изображение на распознавание текста через Tesseract OCR. "
"Поддерживаются форматы: JPEG, PNG, BMP, TIFF.",
responses={
200: {"description": "Текст распознан", "model": VisionResponse},
400: {"description": "Неподдерживаемый формат файла"},
403: {"description": "Неверный или отсутствующий API-ключ"},
500: {"description": "Внутренняя ошибка сервера"},
}
)
async def vision_endpoint(
file: UploadFile = FastAPIFile(...),
_: str = Depends(verify_api_key)
):
"""
Распознаёт текст на изображении через OCR (Tesseract).
Возвращает JSON: {"text": str, "status": "ok" | "no_text_found" | "error"}
"""
orchestrator = 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="Поддерживаются только изображения (jpg, png, bmp, tiff)")
# Сохраняем во временный файл
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
tmp_path = tmp.name
shutil.copyfileobj(file.file, tmp)
try:
# Извлекаем текст через существующий FileService (синхронный метод)
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)
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
@app.post(
"/rag/transcribe",
response_model=TranscribeResponse,
tags=["Media"],
summary="Транскрибировать аудиофайл",
description="Отправляет аудиофайл на распознавание речи через SaluteSpeech. "
"Поддерживаются форматы: OGG, WAV, MP3, AMR, M4A.",
responses={
200: {"description": "Речь распознана", "model": TranscribeResponse},
400: {"description": "Неподдерживаемый формат файла"},
403: {"description": "Неверный или отсутствующий API-ключ"},
500: {"description": "Внутренняя ошибка сервера"},
}
)
async def transcribe_endpoint(
file: UploadFile = FastAPIFile(...),
_: str = Depends(verify_api_key)
):
"""
Транскрибирует аудиофайл через SaluteSpeech.
Возвращает JSON: {"text": str, "status": "ok" | "no_speech_detected" | "error"}
"""
orchestrator = 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="Поддерживаются только аудио (ogg, wav, mp3, amr, m4a)")
# Сохраняем во временный файл
with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
tmp_path = tmp.name
shutil.copyfileobj(file.file, tmp)
try:
# Транскрибируем через существующий FileService (асинхронный метод)
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)
raise HTTPException(status_code=500, detail=str(e))
finally:
if os.path.exists(tmp_path):
os.unlink(tmp_path)
# ============================================================
# ТОЧКА ВХОДА
# ============================================================
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="Уровень логирования"
)
args = parser.parse_args()
# Загружаем конфигурацию через Pydantic
try:
config = load_config(args.profile_dir)
except Exception as e:
print(f"Ошибка загрузки конфигурации: {e}")
sys.exit(1)
# Настраиваем логирование
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)
# Инициализируем оркестратор
try:
init_orchestrator(config)
except Exception as e:
logger.error(f"Ошибка инициализации оркестратора: {e}", exc_info=True)
sys.exit(1)
# Запускаем Uvicorn
uvicorn.run(
"rag.rag_server:app",
host=args.host,
port=args.port,
log_level=args.log_level,
reload=False
)
if __name__ == "__main__":
main()