"""Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
# @generated-id: c007f5ee0325

from __future__ import annotations
from .audiochunk import AudioChunk, AudioChunkTypedDict
from .documenturlchunk import DocumentURLChunk, DocumentURLChunkTypedDict
from .filechunk import FileChunk, FileChunkTypedDict
from .imageurlchunk import ImageURLChunk, ImageURLChunkTypedDict
from .referencechunk import ReferenceChunk, ReferenceChunkTypedDict
from .textchunk import TextChunk, TextChunkTypedDict
from .thinkchunk import ThinkChunk, ThinkChunkTypedDict
from functools import partial
from mistralai.client.types import BaseModel
from mistralai.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,
        FileChunkTypedDict,
        AudioChunkTypedDict,
        DocumentURLChunkTypedDict,
        ThinkChunkTypedDict,
    ],
)


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,
    "document_url": DocumentURLChunk,
    "text": TextChunk,
    "reference": ReferenceChunk,
    "file": FileChunk,
    "thinking": ThinkChunk,
    "input_audio": AudioChunk,
}


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