rag-service/app/schemas.py

"""
Pydantic схемы для API RAG сервиса.
"""

from typing import Any, Dict, Optional

from pydantic import BaseModel, Field


# =============================================================================
# Запросы
# =============================================================================

class ScoreRequest(BaseModel):
    """Запрос на расчет скора."""
    text: str = Field(..., min_length=1, description="Текст поста для оценки")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "text": "Это пример текста поста для оценки скоринга"
            }
        }
    }


class ExampleRequest(BaseModel):
    """Запрос на добавление примера."""
    text: str = Field(..., min_length=1, description="Текст примера")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "text": "Это пример опубликованного/отклоненного поста"
            }
        }
    }


# =============================================================================
# Ответы
# =============================================================================

class ScoreMetadata(BaseModel):
    """Метаданные результата скоринга."""
    positive_examples: int = Field(..., description="Количество положительных примеров")
    negative_examples: int = Field(..., description="Количество отрицательных примеров")
    model: str = Field(..., description="Название модели")
    timestamp: int = Field(..., description="Время расчета (unix timestamp)")


class ScoreResponse(BaseModel):
    """Ответ с результатом скоринга."""
    rag_score: float = Field(..., ge=0.0, le=1.0, description="Основной скор (neg/pos формула)")
    rag_confidence: float = Field(..., ge=0.0, le=1.0, description="Уверенность в оценке")
    rag_score_pos_only: float = Field(..., ge=0.0, le=1.0, description="Скор только по положительным примерам")
    meta: ScoreMetadata = Field(..., description="Метаданные")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "rag_score": 0.7523,
                "rag_confidence": 0.85,
                "rag_score_pos_only": 0.6891,
                "meta": {
                    "positive_examples": 500,
                    "negative_examples": 350,
                    "model": "DeepPavlov/rubert-base-cased",
                    "timestamp": 1706270000
                }
            }
        }
    }


class ExampleResponse(BaseModel):
    """Ответ на добавление примера."""
    success: bool = Field(..., description="Успешность добавления")
    message: str = Field(..., description="Сообщение о результате")
    positive_count: int = Field(..., description="Текущее количество положительных примеров")
    negative_count: int = Field(..., description="Текущее количество отрицательных примеров")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "success": True,
                "message": "Положительный пример добавлен",
                "positive_count": 501,
                "negative_count": 350
            }
        }
    }


class VectorStoreStats(BaseModel):
    """Статистика хранилища векторов."""
    positive_count: int = Field(..., description="Количество положительных примеров")
    negative_count: int = Field(..., description="Количество отрицательных примеров")
    total_count: int = Field(..., description="Общее количество примеров")
    vector_dim: int = Field(..., description="Размерность векторов")
    max_examples: int = Field(..., description="Максимальное количество примеров")
    storage_path: Optional[str] = Field(None, description="Путь к файлу хранилища")


class StatsResponse(BaseModel):
    """Ответ со статистикой сервиса."""
    model_name: str = Field(..., description="Название модели")
    model_loaded: bool = Field(..., description="Загружена ли модель")
    device: Optional[str] = Field(None, description="Устройство (cpu/cuda)")
    cache_dir: str = Field(..., description="Директория кеша модели")
    vector_store: VectorStoreStats = Field(..., description="Статистика хранилища векторов")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "model_name": "DeepPavlov/rubert-base-cased",
                "model_loaded": True,
                "device": "cpu",
                "cache_dir": "data/models",
                "vector_store": {
                    "positive_count": 500,
                    "negative_count": 350,
                    "total_count": 850,
                    "vector_dim": 768,
                    "max_examples": 10000,
                    "storage_path": "data/vectors/vectors.npz"
                }
            }
        }
    }


class WarmupResponse(BaseModel):
    """Ответ на прогрев модели."""
    success: bool = Field(..., description="Успешность загрузки")
    model_loaded: bool = Field(..., description="Загружена ли модель")
    message: str = Field(..., description="Сообщение о результате")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "success": True,
                "model_loaded": True,
                "message": "Модель успешно загружена"
            }
        }
    }


class ErrorResponse(BaseModel):
    """Ответ с ошибкой."""
    detail: str = Field(..., description="Описание ошибки")
    error_type: str = Field(..., description="Тип ошибки")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "detail": "Недостаточно примеров для расчета скора",
                "error_type": "InsufficientExamplesError"
            }
        }
    }


class HealthResponse(BaseModel):
    """Ответ проверки здоровья сервиса."""
    status: str = Field(..., description="Статус сервиса")
    model_loaded: bool = Field(..., description="Загружена ли модель")
    version: str = Field(..., description="Версия сервиса")
    
    model_config = {
        "json_schema_extra": {
            "example": {
                "status": "healthy",
                "model_loaded": True,
                "version": "0.1.0"
            }
        }
    }