Files
fckbot/rag/rag_server.py
Markov Andrey 41f32964fe Update 2 files
- /requirements.txt
- /rag/rag_server.py
2026-06-30 16:03:15 +00:00

807 lines
29 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).
ДОБАВЛЕНО:
- Rate limiting через slowapi (защита от перегрузки)
- Централизованное логирование всех запросов (middleware)
- Аутентификация через API-ключ (X-API-Key)
- Полная OpenAPI-документация с примерами и тегами
"""
import logging
import os
import sys
import tempfile
import shutil
import asyncio
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
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 .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="Какое тактовое время установлено на линии сборки?"
)
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 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 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 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 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="Сообщение об ошибке"
)
# ============================================================
# ГЛОБАЛЬНЫЙ ОРКЕСТРАТОР (инициализируется один раз при старте)
# ============================================================
_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 ПРИЛОЖЕНИЕ С КАСТОМНОЙ OPENAPI, RATE LIMITING И ЛОГИРОВАНИЕМ
# ============================================================
# 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 запросов в минуту
""",
version="3.0.0",
contact={
"name": "Команда Эфцекабот",
"email": "support@fckbot.ru",
},
license_info={
"name": "ISC License",
"url": "https://opensource.org/licenses/ISC",
},
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
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# 5. Централизованное логирование запросов (middleware)
@app.middleware("http")
async def log_requests(request: Request, call_next):
"""Логирует все входящие запросы с временем выполнения и IP-адресом."""
start_time = time.time()
response = await call_next(request)
duration = time.time() - start_time
client_host = request.client.host if request.client else "unknown"
logger.info(
f"{request.method} {request.url.path}{response.status_code} "
f"({duration:.3f}s) from {client_host}"
)
return response
# 6. Кастомизация 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,
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("Оркестратор не инициализирован.")
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 С ТЕГАМИ, ПОДРОБНЫМИ ОПИСАНИЯМИ И 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="Возвращает идентификатор документа и количество чанков"
)
@limiter.limit("5/minute")
async def index_document(
request: Request,
index_req: IndexRequest,
_: str = Depends(verify_api_key)
):
orchestrator = 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')
)
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)
)
@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="Текст, распознанный на изображении"
)
@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()
# Проверяем расширение
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)
@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="Текст, распознанный из аудио"
)
@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()
# Проверяем расширение
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)
@app.get(
"/health",
tags=["Health"],
summary="Проверка состояния сервиса",
description="""
Проверяет, запущен ли RAG-сервер и инициализирован ли оркестратор.
**Ответ:**
- `healthy` сервер работает
- `unhealthy` оркестратор не инициализирован
**Примечание:** этот эндпоинт **не требует** аутентификации и **не имеет** rate limiting.
""",
response_description="Статус сервиса"
)
async def health_check():
"""Проверка здоровья сервиса (без аутентификации)."""
if _orchestrator is None:
return {"status": "unhealthy", "message": "Orchestrator not initialized"}
return {"status": "healthy"}
# ============================================================
# ТОЧКА ВХОДА
# ============================================================
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()