from typing import Annotated, Literal

from typing_extensions import TypeAliasType

from pydantic import BaseModel, Field


class NoopSampler(BaseModel):
    """No-op sampler that returns the full dataset."""
    name: Literal["noop"]


class ChooseNSampler(BaseModel):
    """Sampler that chooses a random subset of annotation tasks."""
    name: Literal["choose_n"]
    n: int = Field(ge=0, description="The number of samples to take")


AnnotationSampler = TypeAliasType(
    "AnnotationSampler",
    Annotated[
        NoopSampler | ChooseNSampler,
        Field(discriminator="name")
    ]
)
