ag2 0.9.7__py3-none-any.whl → 0.9.9__py3-none-any.whl

autogen/oai/ollama.py

@@ -29,12 +29,12 @@ import random
 import re
 import time
 import warnings
-from typing import Any, Literal, Optional, Union
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field, HttpUrl
 
 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 FormatterProtocol, should_hide_tools, validate_parameter
 from .oai_models import ChatCompletion, ChatCompletionMessage, ChatCompletionMessageToolCall, Choice, CompletionUsage
 
@@ -44,10 +44,23 @@ with optional_import_block():
     from ollama import Client
 
 
-@register_llm_config
+class OllamaEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["ollama"]
+    client_host: HttpUrl | None
+    stream: bool
+    num_predict: int
+    num_ctx: int
+    repeat_penalty: float
+    seed: int
+    top_k: int
+    hide_tools: Literal["if_all_run", "if_any_run", "never"]
+    native_tool_calls: bool
+
+
 class OllamaLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["ollama"] = "ollama"
-    client_host: Optional[HttpUrl] = None
+    # TODO: max_tokens
+    client_host: HttpUrl | None = None
     stream: bool = False
     num_predict: int = Field(
         default=-1,
@@ -56,9 +69,7 @@ class OllamaLLMConfigEntry(LLMConfigEntry):
     num_ctx: int = Field(default=2048)
     repeat_penalty: float = Field(default=1.1)
     seed: int = Field(default=0)
-    temperature: float = Field(default=0.8)
     top_k: int = Field(default=40)
-    top_p: float = Field(default=0.9)
     hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
     native_tool_calls: bool = False
 
@@ -101,11 +112,10 @@ class OllamaClient:
     # Override using "manual_tool_call_step2" config parameter
     TOOL_CALL_MANUAL_STEP2 = " (proceed with step 2)"
 
-    def __init__(self, response_format: Optional[Union[BaseModel, dict[str, Any]]] = None, **kwargs):
+    def __init__(self, response_format: BaseModel | dict[str, Any] | None = None, **kwargs):
         """Note that no api_key or environment variable is required for Ollama."""
-
         # Store the response format, if provided (for structured outputs)
-        self._response_format: Optional[Union[BaseModel, dict[str, Any]]] = response_format
+        self._response_format: BaseModel | dict[str, Any] | None = response_format
 
     def message_retrieval(self, response) -> list:
         """Retrieve and return a list of strings or a list of Choice.Message from the response.

autogen/oai/openai_responses.py

@@ -4,7 +4,7 @@
 
 import copy
 import warnings
-from typing import TYPE_CHECKING, Any, Tuple, Union
+from typing import TYPE_CHECKING, Any
 
 from pydantic import BaseModel
 
@@ -48,9 +48,8 @@ VALID_SIZES = {
 
 def calculate_openai_image_cost(
     model: str = "gpt-image-1", size: str = "1024x1024", quality: str = "high"
-) -> Tuple[float, str]:
-    """
-    Calculate the cost for a single image generation.
+) -> tuple[float, str]:
+    """Calculate the cost for a single image generation.
 
     Args:
         model: Model name ("gpt-image-1", "dall-e-3" or "dall-e-2")
@@ -114,7 +113,7 @@ class OpenAIResponsesClient:
     def __init__(
         self,
         client: "OpenAI",
-        response_format: Union[BaseModel, dict[str, Any], None] = None,
+        response_format: BaseModel | dict[str, Any] | None = None,
     ):
         self._oai_client = client  # plain openai.OpenAI instance
         self.response_format = response_format  # kept for parity but unused for now
@@ -195,6 +194,12 @@ class OpenAIResponsesClient:
             delta_messages.append(m)
         return delta_messages[::-1]
 
+    def _parse_params(self, params: dict[str, Any]) -> None:
+        if "verbosity" in params:
+            verbosity = params.pop("verbosity")
+            params["text"] = {"verbosity": verbosity}
+        return params
+
     def create(self, params: dict[str, Any]) -> "Response":
         """Invoke `client.responses.create() or .parse()`.
 
@@ -246,7 +251,14 @@ class OpenAIResponsesClient:
                         "output": content,
                     })
                     break
-            params["input"] = input_items[::-1]
+
+            # Ensure we have at least one valid input item
+            if input_items:
+                params["input"] = input_items[::-1]
+            else:
+                # If no valid input items were created, create a default one
+                # This prevents the API error about missing required parameters
+                params["input"] = [{"role": "user", "content": [{"type": "input_text", "text": "Hello"}]}]
 
         # Initialize tools list
         tools_list = []
@@ -274,6 +286,11 @@ class OpenAIResponsesClient:
                 UserWarning,
             )
 
+        # Validate that we have at least one of the required parameters
+        if not any(key in params for key in ["input", "previous_response_id", "prompt"]):
+            # If we still don't have any required parameters, create a minimal input
+            params["input"] = [{"role": "user", "content": [{"type": "input_text", "text": "Hello"}]}]
+
         # ------------------------------------------------------------------
         # Structured output handling - mimic OpenAIClient behaviour
         # ------------------------------------------------------------------
@@ -306,7 +323,6 @@ class OpenAIResponsesClient:
                     kwargs["text_format"] = type_to_response_format_param(rf)
                 if "response_format" in kwargs:
                     kwargs["text_format"] = kwargs.pop("response_format")
-
                 try:
                     return self._oai_client.responses.parse(**kwargs)
                 except TypeError as e:
@@ -325,22 +341,19 @@ class OpenAIResponsesClient:
             response = _create_or_parse(**params)
             self.previous_response_id = response.id
             return response
-
         # No structured output
+        params = self._parse_params(params)
         response = self._oai_client.responses.create(**params)
         self.previous_response_id = response.id
-
         # Accumulate image costs
         self._add_image_cost(response)
-
         return response
 
-    def message_retrieval(
-        self, response
-    ) -> Union[list[str], list["ModelClient.ModelClientResponseProtocol.Choice.Message"]]:
+    def message_retrieval(self, response) -> list[str] | list["ModelClient.ModelClientResponseProtocol.Choice.Message"]:
         output = getattr(response, "output", [])
-        content = []  # list[dict[str, Union[str, dict[str, Any]]]]]
+        content = []
         tool_calls = []
+
         for item in output:
             # Convert pydantic objects to plain dicts for uniform handling
             if hasattr(item, "model_dump"):
@@ -348,16 +361,26 @@ class OpenAIResponsesClient:
 
             item_type = item.get("type")
 
-            # ------------------------------------------------------------------
-            # 1) Normal messages
-            # ------------------------------------------------------------------
+            # Skip reasoning items - they're not messages
+            if item_type == "reasoning":
+                continue
+
             if item_type == "message":
                 new_item = copy.deepcopy(item)
                 new_item["type"] = "text"
                 new_item["role"] = "assistant"
+
                 blocks = item.get("content", [])
                 if len(blocks) == 1 and blocks[0].get("type") == "output_text":
                     new_item["text"] = blocks[0]["text"]
+                elif len(blocks) > 0:
+                    # Handle multiple content blocks
+                    text_parts = []
+                    for block in blocks:
+                        if block.get("type") == "output_text":
+                            text_parts.append(block.get("text", ""))
+                    new_item["text"] = " ".join(text_parts)
+
                 if "content" in new_item:
                     del new_item["content"]
                 content.append(new_item)

autogen/oai/openai_utils.py

@@ -17,7 +17,7 @@ import time
 import warnings
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Union
 
 from dotenv import find_dotenv, load_dotenv
 from packaging.version import parse
@@ -27,8 +27,9 @@ if TYPE_CHECKING:
     from openai import OpenAI
     from openai.types.beta.assistant import Assistant
 
+    from ..llm_config import LLMConfig
+
 from ..doc_utils import export_module
-from ..llm_config import LLMConfig
 
 NON_CACHE_KEY = [
     "api_key",
@@ -58,6 +59,15 @@ OAI_PRICE1K = {
     # o3
     "o3": (0.0011, 0.0044),
     "o3-mini-2025-01-31": (0.0011, 0.0044),
+    # gpt-5
+    "gpt-5": (0.00125, 0.00125),
+    "gpt-5-2025-08-07": (0.00125, 0.00125),
+    # gpt-5-mini
+    "gpt-5-mini": (0.00025, 0.00025),
+    "gpt-5-mini-2025-08-07": (0.00025, 0.00025),
+    # gpt-5-nano
+    "gpt-5-nano": (0.00005, 0.00005),
+    "gpt-5-nano-2025-08-07": (0.00005, 0.00005),
     # gpt-4o
     "gpt-4o": (0.005, 0.015),
     "gpt-4o-2024-05-13": (0.005, 0.015),
@@ -157,9 +167,9 @@ def is_valid_api_key(api_key: str) -> bool:
 @export_module("autogen")
 def get_config_list(
     api_keys: list[str],
-    base_urls: Optional[list[str]] = None,
-    api_type: Optional[str] = None,
-    api_version: Optional[str] = None,
+    base_urls: list[str] | None = None,
+    api_type: str | None = None,
+    api_version: str | None = None,
 ) -> list[dict[str, Any]]:
     """Get a list of configs for OpenAI API client.
 
@@ -208,7 +218,7 @@ def get_config_list(
 
 @export_module("autogen")
 def get_first_llm_config(
-    llm_config: Union[LLMConfig, dict[str, Any]],
+    llm_config: Union["LLMConfig", dict[str, Any]],
 ) -> dict[str, Any]:
     """Get the first LLM config from the given LLM config.
 
@@ -236,12 +246,12 @@ def get_first_llm_config(
 
 @export_module("autogen")
 def config_list_openai_aoai(
-    key_file_path: Optional[str] = ".",
-    openai_api_key_file: Optional[str] = "key_openai.txt",
-    aoai_api_key_file: Optional[str] = "key_aoai.txt",
-    openai_api_base_file: Optional[str] = "base_openai.txt",
-    aoai_api_base_file: Optional[str] = "base_aoai.txt",
-    exclude: Optional[str] = None,
+    key_file_path: str | None = ".",
+    openai_api_key_file: str | None = "key_openai.txt",
+    aoai_api_key_file: str | None = "key_aoai.txt",
+    openai_api_base_file: str | None = "base_openai.txt",
+    aoai_api_base_file: str | None = "base_aoai.txt",
+    exclude: str | None = None,
 ) -> list[dict[str, Any]]:
     """Get a list of configs for OpenAI API client (including Azure or local model deployments that support OpenAI's chat completion API).
 
@@ -366,12 +376,12 @@ def config_list_openai_aoai(
 
 @export_module("autogen")
 def config_list_from_models(
-    key_file_path: Optional[str] = ".",
-    openai_api_key_file: Optional[str] = "key_openai.txt",
-    aoai_api_key_file: Optional[str] = "key_aoai.txt",
-    aoai_api_base_file: Optional[str] = "base_aoai.txt",
-    exclude: Optional[str] = None,
-    model_list: Optional[list[str]] = None,
+    key_file_path: str | None = ".",
+    openai_api_key_file: str | None = "key_openai.txt",
+    aoai_api_key_file: str | None = "key_aoai.txt",
+    aoai_api_base_file: str | None = "base_aoai.txt",
+    exclude: str | None = None,
+    model_list: list[str] | None = None,
 ) -> list[dict[str, Any]]:
     """Get a list of configs for API calls with models specified in the model list.
 
@@ -433,11 +443,11 @@ def config_list_from_models(
 
 @export_module("autogen")
 def config_list_gpt4_gpt35(
-    key_file_path: Optional[str] = ".",
-    openai_api_key_file: Optional[str] = "key_openai.txt",
-    aoai_api_key_file: Optional[str] = "key_aoai.txt",
-    aoai_api_base_file: Optional[str] = "base_aoai.txt",
-    exclude: Optional[str] = None,
+    key_file_path: str | None = ".",
+    openai_api_key_file: str | None = "key_openai.txt",
+    aoai_api_key_file: str | None = "key_aoai.txt",
+    aoai_api_base_file: str | None = "base_aoai.txt",
+    exclude: str | None = None,
 ) -> list[dict[str, Any]]:
     """Get a list of configs for 'gpt-4' followed by 'gpt-3.5-turbo' API calls.
 
@@ -464,7 +474,7 @@ def config_list_gpt4_gpt35(
 @export_module("autogen")
 def filter_config(
     config_list: list[dict[str, Any]],
-    filter_dict: Optional[dict[str, Union[list[Union[str, None]], set[Union[str, None]]]]],
+    filter_dict: dict[str, list[str | None] | set[str | None]] | None,
     exclude: bool = False,
 ) -> list[dict[str, Any]]:
     """Filter configuration dictionaries based on specified criteria.
@@ -621,8 +631,8 @@ def _satisfies_criteria(config_value: Any, criteria_values: Any) -> bool:
 @export_module("autogen")
 def config_list_from_json(
     env_or_file: str,
-    file_location: Optional[str] = "",
-    filter_dict: Optional[dict[str, Union[list[Union[str, None]], set[Union[str, None]]]]] = None,
+    file_location: str | None = "",
+    filter_dict: dict[str, list[str | None] | set[str | None]] | None = None,
 ) -> list[dict[str, Any]]:
     """Retrieves a list of API configurations from a JSON stored in an environment variable or a file.
 
@@ -690,10 +700,10 @@ def config_list_from_json(
 
 
 def get_config(
-    api_key: Optional[str],
-    base_url: Optional[str] = None,
-    api_type: Optional[str] = None,
-    api_version: Optional[str] = None,
+    api_key: str | None,
+    base_url: str | None = None,
+    api_type: str | None = None,
+    api_version: str | None = None,
 ) -> dict[str, Any]:
     """Constructs a configuration dictionary for a single model with the provided API configurations.
 
@@ -729,10 +739,10 @@ def get_config(
 
 @export_module("autogen")
 def config_list_from_dotenv(
-    dotenv_file_path: Optional[str] = None,
-    model_api_key_map: Optional[dict[str, Any]] = None,
-    filter_dict: Optional[dict[str, Union[list[Union[str, None]], set[Union[str, None]]]]] = None,
-) -> list[dict[str, Union[str, set[str]]]]:
+    dotenv_file_path: str | None = None,
+    model_api_key_map: dict[str, Any] | None = None,
+    filter_dict: dict[str, list[str | None] | set[str | None]] | None = None,
+) -> list[dict[str, str | set[str]]]:
     """Load API configurations from a specified .env file or environment variables and construct a list of configurations.
 
     This function will:
@@ -797,12 +807,12 @@ def config_list_from_dotenv(
             config_without_key_var = {k: v for k, v in config.items() if k != "api_key_env_var"}
             config_dict = get_config(api_key=api_key, **config_without_key_var)
         else:
-            logging.warning(f"Unsupported type {type(config)} for model {model} configuration")
-
-        if not config_dict["api_key"] or config_dict["api_key"].strip() == "":
             logging.warning(
-                f"API key not found or empty for model {model}. Please ensure path to .env file is correct."
+                "Unsupported configuration type encountered for a model. Please check your model_api_key_map."
             )
+
+        if not config_dict["api_key"] or config_dict["api_key"].strip() == "":
+            logging.warning("API key not found or empty for a model. Please ensure path to .env file is correct.")
             continue  # Skip this configuration and continue with the next
 
         # Add model to the configuration and append to the list

autogen/oai/together.py

@@ -33,12 +33,13 @@ import copy
 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
 
@@ -46,22 +47,36 @@ with optional_import_block():
     from together import Together
 
 
-@register_llm_config
+class TogetherEntryDict(LLMConfigEntryDict, total=False):
+    api_type: Literal["together"]
+
+    stream: bool
+    top_k: int | None
+    repetition_penalty: float | None
+    presence_penalty: float | None
+    frequency_penalty: float | None
+    min_p: float | None
+    safety_model: str | None
+    hide_tools: Literal["if_all_run", "if_any_run", "never"]
+    price: list[float] | None
+    tool_choice: str | dict[str, str | dict[str, str]] | None
+
+
 class TogetherLLMConfigEntry(LLMConfigEntry):
     api_type: Literal["together"] = "together"
+
     max_tokens: int = Field(default=512, ge=0)
+
     stream: bool = False
-    temperature: Optional[float] = Field(default=None)
-    top_p: Optional[float] = Field(default=None)
-    top_k: Optional[int] = Field(default=None)
-    repetition_penalty: Optional[float] = Field(default=None)
-    presence_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
-    frequency_penalty: Optional[float] = Field(default=None, ge=-2, le=2)
-    min_p: Optional[float] = Field(default=None, ge=0, le=1)
-    safety_model: Optional[str] = None
+    top_k: int | None = Field(default=None)
+    repetition_penalty: float | None = Field(default=None)
+    presence_penalty: float | None = Field(default=None, ge=-2, le=2)
+    frequency_penalty: float | None = Field(default=None, ge=-2, le=2)
+    min_p: float | None = Field(default=None, ge=0, le=1)
+    safety_model: str | None = None
     hide_tools: Literal["if_all_run", "if_any_run", "never"] = "never"
-    price: Optional[list[float]] = Field(default=None, min_length=2, max_length=2)
-    tool_choice: Optional[Union[str, dict[str, Union[str, dict[str, str]]]]] = (
+    price: list[float] | None = Field(default=None, min_length=2, max_length=2)
+    tool_choice: str | dict[str, str | dict[str, str]] | None = (
         None  # dict is the tool to call: {"type": "function", "function": {"name": "my_function"}}
     )
 
@@ -72,7 +87,7 @@ class TogetherLLMConfigEntry(LLMConfigEntry):
 class TogetherClient:
     """Client for Together.AI's API."""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Unpack[TogetherEntryDict]):
         """Requires api_key or environment variable to be set
 
         Args:

autogen/retrieve_utils.py

@@ -9,7 +9,8 @@ import hashlib
 import logging
 import os
 import re
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any
 from urllib.parse import urlparse
 
 import requests
@@ -158,11 +159,11 @@ def extract_text_from_pdf(file: str) -> str:
 
 
 def split_files_to_chunks(
-    files: list[Union[tuple[str, str], str]],
+    files: list[tuple[str, str] | str],
     max_tokens: int = 4000,
     chunk_mode: str = "multi_lines",
     must_break_at_empty_line: bool = True,
-    custom_text_split_function: Optional[Callable[[str], list[str]]] = None,
+    custom_text_split_function: Callable[[str], list[str]] | None = None,
 ) -> tuple[list[str], list[dict[str, Any]]]:
     """Split a list of files into chunks of max_tokens."""
     chunks = []
@@ -200,9 +201,7 @@ def split_files_to_chunks(
     return chunks, sources
 
 
-def get_files_from_dir(
-    dir_path: Union[str, list[str]], types: list[str] = TEXT_FORMATS, recursive: bool = True
-) -> list[Any]:
+def get_files_from_dir(dir_path: str | list[str], types: list[str] = TEXT_FORMATS, recursive: bool = True) -> list[Any]:
     """Return a list of all the files in a given directory, a url, a file path or a list of them."""
     if len(types) == 0:
         raise ValueError("types cannot be empty.")
@@ -343,7 +342,7 @@ def is_url(string: str):
 
 @require_optional_import("chromadb", "retrievechat")
 def create_vector_db_from_dir(
-    dir_path: Union[str, list[str]],
+    dir_path: str | list[str],
     max_tokens: int = 4000,
     client: "API" = None,
     db_path: str = "tmp/chromadb.db",

autogen/runtime_logging.py

@@ -9,7 +9,8 @@ from __future__ import annotations
 import logging
 import sqlite3
 import uuid
-from typing import TYPE_CHECKING, Any, Callable, Literal, Optional, TypeVar
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 from .logger.base_logger import BaseLogger, LLMConfig
 from .logger.logger_factory import LoggerFactory
@@ -38,9 +39,9 @@ F = TypeVar("F", bound=Callable[..., Any])
 
 
 def start(
-    logger: Optional[BaseLogger] = None,
+    logger: BaseLogger | None = None,
     logger_type: Literal["sqlite", "file"] = "sqlite",
-    config: Optional[dict[str, Any]] = None,
+    config: dict[str, Any] | None = None,
 ) -> str:
     """Start logging for the runtime.
 
@@ -71,7 +72,7 @@ def log_chat_completion(
     wrapper_id: int,
     agent: str | Agent,
     request: dict[str, float | str | list[dict[str, str]]],
-    response: str | "ChatCompletion",
+    response: str | ChatCompletion,
     is_cached: int,
     cost: float,
     start_time: str,

autogen/token_count_utils.py

@@ -8,7 +8,7 @@
 import json
 import logging
 import re
-from typing import Any, Union
+from typing import Any
 
 import tiktoken
 
@@ -63,6 +63,9 @@ def get_max_token_limit(model: str = "gpt-3.5-turbo-0613") -> int:
         "gpt-4o-2024-11-20": 128000,
         "gpt-4o-mini": 128000,
         "gpt-4o-mini-2024-07-18": 128000,
+        "gpt-5": 128000,
+        "gpt-5-mini": 128000,
+        "gpt-5-nano": 128000,
     }
     return max_token_limit[model]
 
@@ -71,7 +74,7 @@ def percentile_used(input, model="gpt-3.5-turbo-0613"):
     return count_token(input) / get_max_token_limit(model)
 
 
-def token_left(input: Union[str, list[str], dict[str, Any]], model="gpt-3.5-turbo-0613") -> int:
+def token_left(input: str | list[str] | dict[str, Any], model="gpt-3.5-turbo-0613") -> int:
     """Count number of tokens left for an OpenAI model.
 
     Args:
@@ -84,7 +87,7 @@ def token_left(input: Union[str, list[str], dict[str, Any]], model="gpt-3.5-turb
     return get_max_token_limit(model) - count_token(input, model=model)
 
 
-def count_token(input: Union[str, list[str], dict[str, Any]], model: str = "gpt-3.5-turbo-0613") -> int:
+def count_token(input: str | list[str] | dict[str, Any], model: str = "gpt-3.5-turbo-0613") -> int:
     """Count number of tokens used by an OpenAI model.
 
     Args:
@@ -112,7 +115,7 @@ def _num_token_from_text(text: str, model: str = "gpt-3.5-turbo-0613"):
     return len(encoding.encode(text))
 
 
-def _num_token_from_messages(messages: Union[list[str], dict[str, Any]], model="gpt-3.5-turbo-0613"):
+def _num_token_from_messages(messages: list[str] | dict[str, Any], model="gpt-3.5-turbo-0613"):
     """Return the number of tokens used by a list of messages.
 
     retrieved from https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb/

autogen/tools/contrib/time/time.py

@@ -26,7 +26,6 @@ class TimeTool(Tool):
         Args:
             date_time_format (str, optional): The format of the date and time. Defaults to "%Y-%m-%d %H:%M:%S".
         """
-
         self._date_time_format = date_time_format
 
         async def get_date_and_time(

autogen/tools/dependency_injection.py

@@ -6,8 +6,9 @@ import functools
 import inspect
 import sys
 from abc import ABC
+from collections.abc import Callable, Iterable
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Iterable, Optional, TypeVar, Union, get_type_hints
+from typing import TYPE_CHECKING, Any, TypeVar, get_type_hints
 
 from ..agentchat import Agent
 from ..doc_utils import export_module
@@ -67,7 +68,7 @@ class ChatContext(BaseContext):
         return self._agent.chat_messages
 
     @property
-    def last_message(self) -> Optional[dict[str, Any]]:
+    def last_message(self) -> dict[str, Any] | None:
         """The last message in the chat.
 
         Returns:
@@ -102,7 +103,7 @@ def Depends(x: Any) -> Any:  # noqa: N802
     return FastDepends(x)
 
 
-def get_context_params(func: Callable[..., Any], subclass: Union[type[BaseContext], type[ChatContext]]) -> list[str]:
+def get_context_params(func: Callable[..., Any], subclass: type[BaseContext] | type[ChatContext]) -> list[str]:
     """Gets the names of the context parameters in a function signature.
 
     Args:
@@ -116,9 +117,7 @@ def get_context_params(func: Callable[..., Any], subclass: Union[type[BaseContex
     return [p.name for p in sig.parameters.values() if _is_context_param(p, subclass=subclass)]
 
 
-def _is_context_param(
-    param: inspect.Parameter, subclass: Union[type[BaseContext], type[ChatContext]] = BaseContext
-) -> bool:
+def _is_context_param(param: inspect.Parameter, subclass: type[BaseContext] | type[ChatContext] = BaseContext) -> bool:
     # param.annotation.__args__[0] is used to handle Annotated[MyContext, Depends(MyContext(b=2))]
     param_annotation = param.annotation.__args__[0] if hasattr(param.annotation, "__args__") else param.annotation
     try: