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

from __future__ import annotations
from .contentchunk import ContentChunk, ContentChunkTypedDict
from .toolcall import ToolCall, ToolCallTypedDict
from mistralai.gcp.client.types import (
    BaseModel,
    Nullable,
    OptionalNullable,
    UNSET,
    UNSET_SENTINEL,
)
from mistralai.gcp.client.utils import validate_const
import pydantic
from pydantic import model_serializer
from pydantic.functional_validators import AfterValidator
from typing import List, Literal, Optional, Union
from typing_extensions import Annotated, NotRequired, TypeAliasType, TypedDict


AssistantMessageContentTypedDict = TypeAliasType(
    "AssistantMessageContentTypedDict", Union[str, List[ContentChunkTypedDict]]
)


AssistantMessageContent = TypeAliasType(
    "AssistantMessageContent", Union[str, List[ContentChunk]]
)


class AssistantMessageTypedDict(TypedDict):
    role: Literal["assistant"]
    content: NotRequired[Nullable[AssistantMessageContentTypedDict]]
    tool_calls: NotRequired[Nullable[List[ToolCallTypedDict]]]
    prefix: NotRequired[bool]
    r"""Set this to `true` when adding an assistant message as prefix to condition the model response. The role of the prefix message is to force the model to start its answer by the content of the message."""


class AssistantMessage(BaseModel):
    role: Annotated[
        Annotated[
            Optional[Literal["assistant"]], AfterValidator(validate_const("assistant"))
        ],
        pydantic.Field(alias="role"),
    ] = "assistant"

    content: OptionalNullable[AssistantMessageContent] = UNSET

    tool_calls: OptionalNullable[List[ToolCall]] = UNSET

    prefix: Optional[bool] = False
    r"""Set this to `true` when adding an assistant message as prefix to condition the model response. The role of the prefix message is to force the model to start its answer by the content of the message."""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(["role", "content", "tool_calls", "prefix"])
        nullable_fields = set(["content", "tool_calls"])
        serialized = handler(self)
        m = {}

        for n, f in type(self).model_fields.items():
            k = f.alias or n
            val = serialized.get(k)
            is_nullable_and_explicitly_set = (
                k in nullable_fields
                and (self.__pydantic_fields_set__.intersection({n}))  # pylint: disable=no-member
            )

            if val != UNSET_SENTINEL:
                if (
                    val is not None
                    or k not in optional_fields
                    or is_nullable_and_explicitly_set
                ):
                    m[k] = val

        return m


try:
    AssistantMessage.model_rebuild()
except NameError:
    pass
