ag2 0.9.7__py3-none-any.whl → 0.9.8.post1__py3-none-any.whl

autogen/oai/gemini.py

@@ -51,14 +51,15 @@ import re
 import time
 import warnings
 from io import BytesIO
-from typing import Any, Literal, Optional, Type, Union
+from typing import Any, Literal
 
 import requests
 from pydantic import BaseModel, Field
+from typing_extensions import Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
 from ..json_utils import resolve_json_references
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import FormatterProtocol
 from .gemini_types import ToolConfig
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
@@ -68,6 +69,7 @@ with optional_import_block():
     import vertexai
     from PIL import Image
     from google.auth.credentials import Credentials
+    from google.genai import types
     from google.genai.types import (
         Content,
         FinishReason,
@@ -100,19 +102,34 @@ with optional_import_block():
 logger = logging.getLogger(__name__)
 
 
-@register_llm_config
+class GeminiEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["google"]
+
+    project_id: str | None
+    location: str | None
+    google_application_credentials: str | None
+    credentials: Any | str | None
+    stream: bool
+    safety_settings: list[dict[str, Any]] | dict[str, Any] | None
+    price: list[float] | None
+    tool_config: ToolConfig | None
+    proxy: str | None
+
+
 class GeminiLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["google"] = "google"
-    project_id: Optional[str] = None
-    location: Optional[str] = None
+    project_id: str | None = None
+    location: str | None = None
     # google_application_credentials points to the path of the JSON Keyfile
-    google_application_credentials: Optional[str] = None
+    google_application_credentials: str | None = None
     # credentials is a google.auth.credentials.Credentials object
-    credentials: Optional[Union[Any, str]] = None
+    credentials: Any | str | None = None
     stream: bool = False
-    safety_settings: Optional[Union[list[dict[str, Any]], dict[str, Any]]] = None
-    price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
-    tool_config: Optional[ToolConfig] = None
+    safety_settings: list[dict[str, Any]] | dict[str, Any] | None = None
+    price: list[float] | None = Field(default=None, min_length=2, max_length=2)
+    tool_config: ToolConfig | None = None
+    proxy: str | None = None
+    """A valid HTTP(S) proxy URL"""
 
     def create_client(self):
         raise NotImplementedError("GeminiLLMConfigEntry.create_client() is not implemented.")
@@ -134,7 +151,7 @@ class GeminiClient:
         "max_output_tokens": "max_output_tokens",
     }
 
-    def _initialize_vertexai(self, **params):
+    def _initialize_vertexai(self, **params: Unpack[GeminiEntryDict]):
         if "google_application_credentials" in params:
             # Path to JSON Keyfile
             os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = params["google_application_credentials"]
@@ -179,9 +196,10 @@ class GeminiClient:
             )
 
         self.api_version = kwargs.get("api_version")
+        self.proxy = kwargs.get("proxy")
 
         # Store the response format, if provided (for structured outputs)
-        self._response_format: Optional[type[BaseModel]] = None
+        self._response_format: type[BaseModel] | None = None
 
     def message_retrieval(self, response) -> list:
         """Retrieve and return a list of strings or a list of Choice.Message from the response.
@@ -237,8 +255,14 @@ class GeminiClient:
                 "See this [LLM configuration tutorial](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration/) for more details."
             )
 
-        params.get("api_type", "google")  # not used
-        http_options = {"api_version": self.api_version} if self.api_version else None
+        http_options = types.HttpOptions()
+        if proxy := params.get("proxy", self.proxy):
+            http_options.client_args = {"proxy": proxy}
+            http_options.async_client_args = {"proxy": proxy}
+
+        if self.api_version:
+            http_options.api_version = self.api_version
+
         messages = params.get("messages", [])
         stream = params.get("stream", False)
         n_response = params.get("n", 1)
@@ -670,9 +694,7 @@ class GeminiClient:
 
     @staticmethod
     def _convert_type_null_to_nullable(schema: Any) -> Any:
-        """
-        Recursively converts all occurrences of {"type": "null"} to {"nullable": True} in a schema.
-        """
+        """Recursively converts all occurrences of {"type": "null"} to {"nullable": True} in a schema."""
         if isinstance(schema, dict):
             # If schema matches {"type": "null"}, replace it
             if schema == {"type": "null"}:

autogen/oai/gemini_types.py

@@ -4,7 +4,7 @@
 
 import enum
 import warnings
-from typing import Any, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing import Any, Optional, TypeVar, Union, get_args, get_origin
 
 from pydantic import BaseModel as BaseModel
 from pydantic import ConfigDict, Field, alias_generators
@@ -15,7 +15,6 @@ def _remove_extra_fields(model: Any, response: dict[str, object]) -> None:
 
     Mutates the response in place.
     """
-
     key_values = list(response.items())
 
     for key, value in key_values:
@@ -64,7 +63,7 @@ class CommonBaseModel(BaseModel):
     )
 
     @classmethod
-    def _from_response(cls: Type[T], *, response: dict[str, object], kwargs: dict[str, object]) -> T:
+    def _from_response(cls: type[T], *, response: dict[str, object], kwargs: dict[str, object]) -> T:
         # To maintain forward compatibility, we need to remove extra fields from
         # the response.
         # We will provide another mechanism to allow users to access these fields.
@@ -91,7 +90,7 @@ class CaseInSensitiveEnum(str, enum.Enum):
                 try:
                     # Creating a enum instance based on the value
                     # We need to use super() to avoid infinite recursion.
-                    unknown_enum_val = super(CaseInSensitiveEnum, cls).__new__(cls, value)
+                    unknown_enum_val = super().__new__(cls, value)
                     unknown_enum_val._name_ = str(value)  # pylint: disable=protected-access
                     unknown_enum_val._value_ = value  # pylint: disable=protected-access
                     return unknown_enum_val
@@ -117,11 +116,11 @@ class LatLng(CommonBaseModel):
     WGS84 standard</a>. Values must be within normalized ranges.
     """
 
-    latitude: Optional[float] = Field(
+    latitude: float | None = Field(
         default=None,
         description="""The latitude in degrees. It must be in the range [-90.0, +90.0].""",
     )
-    longitude: Optional[float] = Field(
+    longitude: float | None = Field(
         default=None,
         description="""The longitude in degrees. It must be in the range [-180.0, +180.0]""",
     )
@@ -130,8 +129,8 @@ class LatLng(CommonBaseModel):
 class FunctionCallingConfig(CommonBaseModel):
     """Function calling config."""
 
-    mode: Optional[FunctionCallingConfigMode] = Field(default=None, description="""Optional. Function calling mode.""")
-    allowed_function_names: Optional[list[str]] = Field(
+    mode: FunctionCallingConfigMode | None = Field(default=None, description="""Optional. Function calling mode.""")
+    allowed_function_names: list[str] | None = Field(
         default=None,
         description="""Optional. Function names to call. Only set when the Mode is ANY. Function names should match [FunctionDeclaration.name]. With mode set to ANY, model will predict a function call from the set of function names provided.""",
     )
@@ -140,8 +139,8 @@ class FunctionCallingConfig(CommonBaseModel):
 class RetrievalConfig(CommonBaseModel):
     """Retrieval config."""
 
-    lat_lng: Optional[LatLng] = Field(default=None, description="""Optional. The location of the user.""")
-    language_code: Optional[str] = Field(default=None, description="""The language code of the user.""")
+    lat_lng: LatLng | None = Field(default=None, description="""Optional. The location of the user.""")
+    language_code: str | None = Field(default=None, description="""The language code of the user.""")
 
 
 class ToolConfig(CommonBaseModel):
@@ -150,7 +149,7 @@ class ToolConfig(CommonBaseModel):
     This config is shared for all tools provided in the request.
     """
 
-    function_calling_config: Optional[FunctionCallingConfig] = Field(
+    function_calling_config: FunctionCallingConfig | None = Field(
         default=None, description="""Optional. Function calling config."""
     )
-    retrieval_config: Optional[RetrievalConfig] = Field(default=None, description="""Optional. Retrieval config.""")
+    retrieval_config: RetrievalConfig | None = Field(default=None, description="""Optional. Retrieval config.""")

autogen/oai/groq.py

@@ -27,12 +27,13 @@ import copy
 import os
 import time
 import warnings
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 from pydantic import Field
+from typing_extensions import Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import should_hide_tools, validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -48,18 +49,26 @@ GROQ_PRICING_1K = {
 }
 
 
-@register_llm_config
+class GroqEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["groq"]
+
+    frequency_penalty: float
+    presence_penalty: float
+    seed: int
+    stream: bool
+    hide_tools: Literal["if_all_run", "if_any_run", "never"]
+    tool_choice: Literal["none", "auto", "required"] | None
+
+
 class GroqLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["groq"] = "groq"
+
     frequency_penalty: float = Field(default=None, ge=-2, le=2)
-    max_tokens: int = Field(default=None, ge=0)
     presence_penalty: float = Field(default=None, ge=-2, le=2)
     seed: int = Field(default=None)
     stream: bool = Field(default=False)
-    temperature: float = Field(default=1, ge=0, le=2)
-    top_p: float = Field(default=None)
     hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
-    tool_choice: Optional[Literal["none", "auto", "required"]] = None
+    tool_choice: Literal["none", "auto", "required"] | None = None
 
     def create_client(self):
         raise NotImplementedError("GroqLLMConfigEntry.create_client is not implemented.")
@@ -68,7 +77,7 @@ class GroqLLMConfigEntry(LLMConfigEntry):
 class GroqClient:
     """Client for Groq's API."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Unpack[GroqEntryDict]):
         """Requires api_key or environment variable to be set
 
         Args:
@@ -126,14 +135,17 @@ class GroqClient:
         groq_params["frequency_penalty"] = validate_parameter(
             params, "frequency_penalty", (int, float), True, None, (-2, 2), None
         )
+
         groq_params["max_tokens"] = validate_parameter(params, "max_tokens", int, True, None, (0, None), None)
+        groq_params["temperature"] = validate_parameter(params, "temperature", (int, float), True, 1, (0, 2), None)
+        groq_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
+
         groq_params["presence_penalty"] = validate_parameter(
             params, "presence_penalty", (int, float), True, None, (-2, 2), None
         )
         groq_params["seed"] = validate_parameter(params, "seed", int, True, None, None, None)
         groq_params["stream"] = validate_parameter(params, "stream", bool, True, False, None, None)
-        groq_params["temperature"] = validate_parameter(params, "temperature", (int, float), True, 1, (0, 2), None)
-        groq_params["top_p"] = validate_parameter(params, "top_p", (int, float), True, None, None, None)
+
         if "tool_choice" in params:
             groq_params["tool_choice"] = validate_parameter(
                 params, "tool_choice", str, True, None, None, ["none", "auto", "required"]

autogen/oai/mistral.py

@@ -29,12 +29,12 @@ import json
 import os
 import time
 import warnings
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal
 
-from pydantic import Field
+from typing_extensions import Unpack
 
 from ..import_utils import optional_import_block, require_optional_import
-from ..llm_config import LLMConfigEntry, register_llm_config
+from ..llm_config.entry import LLMConfigEntry, LLMConfigEntryDict
 from .client_utils import should_hide_tools, validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -53,17 +53,23 @@ with optional_import_block():
     )
 
 
-@register_llm_config
+class MistralEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["mistral"]
+
+    safe_prompt: bool
+    random_seed: int | None
+    stream: bool
+    hide_tools: Literal["if_all_run", "if_any_run", "never"]
+    tool_choice: Literal["none", "auto", "any"] | None
+
+
 class MistralLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["mistral"] = "mistral"
-    temperature: float = Field(default=0.7)
-    top_p: Optional[float] = None
-    max_tokens: Optional[int] = Field(default=None, ge=0)
     safe_prompt: bool = False
-    random_seed: Optional[int] = None
+    random_seed: int | None = None
     stream: bool = False
     hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
-    tool_choice: Optional[Literal["none", "auto", "any"]] = None
+    tool_choice: Literal["none", "auto", "any"] | None = None
 
     def create_client(self):
         raise NotImplementedError("MistralLLMConfigEntry.create_client is not implemented.")
@@ -73,7 +79,7 @@ class MistralLLMConfigEntry(LLMConfigEntry):
 class MistralAIClient:
     """Client for Mistral.AI's API."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Unpack[MistralEntryDict]):
         """Requires api_key or environment variable to be set
 
         Args:
@@ -93,7 +99,7 @@ class MistralAIClient:
 
         self._client = Mistral(api_key=self.api_key)
 
-    def message_retrieval(self, response: ChatCompletion) -> Union[list[str], list[ChatCompletionMessage]]:
+    def message_retrieval(self, response: ChatCompletion) -> list[str] | list[ChatCompletionMessage]:
         """Retrieve the messages from the response."""
         return [choice.message for choice in response.choices]
 

autogen/oai/oai_models/__init__.py

@@ -5,7 +5,19 @@
 from .chat_completion import ChatCompletionExtended as ChatCompletion
 from .chat_completion import Choice
 from .chat_completion_message import ChatCompletionMessage
-from .chat_completion_message_tool_call import ChatCompletionMessageToolCall
+from .chat_completion_message_tool_call import (
+    ChatCompletionMessageCustomToolCall,
+    ChatCompletionMessageFunctionToolCall,
+    ChatCompletionMessageToolCall,
+)
 from .completion_usage import CompletionUsage
 
-__all__ = ["ChatCompletion", "ChatCompletionMessage", "ChatCompletionMessageToolCall", "Choice", "CompletionUsage"]
+__all__ = [
+    "ChatCompletion",
+    "ChatCompletionMessage",
+    "ChatCompletionMessageCustomToolCall",
+    "ChatCompletionMessageFunctionToolCall",
+    "ChatCompletionMessageToolCall",
+    "Choice",
+    "CompletionUsage",
+]

autogen/oai/oai_models/_models.py

@@ -4,13 +4,13 @@
 
 # Taken over from https://github.com/openai/openai-python/blob/main/src/openai/_models.py
 
+
 import pydantic
 import pydantic.generics
 from pydantic import ConfigDict
-from typing_extensions import ClassVar
 
 __all__ = ["BaseModel"]
 
 
 class BaseModel(pydantic.BaseModel):
-    model_config: ClassVar[ConfigDict] = ConfigDict(extra="allow")
+    model_config = ConfigDict(extra="allow")

autogen/oai/oai_models/chat_completion.py

@@ -6,9 +6,8 @@
 
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import Any, Callable, List, Optional
-
-from typing_extensions import Literal
+from collections.abc import Callable
+from typing import Any, Literal
 
 from ._models import BaseModel
 from .chat_completion_message import ChatCompletionMessage
@@ -19,10 +18,10 @@ __all__ = ["ChatCompletion", "Choice", "ChoiceLogprobs"]
 
 
 class ChoiceLogprobs(BaseModel):
-    content: Optional[List[ChatCompletionTokenLogprob]] = None
+    content: list[ChatCompletionTokenLogprob] | None = None
     """A list of message content tokens with log probability information."""
 
-    refusal: Optional[List[ChatCompletionTokenLogprob]] = None
+    refusal: list[ChatCompletionTokenLogprob] | None = None
     """A list of message refusal tokens with log probability information."""
 
 
@@ -40,7 +39,7 @@ class Choice(BaseModel):
     index: int
     """The index of the choice in the list of choices."""
 
-    logprobs: Optional[ChoiceLogprobs] = None
+    logprobs: ChoiceLogprobs | None = None
     """Log probability information for the choice."""
 
     message: ChatCompletionMessage
@@ -51,7 +50,7 @@ class ChatCompletion(BaseModel):
     id: str
     """A unique identifier for the chat completion."""
 
-    choices: List[Choice]
+    choices: list[Choice]
     """A list of chat completion choices.
 
     Can be more than one if `n` is greater than 1.
@@ -66,22 +65,22 @@ class ChatCompletion(BaseModel):
     object: Literal["chat.completion"]
     """The object type, which is always `chat.completion`."""
 
-    service_tier: Optional[Literal["auto", "default", "flex", "scale"]] = None
+    service_tier: Literal["auto", "default", "flex", "scale", "priority"] | None = None
     """The service tier used for processing the request."""
 
-    system_fingerprint: Optional[str] = None
+    system_fingerprint: str | None = None
     """This fingerprint represents the backend configuration that the model runs with.
 
     Can be used in conjunction with the `seed` request parameter to understand when
     backend changes have been made that might impact determinism.
     """
 
-    usage: Optional[CompletionUsage] = None
+    usage: CompletionUsage | None = None
     """Usage statistics for the completion request."""
 
 
 class ChatCompletionExtended(ChatCompletion):
-    message_retrieval_function: Optional[Callable[[Any, "ChatCompletion"], list[ChatCompletionMessage]]] = None
-    config_id: Optional[str] = None
-    pass_filter: Optional[Callable[..., bool]] = None
-    cost: Optional[float] = None
+    message_retrieval_function: Callable[[Any, "ChatCompletion"], list[ChatCompletionMessage]] | None = None
+    config_id: str | None = None
+    pass_filter: Callable[..., bool] | None = None
+    cost: float | None = None

autogen/oai/oai_models/chat_completion_message.py

@@ -6,13 +6,15 @@
 
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import List, Optional
 
-from typing_extensions import Literal
+from typing import Literal
 
 from ._models import BaseModel
 from .chat_completion_audio import ChatCompletionAudio
-from .chat_completion_message_tool_call import ChatCompletionMessageToolCall
+from .chat_completion_message_tool_call import (
+    ChatCompletionMessageCustomToolCall,
+    ChatCompletionMessageFunctionToolCall,
+)
 
 __all__ = ["Annotation", "AnnotationURLCitation", "ChatCompletionMessage", "FunctionCall"]
 
@@ -53,34 +55,34 @@ class FunctionCall(BaseModel):
 
 
 class ChatCompletionMessage(BaseModel):
-    content: Optional[str] = None
+    content: str | None = None
     """The contents of the message."""
 
-    refusal: Optional[str] = None
+    refusal: str | None = None
     """The refusal message generated by the model."""
 
     role: Literal["assistant"]
     """The role of the author of this message."""
 
-    annotations: Optional[List[Annotation]] = None
+    annotations: list[Annotation] | None = None
     """
     Annotations for the message, when applicable, as when using the
     [web search tool](https://platform.openai.com/docs/guides/tools-web-search?api-mode=chat).
     """
 
-    audio: Optional[ChatCompletionAudio] = None
+    audio: ChatCompletionAudio | None = None
     """
     If the audio output modality is requested, this object contains data about the
     audio response from the model.
     [Learn more](https://platform.openai.com/docs/guides/audio).
     """
 
-    function_call: Optional[FunctionCall] = None
+    function_call: FunctionCall | None = None
     """Deprecated and replaced by `tool_calls`.
 
     The name and arguments of a function that should be called, as generated by the
     model.
     """
 
-    tool_calls: Optional[List[ChatCompletionMessageToolCall]] = None
+    tool_calls: list[ChatCompletionMessageFunctionToolCall | ChatCompletionMessageCustomToolCall] | None = None
     """The tool calls generated by the model, such as function calls."""

autogen/oai/oai_models/chat_completion_message_tool_call.py

@@ -6,11 +6,11 @@
 
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing_extensions import Literal
+from typing import Literal
 
 from ._models import BaseModel
 
-__all__ = ["ChatCompletionMessageToolCall", "Function"]
+__all__ = ["ChatCompletionMessageCustomToolCall", "ChatCompletionMessageFunctionToolCall", "Custom", "Function"]
 
 
 class Function(BaseModel):
@@ -26,7 +26,15 @@ class Function(BaseModel):
     """The name of the function to call."""
 
 
-class ChatCompletionMessageToolCall(BaseModel):
+class Custom(BaseModel):
+    input: str
+    """The input to the custom tool."""
+
+    name: str
+    """The name of the custom tool."""
+
+
+class ChatCompletionMessageFunctionToolCall(BaseModel):
     id: str
     """The ID of the tool call."""
 
@@ -35,3 +43,18 @@ class ChatCompletionMessageToolCall(BaseModel):
 
     type: Literal["function"]
     """The type of the tool. Currently, only `function` is supported."""
+
+
+class ChatCompletionMessageCustomToolCall(BaseModel):
+    id: str
+    """The ID of the tool call."""
+
+    custom: Custom
+    """The custom tool that the model called."""
+
+    type: Literal["custom"]
+    """The type of the tool. Currently, only `custom` is supported."""
+
+
+# Backward compatibility alias
+ChatCompletionMessageToolCall = ChatCompletionMessageFunctionToolCall

autogen/oai/oai_models/chat_completion_token_logprob.py

@@ -6,7 +6,6 @@
 
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import List, Optional
 
 from ._models import BaseModel
 
@@ -17,7 +16,7 @@ class TopLogprob(BaseModel):
     token: str
     """The token."""
 
-    bytes: Optional[List[int]] = None
+    bytes: list[int] | None = None
     """A list of integers representing the UTF-8 bytes representation of the token.
 
     Useful in instances where characters are represented by multiple tokens and
@@ -38,7 +37,7 @@ class ChatCompletionTokenLogprob(BaseModel):
     token: str
     """The token."""
 
-    bytes: Optional[List[int]] = None
+    bytes: list[int] | None = None
     """A list of integers representing the UTF-8 bytes representation of the token.
 
     Useful in instances where characters are represented by multiple tokens and
@@ -54,7 +53,7 @@ class ChatCompletionTokenLogprob(BaseModel):
     unlikely.
     """
 
-    top_logprobs: List[TopLogprob]
+    top_logprobs: list[TopLogprob]
     """List of the most likely tokens and their log probability, at this token
     position.
 

autogen/oai/oai_models/completion_usage.py

@@ -6,7 +6,6 @@
 
 # File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
 
-from typing import Optional
 
 from ._models import BaseModel
 
@@ -14,19 +13,19 @@ __all__ = ["CompletionTokensDetails", "CompletionUsage", "PromptTokensDetails"]
 
 
 class CompletionTokensDetails(BaseModel):
-    accepted_prediction_tokens: Optional[int] = None
+    accepted_prediction_tokens: int | None = None
     """
     When using Predicted Outputs, the number of tokens in the prediction that
     appeared in the completion.
     """
 
-    audio_tokens: Optional[int] = None
+    audio_tokens: int | None = None
     """Audio input tokens generated by the model."""
 
-    reasoning_tokens: Optional[int] = None
+    reasoning_tokens: int | None = None
     """Tokens generated by the model for reasoning."""
 
-    rejected_prediction_tokens: Optional[int] = None
+    rejected_prediction_tokens: int | None = None
     """
     When using Predicted Outputs, the number of tokens in the prediction that did
     not appear in the completion. However, like reasoning tokens, these tokens are
@@ -36,10 +35,10 @@ class CompletionTokensDetails(BaseModel):
 
 
 class PromptTokensDetails(BaseModel):
-    audio_tokens: Optional[int] = None
+    audio_tokens: int | None = None
     """Audio input tokens present in the prompt."""
 
-    cached_tokens: Optional[int] = None
+    cached_tokens: int | None = None
     """Cached tokens present in the prompt."""
 
 
@@ -53,8 +52,8 @@ class CompletionUsage(BaseModel):
     total_tokens: int
     """Total number of tokens used in the request (prompt + completion)."""
 
-    completion_tokens_details: Optional[CompletionTokensDetails] = None
+    completion_tokens_details: CompletionTokensDetails | None = None
     """Breakdown of tokens used in a completion."""
 
-    prompt_tokens_details: Optional[PromptTokensDetails] = None
+    prompt_tokens_details: PromptTokensDetails | None = None
     """Breakdown of tokens used in the prompt."""