"""Model types and capability definitions for AI providers."""

from __future__ import annotations

from dataclasses import dataclass
from enum import StrEnum
from typing import Literal


class Capability(StrEnum):
    """Capabilities a model may support."""

    TEXT_GENERATION = "text generation"
    VISION = "vision"
    TEXT_TO_IMAGE = "text-to-image"
    REFERENCE_IMAGES = "reference images"
    CONTROL_IMAGES = "control images"


@dataclass(frozen=True)
class ModelInfo:
    """Describes a supported model and its capabilities."""

    name: str
    provider: str
    type: Literal["text", "image"]
    capabilities: list[Capability]