a3f52888db
- /rag/history_manager.py - /rag/query_processor.py - /rag/rag_orchestrator.py - /rag/rag_api.py - /rag/rag_server.py - /rag/__init__.py
473 lines
19 KiB
Python
473 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
"""
|
||
RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов.
|
||
Предоставляет API для выполнения RAG-запросов и индексации документов.
|
||
|
||
Этот сервер запускается как отдельный процесс и может обслуживать
|
||
несколько ботов и других систем одновременно.
|
||
|
||
ИСПОЛЬЗУЕТ: FastAPI для создания REST API и Uvicorn для запуска.
|
||
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).
|
||
|
||
ДОБАВЛЕНО:
|
||
- Эндпоинт /rag/vision для распознавания текста на изображениях (OCR)
|
||
- Эндпоинт /rag/transcribe для транскрибации аудио (SaluteSpeech)
|
||
"""
|
||
|
||
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 pydantic import BaseModel, Field
|
||
import uvicorn
|
||
|
||
# Импортируем наши модули (теперь из пакета rag)
|
||
from rag.utils.config_loader import load_config
|
||
from rag.utils.logger import setup_logging
|
||
from rag.services.postgres_service import PostgresService
|
||
from rag.services.qdrant_service import QdrantService
|
||
from rag.services.embedding_service import EmbeddingService
|
||
from rag.services.kb_service import KBService
|
||
from rag.services.giga_client import GigaClient
|
||
from rag.services.file_service import FileService
|
||
from rag.rag_orchestrator import RAGOrchestrator
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ============================================================
|
||
# МОДЕЛИ ДАННЫХ (Pydantic) ДЛЯ API
|
||
# ============================================================
|
||
|
||
class QueryRequest(BaseModel):
|
||
"""
|
||
Модель запроса к RAG-серверу на выполнение RAG-запроса.
|
||
ИСТОРИЯ НЕ ПЕРЕДАЁТСЯ – сервер получает её из БД.
|
||
"""
|
||
query: str = Field(..., description="Текст запроса пользователя")
|
||
user_jid: str = Field(..., description="JID пользователя (без ресурса)")
|
||
room_jid: Optional[str] = Field(None, description="JID комнаты (None для личного чата)")
|
||
prompts: Dict[str, str] = Field(
|
||
default_factory=dict,
|
||
description="Словарь промптов (system, synthesis, intent, expand, critique и т.д.)"
|
||
)
|
||
intent_override: Optional[str] = Field(None, description="Принудительное переопределение намерения")
|
||
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="Исходное имя файла")
|
||
file_text: str = Field(..., description="Извлечённый текст документа")
|
||
user_jid: str = Field(..., description="JID владельца документа")
|
||
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: Optional[AppConfig] = None
|
||
|
||
|
||
def get_orchestrator() -> RAGOrchestrator:
|
||
"""Возвращает экземпляр RAGOrchestrator (инициализирует при первом вызове)."""
|
||
global _orchestrator
|
||
if _orchestrator is None:
|
||
raise RuntimeError("RAGOrchestrator не инициализирован. Сервер не запущен корректно.")
|
||
return _orchestrator
|
||
|
||
|
||
def init_orchestrator(config: AppConfig) -> RAGOrchestrator:
|
||
"""Инициализирует все сервисы и создаёт RAGOrchestrator."""
|
||
global _orchestrator, _config
|
||
_config = config
|
||
logger.info("Инициализация RAG-сервера...")
|
||
|
||
# ---- 1. Инициализация сервисов ----
|
||
# PostgreSQL (история, документы)
|
||
db = PostgresService(
|
||
host=config.database.host,
|
||
port=config.database.port,
|
||
user=config.database.user,
|
||
password=config.db_password,
|
||
db_name=config.database.database
|
||
)
|
||
|
||
# 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_certs
|
||
)
|
||
|
||
# База знаний (индексация, поиск)
|
||
kb = KBService(
|
||
db=db,
|
||
qdrant=qdrant,
|
||
embedding=embedding,
|
||
collection_name=config.qdrant.collection,
|
||
chunk_size_tokens=config.chunking.chunk_size_tokens,
|
||
overlap_tokens=config.chunking.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.gigachat.model_generation,
|
||
temperature=config.gigachat.temperature,
|
||
timeout=config.gigachat.timeout,
|
||
verify_ssl=False
|
||
)
|
||
|
||
# Файловый сервис (извлечение текста)
|
||
files = FileService(config)
|
||
|
||
# ---- 2. Создание оркестратора ----
|
||
# Загружаем дефолтные промпты (из config.prompts_content)
|
||
default_prompts = config.prompts_content or {}
|
||
system_prompt = default_prompts.get('system')
|
||
if system_prompt:
|
||
default_prompts['system'] = system_prompt
|
||
|
||
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 ПРИЛОЖЕНИЕ
|
||
# ============================================================
|
||
|
||
app = FastAPI(
|
||
title="RAG-сервер",
|
||
description="Единый сервис для RAG-обработки запросов и индексации документов",
|
||
version="1.0.0"
|
||
)
|
||
|
||
# Разрешаем CORS для возможности запросов с других доменов (для тестов)
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
@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
|
||
# ============================================================
|
||
|
||
@app.post("/rag/query", response_model=QueryResponse)
|
||
async def rag_query(request: QueryRequest):
|
||
"""
|
||
Выполняет 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)
|
||
async def index_document(request: IndexRequest):
|
||
"""Индексирует документ в базу знаний."""
|
||
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")
|
||
async def health_check():
|
||
"""Проверка здоровья сервиса."""
|
||
if _orchestrator is None:
|
||
return {"status": "unhealthy", "message": "Orchestrator not initialized"}
|
||
return {"status": "healthy"}
|
||
|
||
|
||
# ============================================================
|
||
# НОВЫЕ ЭНДПОИНТЫ: OCR и транскрибация
|
||
# ============================================================
|
||
|
||
@app.post("/rag/vision", response_model=VisionResponse)
|
||
async def vision_endpoint(
|
||
file: UploadFile = FastAPIFile(...),
|
||
user_jid: str = None,
|
||
room_jid: str = None
|
||
):
|
||
"""
|
||
Распознаёт текст на изображении через 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)
|
||
async def transcribe_endpoint(
|
||
file: UploadFile = FastAPIFile(...),
|
||
user_jid: str = None,
|
||
room_jid: str = None
|
||
):
|
||
"""
|
||
Транскрибирует аудиофайл через 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()
|
||
|
||
# Загружаем конфигурацию с помощью новой функции load_config
|
||
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)
|
||
# Передаём путь к логу из конфига
|
||
log_file_path = config.log_file_path or Path("/tmp/rag_server.log")
|
||
setup_logging(config.name, log_file_path, log_level=log_level)
|
||
|
||
logger.info(f"Запуск RAG-сервера с профилем {args.profile_dir}")
|
||
logger.info(f"Хост: {args.host}, порт: {args.port}")
|
||
|
||
# Инициализируем оркестратор
|
||
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", # теперь модуль из пакета rag
|
||
host=args.host,
|
||
port=args.port,
|
||
log_level=args.log_level,
|
||
reload=False
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |