# This file was auto-generated by Fern from our API Definition.

from __future__ import annotations
from ..core.unchecked_base_model import UncheckedBaseModel
import typing
from .tool_call import ToolCall
from ..core.pydantic_utilities import IS_PYDANTIC_V2
import pydantic
from .tool_result import ToolResult
import typing_extensions
from ..core.unchecked_base_model import UnionMetadata


class ChatbotMessage(UncheckedBaseModel):
    role: typing.Literal["CHATBOT"] = "CHATBOT"
    message: str
    tool_calls: typing.Optional[typing.List[ToolCall]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class SystemMessage(UncheckedBaseModel):
    role: typing.Literal["SYSTEM"] = "SYSTEM"
    message: str
    tool_calls: typing.Optional[typing.List[ToolCall]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class UserMessage(UncheckedBaseModel):
    role: typing.Literal["USER"] = "USER"
    message: str
    tool_calls: typing.Optional[typing.List[ToolCall]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


class ToolMessage(UncheckedBaseModel):
    role: typing.Literal["TOOL"] = "TOOL"
    tool_results: typing.Optional[typing.List[ToolResult]] = None

    if IS_PYDANTIC_V2:
        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
    else:

        class Config:
            smart_union = True
            extra = pydantic.Extra.allow


Message = typing_extensions.Annotated[
    typing.Union[ChatbotMessage, SystemMessage, UserMessage, ToolMessage], UnionMetadata(discriminant="role")
]