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

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)
 

autogen/mcp/mcp_proxy/mcp_proxy.py

@@ -8,7 +8,7 @@ import json
 import re
 import sys
 import tempfile
-from collections.abc import Iterable, Iterator, Mapping
+from collections.abc import Callable, Iterable, Iterator, Mapping
 from contextlib import contextmanager
 from functools import wraps
 from logging import getLogger
@@ -17,10 +17,7 @@ from types import ModuleType
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
     Literal,
-    Optional,
-    Union,
 )
 
 import requests
@@ -49,7 +46,7 @@ logger = getLogger(__name__)
 
 
 @contextmanager
-def optional_temp_path(path: Optional[str] = None) -> Iterator[Path]:
+def optional_temp_path(path: str | None = None) -> Iterator[Path]:
     if path is None:
         with tempfile.TemporaryDirectory() as temp_dir:
             yield Path(temp_dir)
@@ -74,7 +71,7 @@ def add_to_builtins(new_globals: dict[str, Any]) -> Iterator[None]:
 
 
 class MCPProxy:
-    def __init__(self, servers: list[dict[str, Any]], title: Optional[str] = None, **kwargs: Any) -> None:
+    def __init__(self, servers: list[dict[str, Any]], title: str | None = None, **kwargs: Any) -> None:
         """Proxy class to generate client from OpenAPI schema."""
         self._servers = servers
         self._title = title or "MCP Proxy"
@@ -83,7 +80,7 @@ class MCPProxy:
         self._globals: dict[str, Any] = {}
 
         self._security: dict[str, list[BaseSecurity]] = {}
-        self._security_params: dict[Optional[str], BaseSecurityParameters] = {}
+        self._security_params: dict[str | None, BaseSecurityParameters] = {}
         self._tags: set[str] = set()
 
         self._function_group: dict[str, list[str]] = {}
@@ -100,7 +97,7 @@ class MCPProxy:
         return result
 
     @staticmethod
-    def _get_params(path: str, func: Callable[..., Any]) -> tuple[set[str], set[str], Optional[str], bool]:
+    def _get_params(path: str, func: Callable[..., Any]) -> tuple[set[str], set[str], str | None, bool]:
         sig = inspect.signature(func)
 
         params_names = set(sig.parameters.keys())
@@ -118,11 +115,11 @@ class MCPProxy:
         return q_params, path_params, body, security
 
     def get_mcp(self, **settings: Any) -> "FastMCP":
-        mcp = FastMCP(title=self._title, **settings)
+        mcp = FastMCP(name=self._title, **settings)  # newer mcp
 
         for func in self._registered_funcs:
             try:
-                mcp.tool()(func)  # type: ignore [no-untyped-call]
+                mcp.tool()(func)
             except PydanticInvalidForJsonSchema as e:
                 logger.warning("Could not register function %s: %s", func.__name__, e)
 
@@ -157,7 +154,7 @@ class MCPProxy:
 
         return url, params, body_dict
 
-    def set_security_params(self, security_params: BaseSecurityParameters, name: Optional[str] = None) -> None:
+    def set_security_params(self, security_params: BaseSecurityParameters, name: str | None = None) -> None:
         if name is not None:
             security = self._security.get(name)
             if security is None:
@@ -180,7 +177,7 @@ class MCPProxy:
                 return match_security
         raise ValueError(f"Security parameters {security_params} does not match any given security {security}")
 
-    def _get_security_params(self, name: str) -> tuple[Optional[BaseSecurityParameters], Optional[BaseSecurity]]:
+    def _get_security_params(self, name: str) -> tuple[BaseSecurityParameters | None, BaseSecurity | None]:
         # check if security is set for the method
         security = self._security.get(name)
         if not security:
@@ -203,8 +200,8 @@ class MCPProxy:
         self,
         method: Literal["put", "get", "post", "head", "delete", "patch"],
         path: str,
-        description: Optional[str] = None,
-        security: Optional[list[BaseSecurity]] = None,
+        description: str | None = None,
+        security: list[BaseSecurity] | None = None,
         **kwargs: Any,
     ) -> Callable[..., dict[str, Any]]:
         def decorator(func: Callable[..., Any]) -> Callable[..., dict[str, Any]]:
@@ -275,7 +272,7 @@ class MCPProxy:
         input_text: str,
         output_dir: Path,
         disable_timestamp: bool = False,
-        custom_visitors: Optional[list[Path]] = None,
+        custom_visitors: list[Path] | None = None,
     ) -> str:
         if custom_visitors is None:
             custom_visitors = []
@@ -318,10 +315,10 @@ class MCPProxy:
     def create(
         cls,
         *,
-        openapi_specification: Optional[str] = None,
-        openapi_url: Optional[str] = None,
-        client_source_path: Optional[str] = None,
-        servers: Optional[list[dict[str, Any]]] = None,
+        openapi_specification: str | None = None,
+        openapi_url: str | None = None,
+        client_source_path: str | None = None,
+        servers: list[dict[str, Any]] | None = None,
         rename_functions: bool = False,
         group_functions: bool = False,
         configuration_type: Literal["json", "yaml"] = "json",
@@ -392,7 +389,7 @@ class MCPProxy:
 
     def dump_configurations(self, output_dir: Path) -> None:
         for tag in self._function_group:
-            output_file = output_dir / "mcp_config_{}.json".format(tag)
+            output_file = output_dir / f"mcp_config_{tag}.json"
 
             functions = [
                 registered_function
@@ -458,8 +455,8 @@ class MCPProxy:
 
     def _get_functions_to_register(
         self,
-        functions: Optional[Iterable[Union[str, Mapping[str, Mapping[str, str]]]]] = None,
-    ) -> dict[Callable[..., Any], dict[str, Union[str, None]]]:
+        functions: Iterable[str | Mapping[str, Mapping[str, str]]] | None = None,
+    ) -> dict[Callable[..., Any], dict[str, str | None]]:
         if functions is None:
             return {
                 f: {
@@ -469,7 +466,7 @@ class MCPProxy:
                 for f in self._registered_funcs
             }
 
-        functions_with_name_desc: dict[str, dict[str, Union[str, None]]] = {}
+        functions_with_name_desc: dict[str, dict[str, str | None]] = {}
 
         for f in functions:
             if isinstance(f, str):
@@ -485,7 +482,7 @@ class MCPProxy:
             else:
                 raise ValueError(f"Invalid type {type(f)} for function {f}")
 
-        funcs_to_register: dict[Callable[..., Any], dict[str, Union[str, None]]] = {
+        funcs_to_register: dict[Callable[..., Any], dict[str, str | None]] = {
             f: functions_with_name_desc[f.__name__]
             for f in self._registered_funcs
             if f.__name__ in functions_with_name_desc
@@ -528,7 +525,7 @@ class MCPProxy:
     def _register_for_llm(
         self,
         agent: "ConversableAgent",
-        functions: Optional[Iterable[Union[str, Mapping[str, Mapping[str, str]]]]] = None,
+        functions: Iterable[str | Mapping[str, Mapping[str, str]]] | None = None,
     ) -> None:
         funcs_to_register = self._get_functions_to_register(functions)
 
@@ -543,7 +540,7 @@ class MCPProxy:
     def _register_for_execution(
         self,
         agent: "ConversableAgent",
-        functions: Optional[Iterable[Union[str, Mapping[str, Mapping[str, str]]]]] = None,
+        functions: Iterable[str | Mapping[str, Mapping[str, str]]] | None = None,
     ) -> None:
         funcs_to_register = self._get_functions_to_register(functions)
 

autogen/mcp/mcp_proxy/operation_grouping.py

@@ -5,7 +5,6 @@
 import json
 import logging
 from pathlib import Path
-from typing import Dict, List
 
 from autogen.import_utils import optional_import_block
 
@@ -58,11 +57,11 @@ GROUP_ASSIGNMENT_MESSAGE = (
 )
 
 
-def chunk_list(items: List, size: int) -> List[List]:
+def chunk_list(items: list, size: int) -> list[list]:
     return [items[i : i + size] for i in range(0, len(items), size)]
 
 
-def discover_groups(operations: List["Operation"], chunk_size: int = 30) -> Dict[str, str]:
+def discover_groups(operations: list["Operation"], chunk_size: int = 30) -> dict[str, str]:
     llm_config = LLMConfig.get_current_llm_config().copy()
 
     for config in llm_config.config_list:
@@ -130,13 +129,13 @@ def assign_operation_to_group(operation: "Operation", groups: dict[str, str]) ->
         return groups
 
 
-def refine_group_names(groups: Dict[str, str]) -> Dict[str, str]:
+def refine_group_names(groups: dict[str, str]) -> dict[str, str]:
     # Optional: normalize names, merge similar ones (e.g., using embeddings or string similarity)
     # Placeholder for now:
     return groups
 
 
-def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> Dict[str, object]:
+def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> dict[str, object]:
     operations = sorted(parser.operations.values(), key=lambda op: op.path)
 
     # ---- PASS 1: DISCOVER GROUPS ----

autogen/mcp/mcp_proxy/operation_renaming.py

@@ -4,7 +4,6 @@
 
 import logging
 from pathlib import Path
-from typing import Dict, List
 
 from autogen.import_utils import optional_import_block
 
@@ -28,9 +27,8 @@ SYSTEM_MESSAGE = (
 )
 
 
-def validate_function_name(name: str, taken_names: List[str]) -> str:
-    """
-    Validate the generated function name against length, format, and uniqueness constraints.
+def validate_function_name(name: str, taken_names: list[str]) -> str:
+    """Validate the generated function name against length, format, and uniqueness constraints.
 
     Returns:
         'exit' if the name is valid, or an error message string otherwise.
@@ -44,9 +42,8 @@ def validate_function_name(name: str, taken_names: List[str]) -> str:
     return "exit"
 
 
-def get_new_function_name(operation: "Operation", taken_names: List[str]) -> str:
-    """
-    Ask an AI agent to generate a new function name for a given OpenAPI operation.
+def get_new_function_name(operation: "Operation", taken_names: list[str]) -> str:
+    """Ask an AI agent to generate a new function name for a given OpenAPI operation.
 
     Args:
         operation: The OpenAPI operation that needs renaming.
@@ -88,9 +85,8 @@ def get_new_function_name(operation: "Operation", taken_names: List[str]) -> str
     return response.summary
 
 
-def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> Dict[str, object]:
-    """
-    Visits and optionally renames operations in the OpenAPI parser.
+def custom_visitor(parser: "OpenAPIParser", model_path: Path) -> dict[str, object]:
+    """Visits and optionally renames operations in the OpenAPI parser.
 
     Args:
         parser: An OpenAPIParser instance containing API operations.

autogen/mcp/mcp_proxy/security.py

@@ -4,11 +4,10 @@
 import base64
 import json
 import logging
-from typing import Any, ClassVar, Literal, Optional
+from typing import Any, ClassVar, Literal, TypeAlias
 
 import requests
 from pydantic import BaseModel, model_validator
-from typing_extensions import TypeAlias
 
 # Get the logger
 logger = logging.getLogger(__name__)
@@ -343,7 +342,7 @@ class OAuth2PasswordBearer(BaseSecurity):
 
         username: str = "USERNAME"
         password: str = "PASSWORD"
-        bearer_token: Optional[str] = None
+        bearer_token: str | None = None
         token_url: str = "TOKEN_URL"
 
         # @model_validator(mode="before")