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

autogen/llm_config/entry.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import re
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from httpx import Client as httpxClient
+from pydantic import BaseModel, ConfigDict, Field, HttpUrl, SecretStr, ValidationInfo, field_serializer, field_validator
+from typing_extensions import Required, TypedDict
+
+from .client import ModelClient
+
+
+class LLMConfigEntryDict(TypedDict, total=False):
+    api_type: Required[str]
+    model: str
+    max_tokens: int | None
+    top_p: float | None
+    temperature: float | None
+
+    api_key: SecretStr | None
+    api_version: str | None
+    base_url: HttpUrl | None
+    voice: str | None
+    http_client: httpxClient | None
+    model_client_cls: str | None
+    response_format: str | dict[str, Any] | BaseModel | type[BaseModel] | None
+    default_headers: Mapping[str, Any] | None
+    tags: list[str]
+
+
+class ApplicationConfig(BaseModel):
+    max_tokens: int | None = Field(default=None, ge=0)
+    top_p: float | None = Field(default=None, gt=0, lt=1)
+    temperature: float | None = Field(default=None, ge=0, le=1)
+
+    @field_validator("top_p", mode="before")
+    @classmethod
+    def check_top_p(cls, v: float | None, info: ValidationInfo) -> float | None:
+        if v is not None and info.data.get("temperature") is not None:
+            raise ValueError("temperature and top_p cannot be set at the same time.")
+        return v
+
+    @field_validator("temperature", mode="before")
+    @classmethod
+    def check_temperature(cls, v: float | None, info: ValidationInfo) -> float | None:
+        if v is not None and info.data.get("top_p") is not None:
+            raise ValueError("temperature and top_p cannot be set at the same time.")
+        return v
+
+
+class LLMConfigEntry(ApplicationConfig, ABC):
+    api_type: str
+    model: str = Field(..., min_length=1)
+
+    api_key: SecretStr | None = None
+    api_version: str | None = None
+    base_url: HttpUrl | None = None
+    voice: str | None = None
+    model_client_cls: str | None = None
+    http_client: httpxClient | None = None
+    response_format: str | dict[str, Any] | BaseModel | type[BaseModel] | None = None
+    default_headers: Mapping[str, Any] | None = None
+    tags: list[str] = Field(default_factory=list)
+
+    # Following field is configuration for pydantic to disallow extra fields
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=True)
+
+    def apply_application_config(self, application_config: ApplicationConfig) -> "LLMConfigEntry":
+        """Apply application level configurations."""
+        # TODO: should create a new instance instead of mutating current one
+        self.max_tokens = self.max_tokens or application_config.max_tokens
+        self.top_p = self.top_p or application_config.top_p
+        self.temperature = self.temperature or application_config.temperature
+        return self
+
+    @abstractmethod
+    def create_client(self) -> "ModelClient": ...
+
+    @field_validator("base_url", mode="before")
+    @classmethod
+    def check_base_url(cls, v: HttpUrl | str | None, info: ValidationInfo) -> str | None:
+        if v is None:  # Handle None case explicitly
+            return None
+        if not str(v).startswith("https://") and not str(v).startswith("http://"):
+            return f"http://{str(v)}"
+        return str(v)
+
+    @field_serializer("base_url", when_used="unless-none")  # Ensure serializer also respects None
+    def serialize_base_url(self, v: HttpUrl | None) -> str | None:
+        return str(v) if v is not None else None
+
+    @field_serializer("api_key", when_used="unless-none")
+    def serialize_api_key(self, v: SecretStr) -> str:
+        return v.get_secret_value()
+
+    def model_dump(self, *args: Any, exclude_none: bool = True, **kwargs: Any) -> dict[str, Any]:
+        return BaseModel.model_dump(self, exclude_none=exclude_none, *args, **kwargs)
+
+    def model_dump_json(self, *args: Any, exclude_none: bool = True, **kwargs: Any) -> str:
+        return BaseModel.model_dump_json(self, exclude_none=exclude_none, *args, **kwargs)
+
+    def get(self, key: str, default: Any | None = None) -> Any:
+        val = getattr(self, key, default)
+        if isinstance(val, SecretStr):
+            return val.get_secret_value()
+        return val
+
+    def __getitem__(self, key: str) -> Any:
+        try:
+            val = getattr(self, key)
+            if isinstance(val, SecretStr):
+                return val.get_secret_value()
+            return val
+        except AttributeError:
+            raise KeyError(f"Key '{key}' not found in {self.__class__.__name__}")
+
+    def __setitem__(self, key: str, value: Any) -> None:
+        setattr(self, key, value)
+
+    def __contains__(self, key: str) -> bool:
+        return hasattr(self, key)
+
+    def items(self) -> Iterable[tuple[str, Any]]:
+        d = self.model_dump()
+        return d.items()
+
+    def keys(self) -> Iterable[str]:
+        d = self.model_dump()
+        return d.keys()
+
+    def values(self) -> Iterable[Any]:
+        d = self.model_dump()
+        return d.values()
+
+    def __repr__(self) -> str:
+        # Override to eliminate none values from the repr
+        d = self.model_dump()
+        r = [f"{k}={repr(v)}" for k, v in d.items()]
+
+        s = f"{self.__class__.__name__}({', '.join(r)})"
+
+        # Replace any keys ending with '_key' or '_token' values with stars for security
+        # This regex will match any key ending with '_key' or '_token' and its value, and replace the value with stars
+        # It also captures the type of quote used (single or double) and reuses it in the replacement
+        s = re.sub(r'(\w+_(key|token)\s*=\s*)([\'"]).*?\3', r"\1\3**********\3", s, flags=re.IGNORECASE)
+
+        return s
+
+    def __str__(self) -> str:
+        return repr(self)

autogen/logger/base_logger.py

@@ -9,7 +9,8 @@ from __future__ import annotations
 import sqlite3
 import uuid
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any, Callable, TypeVar, Union
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
 
 if TYPE_CHECKING:
     from openai import AzureOpenAI, OpenAI
@@ -18,8 +19,8 @@ if TYPE_CHECKING:
     from .. import Agent, ConversableAgent, OpenAIWrapper
 
 F = TypeVar("F", bound=Callable[..., Any])
-ConfigItem = dict[str, Union[str, list[str]]]
-LLMConfig = dict[str, Union[None, float, int, ConfigItem, list[ConfigItem]]]
+ConfigItem = dict[str, str | list[str]]
+LLMConfig = dict[str, None | float | int | ConfigItem | list[ConfigItem]]
 
 
 class BaseLogger(ABC):

autogen/logger/file_logger.py

@@ -11,7 +11,8 @@ import logging
 import os
 import threading
 import uuid
-from typing import TYPE_CHECKING, Any, Callable, TypeVar
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from ..doc_utils import export_module
 from .base_logger import BaseLogger, LLMConfig

autogen/logger/logger_factory.py

@@ -4,7 +4,7 @@
 #
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 from ..doc_utils import export_module
 from .base_logger import BaseLogger
@@ -20,7 +20,7 @@ class LoggerFactory:
 
     @staticmethod
     def get_logger(
-        logger_type: Literal["sqlite", "file"] = "sqlite", config: Optional[dict[str, Any]] = None
+        logger_type: Literal["sqlite", "file"] = "sqlite", config: dict[str, Any] | None = None
     ) -> BaseLogger:
         """Factory method to create logger objects.
 

autogen/logger/logger_utils.py

@@ -7,7 +7,7 @@
 import inspect
 from datetime import datetime, timezone
 from pathlib import Path, PurePath
-from typing import Any, Union
+from typing import Any
 
 __all__ = ("get_current_ts", "to_dict")
 
@@ -22,7 +22,7 @@ def get_current_ts() -> str:
 
 
 def to_dict(
-    obj: Union[int, float, str, bool, dict[Any, Any], list[Any], tuple[Any, ...], Any],
+    obj: int | float | str | bool | dict[Any, Any] | list[Any] | tuple[Any, ...] | Any,
     exclude: tuple[str, ...] = (),
     no_recursive: tuple[Any, ...] = (),
 ) -> Any:

autogen/logger/sqlite_logger.py

@@ -12,7 +12,8 @@ import os
 import sqlite3
 import threading
 import uuid
-from typing import TYPE_CHECKING, Any, Callable, TypeVar
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, TypeVar
 
 from ..doc_utils import export_module
 from .base_logger import BaseLogger, LLMConfig

autogen/math_utils.py

@@ -4,10 +4,9 @@
 #
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
-from typing import Optional
 
 
-def remove_boxed(string: str) -> Optional[str]:
+def remove_boxed(string: str) -> str | None:
     """Source: https://github.com/hendrycks/math
     Extract the text within a \\boxed`{...}` environment.
 
@@ -28,7 +27,7 @@ def remove_boxed(string: str) -> Optional[str]:
         return None
 
 
-def last_boxed_only_string(string: str) -> Optional[str]:
+def last_boxed_only_string(string: str) -> str | None:
     """Source: https://github.com/hendrycks/math
     Extract the last \\boxed`{...}` or \\fbox`{...}` element from a string.
     """
@@ -239,7 +238,7 @@ def _strip_string(string: str) -> str:
     return string
 
 
-def get_answer(solution: Optional[str]) -> Optional[str]:
+def get_answer(solution: str | None) -> str | None:
     if solution is None:
         return None
     last_boxed = last_boxed_only_string(solution)
@@ -251,7 +250,7 @@ def get_answer(solution: Optional[str]) -> Optional[str]:
     return answer
 
 
-def is_equiv(str1: Optional[str], str2: Optional[str]) -> float:
+def is_equiv(str1: str | None, str2: str | None) -> float:
     """Returns (as a float) whether two strings containing math are equivalent up to differences of formatting in
     - units
     - fractions

autogen/mcp/__main__.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import logging
-from typing import Annotated, Literal, Optional
+from typing import Annotated, Literal
 
 from .. import __version__
 from ..import_utils import optional_import_block, require_optional_import
@@ -27,7 +27,7 @@ def create_typer_app() -> "typer.Typer":
     @app.callback()
     def callback(
         version: Annotated[
-            Optional[bool],
+            bool | None,
             typer.Option("--version", help="Show the version and exit.", callback=version_callback),
         ] = None,
     ) -> None:
@@ -41,19 +41,19 @@ def create_typer_app() -> "typer.Typer":
     @app.command()
     def create(
         openapi_specification: Annotated[
-            Optional[str],
+            str | None,
             "Specification of the OpenAPI to use for the proxy generation.",
         ] = None,
         openapi_url: Annotated[
-            Optional[str],
+            str | None,
             "URL to the OpenAPI specification to use for the proxy generation.",
         ] = None,
         client_source_path: Annotated[
-            Optional[str],
+            str | None,
             "Path to the generated proxy client source code.",
         ] = None,
         server_url: Annotated[
-            Optional[str],
+            str | None,
             "Comma-separated list of server URLs to use for the proxy generation.",
         ] = None,
         configuration_type: Annotated[

autogen/mcp/helpers.py

@@ -5,21 +5,21 @@ import asyncio
 import os
 import signal
 from asyncio.subprocess import PIPE, Process, create_subprocess_exec
+from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import AsyncGenerator, Dict, Optional
 
 
 @asynccontextmanager
 async def run_streamable_http_client(
-    *, mcp_server_path: str, env_vars: Optional[Dict[str, str]] = None, startup_wait_secs: float = 5.0
+    *, mcp_server_path: str, env_vars: dict[str, str] | None = None, startup_wait_secs: float = 5.0
 ) -> AsyncGenerator[Process, None]:
-    """
-    Async context manager to run a Python subprocess for streamable-http with custom env vars.
+    """Async context manager to run a Python subprocess for streamable-http with custom env vars.
 
     Args:
         mcp_server_path: Path to the Python script to run.
         env_vars: Environment variables to export to the subprocess.
         startup_wait_secs: Time to wait for the server to start (in seconds).
+
     Yields:
         An asyncio.subprocess.Process object.
     """

autogen/mcp/mcp_client.py

@@ -3,38 +3,153 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-import sys
-from datetime import datetime
+from collections.abc import AsyncIterator
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from datetime import datetime, timedelta
 from pathlib import Path
-from typing import Annotated, Any, Optional, Union
+from typing import Annotated, Any, Literal, Protocol, cast
 
 import anyio
-from pydantic import BaseModel
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp.client.session import ClientSession
+from mcp.client.sse import sse_client
+from mcp.client.stdio import StdioServerParameters, stdio_client
+from pydantic import AnyUrl, BaseModel, Field
 
 from ..doc_utils import export_module
 from ..import_utils import optional_import_block, require_optional_import
 from ..tools import Tool, Toolkit
 
 with optional_import_block():
-    from mcp import ClientSession
-    from mcp.types import (
-        CallToolResult,
-        ReadResourceResult,
-        ResourceTemplate,
-        TextContent,
-    )
-    from mcp.types import (
-        Tool as MCPTool,
-    )
+    from mcp.shared.message import SessionMessage
+    from mcp.types import CallToolResult, ReadResourceResult, ResourceTemplate, TextContent
+    from mcp.types import Tool as MCPTool
 
 __all__ = ["ResultSaved", "create_toolkit"]
 
+# Type definitions
+EncodingErrorHandlerType = Literal["strict", "ignore", "replace"]
+
+# Default constants
+DEFAULT_TEXT_ENCODING = "utf-8"
+DEFAULT_TEXT_ENCODING_ERROR_HANDLER: EncodingErrorHandlerType = "strict"
+DEFAULT_HTTP_REQUEST_TIMEOUT = 5
+DEFAULT_SSE_EVENT_READ_TIMEOUT = 60 * 5
+DEFAULT_STREAMABLE_HTTP_REQUEST_TIMEOUT = timedelta(seconds=30)
+DEFAULT_STREAMABLE_HTTP_SSE_EVENT_READ_TIMEOUT = timedelta(seconds=60 * 5)
+
+
+class SessionConfigProtocol(Protocol):
+    """Protocol for session configuration classes that can create MCP sessions."""
+
+    server_name: str
+
+    @asynccontextmanager
+    async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
+        """Create a session using the given exit stack."""
+        try:
+            yield cast(ClientSession, None)  # placeholder yield to satisfy AsyncIterator type
+        except Exception:
+            raise NotImplementedError
+
+
+class BasicSessionConfig(BaseModel):
+    """Basic session configuration."""
+
+    server_name: str = Field(..., description="Name of the server")
+
+    async def initialize(
+        self,
+        client: AbstractAsyncContextManager[
+            tuple[
+                MemoryObjectReceiveStream[SessionMessage | Exception],
+                MemoryObjectSendStream[SessionMessage],
+            ]
+        ],
+        exit_stack: AsyncExitStack,
+    ) -> ClientSession:
+        """Initialize the session."""
+        reader, writer = await exit_stack.enter_async_context(client)
+        session = cast(
+            ClientSession,
+            await exit_stack.enter_async_context(ClientSession(reader, writer)),
+        )
+        return session
+
+
+class SseConfig(BasicSessionConfig):
+    """Configuration for a single SSE MCP server."""
+
+    url: str = Field(..., description="URL of the SSE server")
+    headers: dict[str, Any] | None = Field(default=None, description="HTTP headers to send to the SSE endpoint")
+    timeout: float = Field(default=DEFAULT_HTTP_REQUEST_TIMEOUT, description="HTTP timeout")
+    sse_read_timeout: float = Field(default=DEFAULT_SSE_EVENT_READ_TIMEOUT, description="SSE read timeout")
+
+    @asynccontextmanager
+    async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
+        """
+        Create a new session to an MCP server using SSE transport.
+
+        Args:
+            exit_stack: AsyncExitStack for managing async resources
+
+        Yields:
+            ClientSession: The MCP client session
+        """
+        # Create and store the connection
+        client = sse_client(self.url, self.headers, self.timeout, self.sse_read_timeout)
+        yield await self.initialize(client, exit_stack)
+
+
+class StdioConfig(BasicSessionConfig):
+    """Configuration for a single stdio MCP server."""
+
+    command: str = Field(..., description="Command to execute")
+    args: list[str] = Field(..., description="Arguments for the command")
+    transport: Literal["stdio"] = Field(default="stdio", description="Transport type")
+    environment: dict[str, str] | None = Field(default=None, description="Environment variables")
+    working_dir: str | Path | None = Field(default=None, description="Working directory")
+    encoding: str = Field(default=DEFAULT_TEXT_ENCODING, description="Character encoding")
+    encoding_error_handler: EncodingErrorHandlerType = Field(
+        default=DEFAULT_TEXT_ENCODING_ERROR_HANDLER, description="How to handle encoding errors"
+    )
+    session_options: dict[str, Any] | None = Field(default=None, description="Additional session options")
+
+    @asynccontextmanager
+    async def create_session(self, exit_stack: AsyncExitStack) -> AsyncIterator[ClientSession]:
+        """
+        Create a new session to an MCP server using stdio transport.
+
+        Args:
+            exit_stack: AsyncExitStack for managing async resources
+
+        Yields:
+            ClientSession: The MCP client session
+        """
+        client = stdio_client(
+            StdioServerParameters(
+                command=self.command,
+                args=self.args,
+                env=self.environment,
+                encoding=self.encoding,
+                encoding_error_handler=self.encoding_error_handler,
+            )
+        )
+        yield await self.initialize(client, exit_stack)
+
+
+class MCPConfig(BaseModel):
+    """Configuration for multiple MCP sessions using stdio transport."""
+
+    # we should use final classes to allow pydantic to validate the type
+    servers: list[SseConfig | StdioConfig] = Field(..., description="List of stdio & sse server configurations")
+
 
 class MCPClient:
     @staticmethod
     def _convert_call_tool_result(  # type: ignore[no-any-unimported]
         call_tool_result: "CallToolResult",  # type: ignore[no-any-unimported]
-    ) -> tuple[Union[str, list[str]], Any]:
+    ) -> tuple[str | list[str], Any]:
         text_contents: list[TextContent] = []  # type: ignore[no-any-unimported]
         non_text_contents = []
         for content in call_tool_result.content:
@@ -43,7 +158,7 @@ class MCPClient:
             else:
                 non_text_contents.append(content)
 
-        tool_content: Union[str, list[str]] = [content.text for content in text_contents]
+        tool_content: str | list[str] = [content.text for content in text_contents]
         if len(text_contents) == 1:
             tool_content = tool_content[0]
 
@@ -65,7 +180,7 @@ class MCPClient:
 
         async def call_tool(  # type: ignore[no-any-unimported]
             **arguments: dict[str, Any],
-        ) -> tuple[Union[str, list[str]], Any]:
+        ) -> tuple[str | list[str], Any]:
             call_tool_result = await session.call_tool(tool.name, arguments)
             return MCPClient._convert_call_tool_result(call_tool_result)
 
@@ -83,7 +198,7 @@ class MCPClient:
         cls,
         resource_template: Any,
         session: "ClientSession",
-        resource_download_folder: Optional[Path],
+        resource_download_folder: Path | None,
         **kwargs: Any,
     ) -> Tool:
         if not isinstance(resource_template, ResourceTemplate):
@@ -97,8 +212,8 @@ Here is the correct format for the URI template:
 {mcp_resource.uriTemplate}
 """
 
-        async def call_resource(uri: Annotated[str, uri_description]) -> Union[ReadResourceResult, ResultSaved]:  # type: ignore[no-any-unimported]
-            result = await session.read_resource(uri)
+        async def call_resource(uri: Annotated[str, uri_description]) -> ReadResourceResult | ResultSaved:  # type: ignore[no-any-unimported]
+            result = await session.read_resource(AnyUrl(uri))
 
             if not resource_download_folder:
                 return result
@@ -131,7 +246,7 @@ Here is the correct format for the URI template:
         *,
         use_mcp_tools: bool,
         use_mcp_resources: bool,
-        resource_download_folder: Optional[Path],
+        resource_download_folder: Path | None,
     ) -> Toolkit:  # type: ignore[no-any-unimported]
         """Load all available MCP tools and convert them to AG2 Toolkit."""
         all_ag2_tools: list[Tool] = []
@@ -156,10 +271,7 @@ Here is the correct format for the URI template:
         return Toolkit(tools=all_ag2_tools)
 
     @classmethod
-    def get_unsupported_reason(cls) -> Optional[str]:
-        if sys.version_info < (3, 10):
-            return "This submodule is only supported for Python versions 3.10 and above"
-
+    def get_unsupported_reason(cls) -> str | None:
         with optional_import_block() as result:
             import mcp  # noqa: F401
 
@@ -169,13 +281,42 @@ Here is the correct format for the URI template:
         return None
 
 
+class MCPClientSessionManager:
+    """
+    A class to manage MCP client sessions.
+    """
+
+    def __init__(self) -> None:
+        """Initialize the MCP client session manager."""
+        self.exit_stack = AsyncExitStack()
+        self.sessions: dict[str, ClientSession] = {}
+
+    @asynccontextmanager
+    async def open_session(
+        self,
+        config: SessionConfigProtocol,
+    ) -> AsyncIterator[ClientSession]:
+        """Open a new session to an MCP server based on configuration.
+
+        Args:
+            config: SessionConfigProtocol object containing session configuration
+
+        Yields:
+            ClientSession: The MCP client session
+        """
+        async with config.create_session(self.exit_stack) as session:
+            await session.initialize()
+            self.sessions[config.server_name] = session
+            yield session
+
+
 @export_module("autogen.mcp")
 async def create_toolkit(
     session: "ClientSession",
     *,
     use_mcp_tools: bool = True,
     use_mcp_resources: bool = True,
-    resource_download_folder: Optional[Union[Path, str]] = None,
+    resource_download_folder: Path | str | None = None,
 ) -> Toolkit:  # type: ignore[no-any-unimported]
     """Create a toolkit from the MCP client session.
 
@@ -184,12 +325,12 @@ async def create_toolkit(
         use_mcp_tools (bool): Whether to include MCP tools in the toolkit.
         use_mcp_resources (bool): Whether to include MCP resources in the toolkit.
         resource_download_folder (Optional[Union[Path, str]]): The folder to download files to.
+
     Returns:
         Toolkit: The toolkit containing the converted tools.
     """
-    if resource_download_folder:
-        if isinstance(resource_download_folder, str):
-            resource_download_folder = Path(resource_download_folder)
+    if resource_download_folder is not None:
+        resource_download_folder = Path(resource_download_folder)
         await anyio.to_thread.run_sync(lambda: resource_download_folder.mkdir(parents=True, exist_ok=True))
 
     return await MCPClient.load_mcp_toolkit(

autogen/mcp/mcp_proxy/fastapi_code_generator_helpers.py

@@ -4,7 +4,6 @@
 from collections.abc import Iterator
 from contextlib import contextmanager
 from functools import cached_property
-from typing import Optional, Union
 
 from ...import_utils import optional_import_block
 
@@ -24,7 +23,7 @@ __all__ = ["SUCCESFUL_IMPORT", "patch_get_parameter_type"]
 @contextmanager
 def patch_get_parameter_type() -> Iterator[None]:
     class ArgumentWithDescription(Argument):  # type: ignore[misc]
-        description: Optional[str] = None
+        description: str | None = None
 
         @cached_property
         def argument(self) -> str:
@@ -43,10 +42,10 @@ def patch_get_parameter_type() -> Iterator[None]:
 
     def get_parameter_type(
         self: OpenAPIParser,
-        parameters: Union[ReferenceObject, ParameterObject],
+        parameters: ReferenceObject | ParameterObject,
         snake_case: bool,
         path: list[str],
-    ) -> Optional[Argument]:
+    ) -> Argument | None:
         # get the original argument
         argument = original_get_parameter_type(self, parameters, snake_case, path)