from typing import Literal

from pydantic import BaseModel


class PositionInput(BaseModel):
    position_text: str
    length: int
    offset: int


class EntityInput(BaseModel):
    entity_id: int
    entity_text: str
    entity_type: Literal["company", "location"]
    positions: list[PositionInput]


class SampleInput(BaseModel):
    id: int
    text: str
    entities: list[EntityInput]