llm-chat/src/llm_chat/models.py

from enum import StrEnum, auto

from pydantic import BaseModel, ConfigDict


class Role(StrEnum):
    """Role of a user in the chat."""

    ASSISTANT = auto()
    SYSTEM = auto()
    USER = auto()


class Message(BaseModel):
    """Message in the conversation."""

    role: Role
    content: str

    model_config: ConfigDict = ConfigDict(  # type: ignore[misc]
        frozen=True,
        use_enum_values=True,
    )


class Conversation(BaseModel):
    """Conversation in the chat."""

    messages: list[Message]

    model_config: ConfigDict = ConfigDict(  # type: ignore[misc]
        frozen=False,
        use_enum_values=True,
    )