models.py

"""Pydantic models for Nemo Seed Generator API."""

from enum import Enum
from typing import List, Optional, Any

from pydantic import BaseModel, Field


class HFInferenceModel(str, Enum):
    """Available HuggingFace Inference models."""

    QWEN_2_5_7B = "Qwen/Qwen2.5-7B-Instruct"


class HFInferenceProvider(str, Enum):
    """Available HuggingFace Inference providers."""

    TOGETHER = "together"
    INFERENCE_SERVER = "inference-server"


class SeedGenerateRequest(BaseModel):
    """Request model for generating seed Q&A pairs."""

    num_records: int = Field(
        default=1,
        ge=1,
        le=50,
        description="Number of Q&A pairs to generate (max 50)"
    )
    topic: str = Field(
        ...,
        min_length=1,
        max_length=200,
        description="Topic for the Q&A pairs (e.g., 'machine learning')"
    )
    difficulty: str = Field(
        ...,
        pattern=r"^(beginner|intermediate|advanced)$",
        description="Difficulty level: beginner, intermediate, or advanced"
    )
    model: HFInferenceModel = Field(
        default=HFInferenceModel.QWEN_2_5_7B,
        description="HuggingFace model to use"
    )
    provider: HFInferenceProvider = Field(
        default=HFInferenceProvider.TOGETHER,
        description="HF Inference API provider"
    )


class SeedRecord(BaseModel):
    """A single Q&A seed record."""

    topic: str = Field(..., description="The topic of the Q&A")
    difficulty: str = Field(..., description="Difficulty level")
    question: str = Field(..., description="The question")
    answer: str = Field(..., description="The answer")
    model: str = Field(..., description="Model used to generate")
    provider: str = Field(..., description="Provider used")


class SeedGenerateResponse(BaseModel):
    """Response model for seed generation."""

    success: bool = Field(..., description="Whether generation succeeded")
    data: Optional[List[SeedRecord]] = Field(
        default=None,
        description="Generated Q&A pairs if successful"
    )
    record_count: int = Field(default=0, description="Number of records generated")
    error: Optional[str] = Field(default=None, description="Error message if failed")


class HealthResponse(BaseModel):
    """Health check response."""

    status: str = Field(..., description="Health status")
    model: str = Field(..., description="Current model")
    provider: str = Field(..., description="Current provider")