# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pyformat: disable

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

from __future__ import annotations
from .. import BaseModel
from ...utils.unions import parse_open_union
from .argumentsdelta import ArgumentsDelta, ArgumentsDeltaTypedDict
from .audiodelta import AudioDelta, AudioDeltaTypedDict
from .codeexecutioncalldelta import (
    CodeExecutionCallDelta,
    CodeExecutionCallDeltaTypedDict,
)
from .codeexecutionresultdelta import (
    CodeExecutionResultDelta,
    CodeExecutionResultDeltaTypedDict,
)
from .documentdelta import DocumentDelta, DocumentDeltaTypedDict
from .filesearchcalldelta import FileSearchCallDelta, FileSearchCallDeltaTypedDict
from .filesearchresultdelta import FileSearchResultDelta, FileSearchResultDeltaTypedDict
from .functionresultdelta import FunctionResultDelta, FunctionResultDeltaTypedDict
from .googlemapscalldelta import GoogleMapsCallDelta, GoogleMapsCallDeltaTypedDict
from .googlemapsresultdelta import GoogleMapsResultDelta, GoogleMapsResultDeltaTypedDict
from .googlesearchcalldelta import GoogleSearchCallDelta, GoogleSearchCallDeltaTypedDict
from .googlesearchresultdelta import (
    GoogleSearchResultDelta,
    GoogleSearchResultDeltaTypedDict,
)
from .imagedelta import ImageDelta, ImageDeltaTypedDict
from .mcpservertoolcalldelta import (
    MCPServerToolCallDelta,
    MCPServerToolCallDeltaTypedDict,
)
from .mcpservertoolresultdelta import (
    MCPServerToolResultDelta,
    MCPServerToolResultDeltaTypedDict,
)
from .retrievalcalldelta import RetrievalCallDelta, RetrievalCallDeltaTypedDict
from .retrievalresultdelta import RetrievalResultDelta, RetrievalResultDeltaTypedDict
from .textannotationdelta import TextAnnotationDelta, TextAnnotationDeltaTypedDict
from .textdelta import TextDelta, TextDeltaTypedDict
from .thoughtsignaturedelta import ThoughtSignatureDelta, ThoughtSignatureDeltaTypedDict
from .thoughtsummarydelta import ThoughtSummaryDelta, ThoughtSummaryDeltaTypedDict
from .urlcontextcalldelta import URLContextCallDelta, URLContextCallDeltaTypedDict
from .urlcontextresultdelta import URLContextResultDelta, URLContextResultDeltaTypedDict
from .videodelta import VideoDelta, VideoDeltaTypedDict
from functools import partial
from pydantic import ConfigDict
from pydantic.functional_validators import BeforeValidator
from typing import Any, Literal, Union
from typing_extensions import Annotated, TypeAliasType


StepDeltaDataTypedDict = TypeAliasType(
    "StepDeltaDataTypedDict",
    Union[
        ThoughtSignatureDeltaTypedDict,
        TextDeltaTypedDict,
        ArgumentsDeltaTypedDict,
        FileSearchCallDeltaTypedDict,
        TextAnnotationDeltaTypedDict,
        ThoughtSummaryDeltaTypedDict,
        FileSearchResultDeltaTypedDict,
        GoogleMapsCallDeltaTypedDict,
        GoogleMapsResultDeltaTypedDict,
        CodeExecutionCallDeltaTypedDict,
        URLContextCallDeltaTypedDict,
        GoogleSearchCallDeltaTypedDict,
        RetrievalResultDeltaTypedDict,
        MCPServerToolCallDeltaTypedDict,
        RetrievalCallDeltaTypedDict,
        CodeExecutionResultDeltaTypedDict,
        URLContextResultDeltaTypedDict,
        GoogleSearchResultDeltaTypedDict,
        MCPServerToolResultDeltaTypedDict,
        DocumentDeltaTypedDict,
        VideoDeltaTypedDict,
        ImageDeltaTypedDict,
        FunctionResultDeltaTypedDict,
        AudioDeltaTypedDict,
    ],
)


class UnknownStepDeltaData(BaseModel):
    r"""A StepDeltaData 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)


_STEP_DELTA_DATA_VARIANTS: dict[str, Any] = {
    "text": TextDelta,
    "image": ImageDelta,
    "audio": AudioDelta,
    "document": DocumentDelta,
    "video": VideoDelta,
    "thought_summary": ThoughtSummaryDelta,
    "thought_signature": ThoughtSignatureDelta,
    "text_annotation_delta": TextAnnotationDelta,
    "arguments_delta": ArgumentsDelta,
    "code_execution_call": CodeExecutionCallDelta,
    "url_context_call": URLContextCallDelta,
    "google_search_call": GoogleSearchCallDelta,
    "mcp_server_tool_call": MCPServerToolCallDelta,
    "file_search_call": FileSearchCallDelta,
    "google_maps_call": GoogleMapsCallDelta,
    "retrieval_call": RetrievalCallDelta,
    "code_execution_result": CodeExecutionResultDelta,
    "url_context_result": URLContextResultDelta,
    "google_search_result": GoogleSearchResultDelta,
    "mcp_server_tool_result": MCPServerToolResultDelta,
    "file_search_result": FileSearchResultDelta,
    "google_maps_result": GoogleMapsResultDelta,
    "retrieval_result": RetrievalResultDelta,
    "function_result": FunctionResultDelta,
}


StepDeltaData = Annotated[
    Union[
        TextDelta,
        ImageDelta,
        AudioDelta,
        DocumentDelta,
        VideoDelta,
        ThoughtSummaryDelta,
        ThoughtSignatureDelta,
        TextAnnotationDelta,
        ArgumentsDelta,
        CodeExecutionCallDelta,
        URLContextCallDelta,
        GoogleSearchCallDelta,
        MCPServerToolCallDelta,
        FileSearchCallDelta,
        GoogleMapsCallDelta,
        RetrievalCallDelta,
        CodeExecutionResultDelta,
        URLContextResultDelta,
        GoogleSearchResultDelta,
        MCPServerToolResultDelta,
        FileSearchResultDelta,
        GoogleMapsResultDelta,
        RetrievalResultDelta,
        FunctionResultDelta,
        UnknownStepDeltaData,
    ],
    BeforeValidator(
        partial(
            parse_open_union,
            disc_key="type",
            variants=_STEP_DELTA_DATA_VARIANTS,
            unknown_cls=UnknownStepDeltaData,
            union_name="StepDeltaData",
            lenient=True,
        )
    ),
]
