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

from __future__ import annotations
from .imageurlchunk import ImageURLChunk, ImageURLChunkTypedDict
from .referencechunk import ReferenceChunk, ReferenceChunkTypedDict
from .textchunk import TextChunk, TextChunkTypedDict
from functools import partial
from mistralai.gcp.client.types import BaseModel
from mistralai.gcp.client.utils.unions import parse_open_union
from pydantic import ConfigDict
from pydantic.functional_validators import BeforeValidator
from typing import Any, Literal, Union
from typing_extensions import Annotated, TypeAliasType


ContentChunkTypedDict = TypeAliasType(
    "ContentChunkTypedDict",
    Union[TextChunkTypedDict, ImageURLChunkTypedDict, ReferenceChunkTypedDict],
)


class UnknownContentChunk(BaseModel):
    r"""A ContentChunk variant the SDK doesn't recognize. Preserves the raw payload."""

    type: Literal["UNKNOWN"] = "UNKNOWN"
    raw: Any
    is_unknown: Literal[True] = True

    model_config = ConfigDict(frozen=True)


_CONTENT_CHUNK_VARIANTS: dict[str, Any] = {
    "image_url": ImageURLChunk,
    "text": TextChunk,
    "reference": ReferenceChunk,
}


ContentChunk = Annotated[
    Union[ImageURLChunk, TextChunk, ReferenceChunk, UnknownContentChunk],
    BeforeValidator(
        partial(
            parse_open_union,
            disc_key="type",
            variants=_CONTENT_CHUNK_VARIANTS,
            unknown_cls=UnknownContentChunk,
            union_name="ContentChunk",
        )
    ),
]