app/models.py

"""Pydantic models for API requests and responses."""

from typing import Any

from pydantic import BaseModel, Field


class InferenceRequest(BaseModel):
    """Request model for inference endpoint."""

    inputs: str | list[str] = Field(..., description="Text input(s) for inference")
    parameters: dict[str, Any] = Field(
        default_factory=dict, description="Optional model parameters"
    )


class InferenceResponse(BaseModel):
    """Response model for inference endpoint."""

    predictions: list[Any] = Field(..., description="Model predictions")
    model_name: str = Field(..., description="Name of the model used")


class HealthResponse(BaseModel):
    """Response model for health check endpoint."""

    status: str = "ok"
    model_loaded: bool = False
    model_name: str | None = None