manager/models.py

"""Pydantic models for the token manager API."""
from pydantic import BaseModel, Field

VALID_SERVICES = ("tavily", "firecrawl", "exa")


# ── API Key models ──

class ApiKeyCreate(BaseModel):
    email: str
    password: str = ""
    api_key: str
    service: str | None = None
    created_at: str = ""


class ApiKeyImport(BaseModel):
    keys: list[ApiKeyCreate]


class ApiKeyResponse(BaseModel):
    id: int
    email: str
    password: str
    api_key: str
    service: str = "tavily"
    status: str
    created_at: str
    last_checked: str | None = None
    quota_remaining: int | None = None
    use_count: int | None = 0


class ApiKeyListResponse(BaseModel):
    total: int
    active: int
    keys: list[ApiKeyResponse]


class ServiceStats(BaseModel):
    total_keys: int = 0
    active_keys: int = 0
    inactive_keys: int = 0
    exhausted_keys: int = 0
    total_usage: int = 0
    total_quota_remaining: int | None = None


class StatsResponse(BaseModel):
    total_keys: int
    active_keys: int
    inactive_keys: int
    exhausted_keys: int
    total_usage: int = 0
    total_quota_remaining: int | None = None
    last_registration: str | None = None
    last_healthcheck: str | None = None
    by_service: dict[str, ServiceStats] = Field(default_factory=dict)


class HealthCheckResult(BaseModel):
    id: int
    api_key: str
    status: str
    message: str
    quota_remaining: int | None = None


class BatchIds(BaseModel):
    ids: list[int]


class BatchStatus(BaseModel):
    ids: list[int]
    status: str


# ── Access Token models ──

class AccessTokenCreate(BaseModel):
    name: str = ""
    token: str = ""
    quota_limit: int = 1000
    is_admin: bool = False
    expires_at: str | None = None


class AccessTokenUpdate(BaseModel):
    name: str | None = None
    quota_limit: int | None = None
    status: str | None = None
    expires_at: str | None = None


class AccessTokenResponse(BaseModel):
    id: int
    token: str
    name: str
    quota_limit: int
    quota_used: int
    is_admin: bool
    status: str
    created_at: str
    last_used: str | None = None
    expires_at: str | None = None


# ── Config models ──

class ConfigUpdate(BaseModel):
    configs: dict[str, str]