from datetime import datetime
from typing import Annotated, Literal

from pydantic import BaseModel, Field


class ToolFunctionCall(BaseModel):
    name: str
    arguments: str


class ToolCall(BaseModel):
    id: str
    type: Literal["function"] = "function"
    function: ToolFunctionCall


class UserMessage(BaseModel):
    role: Literal["user"] = "user"
    content: str
    created_at: datetime


class AssistantMessage(BaseModel):
    role: Literal["assistant"] = "assistant"
    content: str | None = None
    tool_calls: list[ToolCall] = Field(default_factory=list)
    created_at: datetime


class ToolMessage(BaseModel):
    role: Literal["tool"] = "tool"
    content: str
    tool_call_id: str
    name: str
    is_error: bool = False
    created_at: datetime


Message = Annotated[
    UserMessage | AssistantMessage | ToolMessage,
    Field(discriminator="role"),
]