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

from __future__ import annotations
from .assistantmessage import AssistantMessage, AssistantMessageTypedDict
from .systemmessage import SystemMessage, SystemMessageTypedDict
from .toolmessage import ToolMessage, ToolMessageTypedDict
from .usermessage import UserMessage, UserMessageTypedDict
from mistralai.types import BaseModel
from mistralai.utils import get_discriminator
import pydantic
from pydantic import Discriminator, Tag
from typing import List, Optional, Union
from typing_extensions import Annotated, NotRequired, TypeAliasType, TypedDict


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


Two = 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")),
]


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


One = 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")),
]


ChatModerationRequestInputsTypedDict = TypeAliasType(
    "ChatModerationRequestInputsTypedDict",
    Union[List[OneTypedDict], List[List[TwoTypedDict]]],
)
r"""Chat to classify"""


ChatModerationRequestInputs = TypeAliasType(
    "ChatModerationRequestInputs", Union[List[One], List[List[Two]]]
)
r"""Chat to classify"""


class ChatModerationRequestTypedDict(TypedDict):
    model: str
    inputs: ChatModerationRequestInputsTypedDict
    r"""Chat to classify"""
    truncate_for_context_length: NotRequired[bool]


class ChatModerationRequest(BaseModel):
    model: str

    inputs: Annotated[ChatModerationRequestInputs, pydantic.Field(alias="input")]
    r"""Chat to classify"""

    truncate_for_context_length: Optional[bool] = False