from typing import Literal

from pydantic import BaseModel


class EntityOutput(BaseModel):
    entity_id: int
    entity_text: str
    classification: Literal["positive", "negative", "neutral"]


class SampleOutput(BaseModel):
    id: int
    entities: list[EntityOutput]