# 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, UNSET_SENTINEL
from .groundingtoolcount import GroundingToolCount, GroundingToolCountTypedDict
from .modalitytokens import ModalityTokens, ModalityTokensTypedDict
from pydantic import model_serializer
from typing import List, Optional
from typing_extensions import NotRequired, TypedDict


class UsageTypedDict(TypedDict):
    r"""Statistics on the interaction request's token usage."""

    cached_tokens_by_modality: NotRequired[List[ModalityTokensTypedDict]]
    r"""A breakdown of cached token usage by modality."""
    grounding_tool_count: NotRequired[List[GroundingToolCountTypedDict]]
    r"""Grounding tool count."""
    input_tokens_by_modality: NotRequired[List[ModalityTokensTypedDict]]
    r"""A breakdown of input token usage by modality."""
    output_tokens_by_modality: NotRequired[List[ModalityTokensTypedDict]]
    r"""A breakdown of output token usage by modality."""
    tool_use_tokens_by_modality: NotRequired[List[ModalityTokensTypedDict]]
    r"""A breakdown of tool-use token usage by modality."""
    total_cached_tokens: NotRequired[int]
    r"""Number of tokens in the cached part of the prompt (the cached content)."""
    total_input_tokens: NotRequired[int]
    r"""Number of tokens in the prompt (context)."""
    total_output_tokens: NotRequired[int]
    r"""Total number of tokens across all the generated responses."""
    total_thought_tokens: NotRequired[int]
    r"""Number of tokens of thoughts for thinking models."""
    total_tokens: NotRequired[int]
    r"""Total token count for the interaction request (prompt + responses + other
    internal tokens).
    """
    total_tool_use_tokens: NotRequired[int]
    r"""Number of tokens present in tool-use prompt(s)."""


class Usage(BaseModel):
    r"""Statistics on the interaction request's token usage."""

    cached_tokens_by_modality: Optional[List[ModalityTokens]] = None
    r"""A breakdown of cached token usage by modality."""

    grounding_tool_count: Optional[List[GroundingToolCount]] = None
    r"""Grounding tool count."""

    input_tokens_by_modality: Optional[List[ModalityTokens]] = None
    r"""A breakdown of input token usage by modality."""

    output_tokens_by_modality: Optional[List[ModalityTokens]] = None
    r"""A breakdown of output token usage by modality."""

    tool_use_tokens_by_modality: Optional[List[ModalityTokens]] = None
    r"""A breakdown of tool-use token usage by modality."""

    total_cached_tokens: Optional[int] = None
    r"""Number of tokens in the cached part of the prompt (the cached content)."""

    total_input_tokens: Optional[int] = None
    r"""Number of tokens in the prompt (context)."""

    total_output_tokens: Optional[int] = None
    r"""Total number of tokens across all the generated responses."""

    total_thought_tokens: Optional[int] = None
    r"""Number of tokens of thoughts for thinking models."""

    total_tokens: Optional[int] = None
    r"""Total token count for the interaction request (prompt + responses + other
    internal tokens).
    """

    total_tool_use_tokens: Optional[int] = None
    r"""Number of tokens present in tool-use prompt(s)."""

    @model_serializer(mode="wrap")
    def serialize_model(self, handler):
        optional_fields = set(
            [
                "cached_tokens_by_modality",
                "grounding_tool_count",
                "input_tokens_by_modality",
                "output_tokens_by_modality",
                "tool_use_tokens_by_modality",
                "total_cached_tokens",
                "total_input_tokens",
                "total_output_tokens",
                "total_thought_tokens",
                "total_tokens",
                "total_tool_use_tokens",
            ]
        )
        serialized = handler(self)
        m = {}

        for n, f in type(self).model_fields.items():
            k = f.alias or n
            val = serialized.get(k, serialized.get(n))

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

        return m
