"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""

from __future__ import annotations

from typing import Union

from pydantic import Discriminator, Tag
from typing_extensions import Annotated, TypeAliasType

from friendli.utils import get_discriminator

from .assistantmessage import AssistantMessage, AssistantMessageTypedDict
from .systemmessage import SystemMessage, SystemMessageTypedDict
from .toolmessage import ToolMessage, ToolMessageTypedDict
from .usermessage import UserMessage, UserMessageTypedDict

MessageTypedDict = TypeAliasType(
    "MessageTypedDict",
    Union[
        SystemMessageTypedDict,
        UserMessageTypedDict,
        AssistantMessageTypedDict,
        ToolMessageTypedDict,
    ],
)


Message = Annotated[
    Union[
        Annotated[AssistantMessage, Tag("assistant")],
        Annotated[SystemMessage, Tag("system")],
        Annotated[ToolMessage, Tag("tool")],
        Annotated[UserMessage, Tag("user")],
    ],
    Discriminator(lambda m: get_discriminator(m, "role", "role")),
]
