Add new file
This commit is contained in:
Markov Andrey
379
core/rag_server.py
Normal file
379
core/rag_server.py
Normal file
@@ -0,0 +1,379 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
RAG-сервер – отдельный HTTP-сервис для обработки RAG-запросов.
|
||||
Предоставляет API для выполнения RAG-запросов и индексации документов.
|
||||
|
||||
Этот сервер запускается как отдельный процесс и может обслуживать
|
||||
несколько ботов и других систем одновременно.
|
||||
|
||||
ИСПОЛЬЗУЕТ: FastAPI для создания REST API и Uvicorn для запуска.
|
||||
ИСТОРИЯ ХРАНИТСЯ НА СЕРВЕРЕ (в PostgreSQL).
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional, Dict, List, Any
|
||||
|
||||
# Добавляем путь к core, если запускаем из корня проекта
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from fastapi import FastAPI, HTTPException, Depends
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
import uvicorn
|
||||
|
||||
# Импортируем наши модули
|
||||
from core.utils.config_loader import BotConfig
|
||||
from core.utils.logger import setup_logging
|
||||
from core.services.postgres_service import PostgresService
|
||||
from core.services.qdrant_service import QdrantService
|
||||
from core.services.embedding_service import EmbeddingService
|
||||
from core.services.kb_service import KBService
|
||||
from core.services.giga_client import GigaClient
|
||||
from core.services.file_service import FileService
|
||||
from core.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 для личного чата)")
|
||||
|
||||
# history: List[Dict[str, str]] = Field(...) # УБРАНО!
|
||||
|
||||
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="Сообщение об ошибке, если есть")
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ГЛОБАЛЬНЫЙ ОРКЕСТРАТОР (инициализируется один раз при старте)
|
||||
# ============================================================
|
||||
|
||||
_orchestrator: Optional[RAGOrchestrator] = None
|
||||
_config: Optional[BotConfig] = None
|
||||
|
||||
|
||||
def get_orchestrator() -> RAGOrchestrator:
|
||||
"""Возвращает экземпляр RAGOrchestrator (инициализирует при первом вызове)."""
|
||||
global _orchestrator
|
||||
if _orchestrator is None:
|
||||
raise RuntimeError("RAGOrchestrator не инициализирован. Сервер не запущен корректно.")
|
||||
return _orchestrator
|
||||
|
||||
|
||||
def init_orchestrator(config: BotConfig) -> RAGOrchestrator:
|
||||
"""Инициализирует все сервисы и создаёт 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 ПРИЛОЖЕНИЕ
|
||||
# ============================================================
|
||||
|
||||
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"}
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ТОЧКА ВХОДА
|
||||
# ============================================================
|
||||
|
||||
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()
|
||||
|
||||
# Загружаем конфигурацию
|
||||
try:
|
||||
config = BotConfig(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, 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(
|
||||
"core.rag_server:app",
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
log_level=args.log_level,
|
||||
reload=False
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user