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

autogen/agentchat/contrib/capabilities/transforms.py

@@ -6,7 +6,7 @@
 # SPDX-License-Identifier: MIT
 import copy
 import sys
-from typing import Any, Optional, Protocol, Union
+from typing import Any, Protocol
 
 import tiktoken
 from termcolor import colored
@@ -62,9 +62,9 @@ class MessageHistoryLimiter:
 
     def __init__(
         self,
-        max_messages: Optional[int] = None,
+        max_messages: int | None = None,
         keep_first_message: bool = False,
-        exclude_names: Optional[list[str]] = None,
+        exclude_names: list[str] | None = None,
     ):
         """Args:
         max_messages Optional[int]: Maximum number of messages to keep in the context. Must be greater than 0 if not None.
@@ -91,7 +91,6 @@ class MessageHistoryLimiter:
         Returns:
             List[Dict]: A new list containing the most recent messages up to the specified maximum.
         """
-
         exclude_names = getattr(self, "_exclude_names", None)
 
         filtered = [msg for msg in messages if msg.get("name") not in exclude_names] if exclude_names else messages
@@ -136,7 +135,7 @@ class MessageHistoryLimiter:
             return logs_str, True
         return "No messages were removed.", False
 
-    def _validate_max_messages(self, max_messages: Optional[int]):
+    def _validate_max_messages(self, max_messages: int | None):
         if max_messages is not None and max_messages < 1:
             raise ValueError("max_messages must be None or greater than 1")
 
@@ -171,11 +170,11 @@ class MessageTokenLimiter:
 
     def __init__(
         self,
-        max_tokens_per_message: Optional[int] = None,
-        max_tokens: Optional[int] = None,
-        min_tokens: Optional[int] = None,
+        max_tokens_per_message: int | None = None,
+        max_tokens: int | None = None,
+        min_tokens: int | None = None,
         model: str = "gpt-3.5-turbo-0613",
-        filter_dict: Optional[dict[str, Any]] = None,
+        filter_dict: dict[str, Any] | None = None,
         exclude_filter: bool = True,
     ):
         """Args:
@@ -268,7 +267,7 @@ class MessageTokenLimiter:
             return logs_str, True
         return "No tokens were truncated.", False
 
-    def _truncate_str_to_tokens(self, contents: Union[str, list], n_tokens: int) -> Union[str, list]:
+    def _truncate_str_to_tokens(self, contents: str | list, n_tokens: int) -> str | list:
         if isinstance(contents, str):
             return self._truncate_tokens(contents, n_tokens)
         elif isinstance(contents, list):
@@ -296,7 +295,7 @@ class MessageTokenLimiter:
 
         return truncated_text
 
-    def _validate_max_tokens(self, max_tokens: Optional[int] = None) -> Optional[int]:
+    def _validate_max_tokens(self, max_tokens: int | None = None) -> int | None:
         if max_tokens is not None and max_tokens < 0:
             raise ValueError("max_tokens and max_tokens_per_message must be None or greater than or equal to 0")
 
@@ -317,7 +316,7 @@ class MessageTokenLimiter:
 
         return max_tokens if max_tokens is not None else sys.maxsize
 
-    def _validate_min_tokens(self, min_tokens: Optional[int], max_tokens: Optional[int]) -> int:
+    def _validate_min_tokens(self, min_tokens: int | None, max_tokens: int | None) -> int:
         if min_tokens is None:
             return 0
         if min_tokens < 0:
@@ -336,11 +335,11 @@ class TextMessageCompressor:
 
     def __init__(
         self,
-        text_compressor: Optional[TextCompressor] = None,
-        min_tokens: Optional[int] = None,
+        text_compressor: TextCompressor | None = None,
+        min_tokens: int | None = None,
         compression_params: dict = dict(),
-        cache: Optional[AbstractCache] = None,
-        filter_dict: Optional[dict[str, Any]] = None,
+        cache: AbstractCache | None = None,
+        filter_dict: dict[str, Any] | None = None,
         exclude_filter: bool = True,
     ):
         """Args:
@@ -466,7 +465,7 @@ class TextMessageCompressor:
 
         return compressed_text["compressed_prompt"], savings
 
-    def _validate_min_tokens(self, min_tokens: Optional[int]):
+    def _validate_min_tokens(self, min_tokens: int | None):
         if min_tokens is not None and min_tokens <= 0:
             raise ValueError("min_tokens must be greater than 0 or None")
 
@@ -497,7 +496,7 @@ class TextMessageContentName:
         position: str = "start",
         format_string: str = "{name}:\n",
         deduplicate: bool = True,
-        filter_dict: Optional[dict[str, Any]] = None,
+        filter_dict: dict[str, Any] | None = None,
         exclude_filter: bool = True,
     ):
         """Args:

autogen/agentchat/contrib/capabilities/transforms_util.py

@@ -5,7 +5,7 @@
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
 from collections.abc import Hashable
-from typing import Any, Optional
+from typing import Any
 
 from .... import token_count_utils
 from ....cache.abstract_cache_base import AbstractCache
@@ -24,7 +24,7 @@ def cache_key(content: MessageContentType, *args: Hashable) -> str:
     return "".join(str_keys)
 
 
-def cache_content_get(cache: Optional[AbstractCache], key: str) -> Optional[tuple[MessageContentType, ...]]:
+def cache_content_get(cache: AbstractCache | None, key: str) -> tuple[MessageContentType, ...] | None:
     """Retrieves cached content from the cache.
 
     Args:
@@ -37,7 +37,7 @@ def cache_content_get(cache: Optional[AbstractCache], key: str) -> Optional[tupl
             return cached_value
 
 
-def cache_content_set(cache: Optional[AbstractCache], key: str, content: MessageContentType, *extra_values):
+def cache_content_set(cache: AbstractCache | None, key: str, content: MessageContentType, *extra_values):
     """Sets content into the cache.
 
     Args:
@@ -51,7 +51,7 @@ def cache_content_set(cache: Optional[AbstractCache], key: str, content: Message
         cache.set(key, cache_value)
 
 
-def min_tokens_reached(messages: list[dict[str, Any]], min_tokens: Optional[int]) -> bool:
+def min_tokens_reached(messages: list[dict[str, Any]], min_tokens: int | None) -> bool:
     """Returns True if the total number of tokens in the messages is greater than or equal to the specified value.
 
     Args:
@@ -108,7 +108,7 @@ def is_content_text_empty(content: MessageContentType) -> bool:
         return True
 
 
-def should_transform_message(message: dict[str, Any], filter_dict: Optional[dict[str, Any]], exclude: bool) -> bool:
+def should_transform_message(message: dict[str, Any], filter_dict: dict[str, Any] | None, exclude: bool) -> bool:
     """Validates whether the transform should be applied according to the filter dictionary.
 
     Args:

autogen/agentchat/contrib/capabilities/vision_capability.py

@@ -5,7 +5,8 @@
 # Portions derived from  https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
 import copy
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any
 
 from ....code_utils import content_str
 from ....oai.client import OpenAIWrapper
@@ -47,7 +48,7 @@ class VisionCapability(AgentCapability):
     def __init__(
         self,
         lmm_config: dict[str, Any],
-        description_prompt: Optional[str] = DEFAULT_DESCRIPTION_PROMPT,
+        description_prompt: str | None = DEFAULT_DESCRIPTION_PROMPT,
         custom_caption_func: Callable = None,
     ) -> None:
         """Initializes a new instance, setting up the configuration for interacting with
@@ -101,7 +102,7 @@ class VisionCapability(AgentCapability):
         # Register a hook for processing the last message.
         agent.register_hook(hookable_method="process_last_received_message", hook=self.process_last_received_message)
 
-    def process_last_received_message(self, content: Union[str, list[dict[str, Any]]]) -> str:
+    def process_last_received_message(self, content: str | list[dict[str, Any]]) -> str:
         """Processes the last received message content by normalizing and augmenting it
         with descriptions of any included images. The function supports input content
         as either a string or a list of dictionaries, where each dictionary represents

autogen/agentchat/contrib/captainagent/agent_builder.py

@@ -11,7 +11,7 @@ import logging
 import re
 import subprocess as sp
 import time
-from typing import Any, Optional, Union
+from typing import Any
 
 from termcolor import colored
 
@@ -184,14 +184,14 @@ Match roles in the role set to each expert in expert set.
 
     def __init__(
         self,
-        config_file_or_env: Optional[str] = "OAI_CONFIG_LIST",
-        config_file_location: Optional[str] = "",
-        llm_config: Optional[Union[LLMConfig, dict[str, Any]]] = None,
-        builder_model: Optional[Union[str, list]] = [],
-        agent_model: Optional[Union[str, list]] = [],
-        builder_model_tags: Optional[list] = [],
-        agent_model_tags: Optional[list] = [],
-        max_agents: Optional[int] = 5,
+        config_file_or_env: str | None = "OAI_CONFIG_LIST",
+        config_file_location: str | None = "",
+        llm_config: LLMConfig | dict[str, Any] | None = None,
+        builder_model: str | list | None = [],
+        agent_model: str | list | None = [],
+        builder_model_tags: list | None = [],
+        agent_model_tags: list | None = [],
+        max_agents: int | None = 5,
     ):
         """(These APIs are experimental and may change in the future.)
 
@@ -259,8 +259,8 @@ Match roles in the role set to each expert in expert set.
         self,
         agent_config: dict[str, Any],
         member_name: list[str],
-        llm_config: Union[LLMConfig, dict[str, Any]],
-        use_oai_assistant: Optional[bool] = False,
+        llm_config: LLMConfig | dict[str, Any],
+        use_oai_assistant: bool | None = False,
     ) -> AssistantAgent:
         """Create a group chat participant agent.
 
@@ -357,7 +357,7 @@ Match roles in the role set to each expert in expert set.
         self.agent_procs_assign[agent_name] = (agent, server_id)
         return agent
 
-    def clear_agent(self, agent_name: str, recycle_endpoint: Optional[bool] = True):
+    def clear_agent(self, agent_name: str, recycle_endpoint: bool | None = True):
         """Clear a specific agent by name.
 
         Args:
@@ -378,7 +378,7 @@ Match roles in the role set to each expert in expert set.
                 self.open_ports.append(server_id.split("_")[-1])
         print(colored(f"Agent {agent_name} has been cleared.", "yellow"), flush=True)
 
-    def clear_all_agents(self, recycle_endpoint: Optional[bool] = True):
+    def clear_all_agents(self, recycle_endpoint: bool | None = True):
         """Clear all cached agents."""
         for agent_name in [agent_name for agent_name in self.agent_procs_assign]:
             self.clear_agent(agent_name, recycle_endpoint)
@@ -387,12 +387,12 @@ Match roles in the role set to each expert in expert set.
     def build(
         self,
         building_task: str,
-        default_llm_config: Union[LLMConfig, dict[str, Any]],
-        coding: Optional[bool] = None,
-        code_execution_config: Optional[dict[str, Any]] = None,
-        use_oai_assistant: Optional[bool] = False,
-        user_proxy: Optional[ConversableAgent] = None,
-        max_agents: Optional[int] = None,
+        default_llm_config: LLMConfig | dict[str, Any],
+        coding: bool | None = None,
+        code_execution_config: dict[str, Any] | None = None,
+        use_oai_assistant: bool | None = False,
+        user_proxy: ConversableAgent | None = None,
+        max_agents: int | None = None,
         **kwargs: Any,
     ) -> tuple[list[ConversableAgent], dict[str, Any]]:
         """Auto build agents based on the building task.
@@ -515,13 +515,13 @@ Match roles in the role set to each expert in expert set.
         self,
         building_task: str,
         library_path_or_json: str,
-        default_llm_config: Union[LLMConfig, dict[str, Any]],
+        default_llm_config: LLMConfig | dict[str, Any],
         top_k: int = 3,
-        coding: Optional[bool] = None,
-        code_execution_config: Optional[dict[str, Any]] = None,
-        use_oai_assistant: Optional[bool] = False,
-        embedding_model: Optional[str] = "all-mpnet-base-v2",
-        user_proxy: Optional[ConversableAgent] = None,
+        coding: bool | None = None,
+        code_execution_config: dict[str, Any] | None = None,
+        use_oai_assistant: bool | None = False,
+        embedding_model: str | None = "all-mpnet-base-v2",
+        user_proxy: ConversableAgent | None = None,
         **kwargs: Any,
     ) -> tuple[list[ConversableAgent], dict[str, Any]]:
         """Build agents from a library.
@@ -668,7 +668,7 @@ Match roles in the role set to each expert in expert set.
         return self._build_agents(use_oai_assistant, user_proxy=user_proxy, **kwargs)
 
     def _build_agents(
-        self, use_oai_assistant: Optional[bool] = False, user_proxy: Optional[ConversableAgent] = None, **kwargs
+        self, use_oai_assistant: bool | None = False, user_proxy: ConversableAgent | None = None, **kwargs
     ) -> tuple[list[ConversableAgent], dict[str, Any]]:
         """Build agents with generated configs.
 
@@ -711,7 +711,7 @@ Match roles in the role set to each expert in expert set.
 
         return agent_list, self.cached_configs.copy()
 
-    def save(self, filepath: Optional[str] = None) -> str:
+    def save(self, filepath: str | None = None) -> str:
         """Save building configs. If the filepath is not specific, this function will create a filename by encrypt the
         building_task string by md5 with "save_config_" prefix, and save config to the local path.
 
@@ -731,9 +731,9 @@ Match roles in the role set to each expert in expert set.
 
     def load(
         self,
-        filepath: Optional[str] = None,
-        config_json: Optional[str] = None,
-        use_oai_assistant: Optional[bool] = False,
+        filepath: str | None = None,
+        config_json: str | None = None,
+        use_oai_assistant: bool | None = False,
         **kwargs: Any,
     ) -> tuple[list[ConversableAgent], dict[str, Any]]:
         """Load building configs and call the build function to complete building without calling online LLMs' api.

autogen/agentchat/contrib/captainagent/captainagent.py

@@ -4,7 +4,8 @@
 import hashlib
 import json
 import os
-from typing import Any, Callable, Literal, Optional, Union
+from collections.abc import Callable
+from typing import Any, Literal
 
 from termcolor import colored
 
@@ -135,17 +136,17 @@ Note that the previous experts will forget everything after you obtain the respo
     def __init__(
         self,
         name: str,
-        system_message: Optional[str] = None,
-        llm_config: Optional[Union[LLMConfig, dict[str, Any], Literal[False]]] = None,
-        is_termination_msg: Optional[Callable[[dict[str, Any]], bool]] = None,
-        max_consecutive_auto_reply: Optional[int] = None,
-        human_input_mode: Optional[str] = "NEVER",
-        code_execution_config: Optional[Union[dict[str, Any], Literal[False]]] = False,
-        nested_config: Optional[dict[str, Any]] = None,
-        agent_lib: Optional[str] = None,
-        tool_lib: Optional[str] = None,
-        agent_config_save_path: Optional[str] = None,
-        description: Optional[str] = DEFAULT_DESCRIPTION,
+        system_message: str | None = None,
+        llm_config: LLMConfig | dict[str, Any] | Literal[False] | None = None,
+        is_termination_msg: Callable[[dict[str, Any]], bool] | None = None,
+        max_consecutive_auto_reply: int | None = None,
+        human_input_mode: str | None = "NEVER",
+        code_execution_config: dict[str, Any] | Literal[False] | None = False,
+        nested_config: dict[str, Any] | None = None,
+        agent_lib: str | None = None,
+        tool_lib: str | None = None,
+        agent_config_save_path: str | None = None,
+        description: str | None = DEFAULT_DESCRIPTION,
         **kwargs: Any,
     ):
         """Args:
@@ -227,7 +228,7 @@ Note that the previous experts will forget everything after you obtain the respo
         )
 
     @staticmethod
-    def _update_config(default_dict: dict[str, Any], update_dict: Optional[dict[str, Any]]) -> dict[str, Any]:
+    def _update_config(default_dict: dict[str, Any], update_dict: dict[str, Any] | None) -> dict[str, Any]:
         """Recursively updates the default_dict with values from update_dict."""
         if update_dict is None:
             return default_dict
@@ -297,14 +298,14 @@ Collect information from the general task, follow the suggestions from manager t
         name: str,
         nested_config: dict[str, Any],
         agent_config_save_path: str = None,
-        is_termination_msg: Optional[Callable[[dict[str, Any]], bool]] = None,
-        max_consecutive_auto_reply: Optional[int] = None,
-        human_input_mode: Optional[str] = "NEVER",
-        code_execution_config: Optional[Union[dict[str, Any], Literal[False]]] = None,
-        default_auto_reply: Optional[Union[str, dict[str, Any]]] = DEFAULT_AUTO_REPLY,
-        llm_config: Optional[Union[LLMConfig, dict[str, Any], Literal[False]]] = False,
-        system_message: Optional[Union[str, list]] = "",
-        description: Optional[str] = None,
+        is_termination_msg: Callable[[dict[str, Any]], bool] | None = None,
+        max_consecutive_auto_reply: int | None = None,
+        human_input_mode: str | None = "NEVER",
+        code_execution_config: dict[str, Any] | Literal[False] | None = None,
+        default_auto_reply: str | dict[str, Any] | None = DEFAULT_AUTO_REPLY,
+        llm_config: LLMConfig | dict[str, Any] | Literal[False] | None = False,
+        system_message: str | list | None = "",
+        description: str | None = None,
     ):
         """Args:
         name (str): name of the agent.

autogen/agentchat/contrib/captainagent/tool_retriever.py

@@ -14,7 +14,6 @@ import traceback
 from hashlib import md5
 from pathlib import Path
 from textwrap import dedent, indent
-from typing import Optional, Union
 
 from .... import AssistantAgent, UserProxyAgent
 from ....coding import CodeExecutor, CodeExtractor, LocalCommandLineCodeExecutor, MarkdownCodeExtractor
@@ -76,7 +75,7 @@ You have access to the following functions. You can write python code to call th
         agent.update_system_message(sys_message)
         return
 
-    def bind_user_proxy(self, agent: UserProxyAgent, tool_root: Union[str, list]):
+    def bind_user_proxy(self, agent: UserProxyAgent, tool_root: str | list):
         """Updates user proxy agent with a executor so that code executor can successfully execute function-related code.
         Returns an updated user proxy.
         """
@@ -164,7 +163,7 @@ class LocalExecutorWithTools(CodeExecutor):
         """(Experimental) Export a code extractor that can be used by an agent."""
         return MarkdownCodeExtractor()
 
-    def __init__(self, tools: Optional[list[Tool]] = None, work_dir: Union[Path, str] = Path()):
+    def __init__(self, tools: list[Tool] | None = None, work_dir: Path | str = Path()):
         self.tools = tools if tools is not None else []
         self.work_dir = work_dir
         if not os.path.exists(work_dir):

autogen/agentchat/contrib/gpt_assistant_agent.py

@@ -9,7 +9,7 @@ import json
 import logging
 import time
 from collections import defaultdict
-from typing import Any, Optional, Union
+from typing import Any
 
 from ... import OpenAIWrapper
 from ...llm_config import LLMConfig
@@ -31,9 +31,9 @@ class GPTAssistantAgent(ConversableAgent):
     def __init__(
         self,
         name="GPT Assistant",
-        instructions: Optional[str] = None,
-        llm_config: Optional[Union[LLMConfig, dict[str, Any], bool]] = None,
-        assistant_config: Optional[dict[str, Any]] = None,
+        instructions: str | None = None,
+        llm_config: LLMConfig | dict[str, Any] | bool | None = None,
+        assistant_config: dict[str, Any] | None = None,
         overwrite_instructions: bool = False,
         overwrite_tools: bool = False,
         **kwargs: Any,
@@ -182,10 +182,10 @@ class GPTAssistantAgent(ConversableAgent):
 
     def _invoke_assistant(
         self,
-        messages: Optional[list[dict[str, Any]]] = None,
-        sender: Optional[Agent] = None,
-        config: Optional[Any] = None,
-    ) -> tuple[bool, Optional[Union[str, dict[str, Any]]]]:
+        messages: list[dict[str, Any]] | None = None,
+        sender: Agent | None = None,
+        config: Any | None = None,
+    ) -> tuple[bool, str | dict[str, Any] | None]:
         """Invokes the OpenAI assistant to generate a reply based on the given messages.
 
         Args:
@@ -392,7 +392,7 @@ class GPTAssistantAgent(ConversableAgent):
         # Clear the record of unread messages
         self._unread_index.clear()
 
-    def clear_history(self, agent: Optional[Agent] = None):
+    def clear_history(self, agent: Agent | None = None):
         """Clear the chat history of the agent.
 
         Args:

autogen/agentchat/contrib/graph_rag/document.py

@@ -6,7 +6,7 @@
 # SPDX-License-Identifier: MIT
 from dataclasses import dataclass, field
 from enum import Enum, auto
-from typing import Any, Optional
+from typing import Any
 
 __all__ = ["Document", "DocumentType"]
 
@@ -25,5 +25,5 @@ class Document:
     """A wrapper of graph store query results."""
 
     doctype: DocumentType
-    data: Optional[Any] = None
-    path_or_url: Optional[str] = field(default_factory=lambda: "")
+    data: Any | None = None
+    path_or_url: str | None = field(default_factory=lambda: "")

autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py

@@ -28,8 +28,8 @@ class FalkorGraphQueryEngine:
         name: str,
         host: str = "127.0.0.1",
         port: int = 6379,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
+        username: str | None = None,
+        password: str | None = None,
         model: Optional["GenerativeModel"] = None,
         ontology: Optional["Ontology"] = None,
     ):
@@ -57,7 +57,7 @@ class FalkorGraphQueryEngine:
         self.model = model or OpenAiGenerativeModel("gpt-4o")
         self.model_config = KnowledgeGraphModelConfig.with_model(model)
         self.ontology = ontology
-        self.knowledge_graph: Optional["KnowledgeGraph"] = None  # type: ignore[no-any-unimported]
+        self.knowledge_graph: KnowledgeGraph | None = None  # type: ignore[no-any-unimported]
         self.falkordb = FalkorDB(host=self.host, port=self.port, username=self.username, password=self.password)
 
     def connect_db(self) -> None:

autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Optional, Union
+from typing import Any
 
 from .... import Agent, ConversableAgent
 from .falkor_graph_query_engine import FalkorGraphQueryEngine
@@ -50,10 +50,10 @@ class FalkorGraphRagCapability(GraphRagCapability):
     def _reply_using_falkordb_query(
         self,
         recipient: ConversableAgent,
-        messages: Optional[list[dict[str, Any]]] = None,
-        sender: Optional[Agent] = None,
-        config: Optional[Any] = None,
-    ) -> tuple[bool, Optional[Union[str, dict[str, Any]]]]:
+        messages: list[dict[str, Any]] | None = None,
+        sender: Agent | None = None,
+        config: Any | None = None,
+    ) -> tuple[bool, str | dict[str, Any] | None]:
         """Query FalkorDB and return the message. Internally, it utilises OpenAI to generate a reply based on the given messages.
         The history with FalkorDB is also logged and updated.
 
@@ -76,7 +76,7 @@ class FalkorGraphRagCapability(GraphRagCapability):
 
         return True, result.answer if result.answer else "I'm sorry, I don't have an answer for that."
 
-    def _messages_summary(self, messages: Union[dict[str, Any], str], system_message: str) -> str:
+    def _messages_summary(self, messages: dict[str, Any] | str, system_message: str) -> str:
         """Summarize the messages in the conversation history. Excluding any message with 'tool_calls' and 'tool_responses'
         Includes the 'name' (if it exists) and the 'content', with a new line between each one, like:
         customer:

autogen/agentchat/contrib/graph_rag/graph_query_engine.py

@@ -5,7 +5,7 @@
 # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
 # SPDX-License-Identifier: MIT
 from dataclasses import dataclass, field
-from typing import Any, Optional, Protocol, runtime_checkable
+from typing import Any, Protocol, runtime_checkable
 
 from .document import Document
 
@@ -20,7 +20,7 @@ class GraphStoreQueryResult:
     results: intermediate results to question/query, e.g. node entities.
     """
 
-    answer: Optional[str] = None
+    answer: str | None = None
     results: list[Any] = field(default_factory=list)
 
 
@@ -31,7 +31,7 @@ class GraphQueryEngine(Protocol):
     This interface defines the basic methods for graph-based RAG.
     """
 
-    def init_db(self, input_doc: Optional[list[Document]] = None) -> None:
+    def init_db(self, input_doc: list[Document] | None = None) -> None:
         """This method initializes graph database with the input documents or records.
         Usually, it takes the following steps,
         1. connecting to a graph database.

autogen/agentchat/contrib/graph_rag/neo4j_graph_query_engine.py

@@ -2,13 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 import os
-import sys
-from typing import Any, Optional, Union
-
-if sys.version_info >= (3, 10):
-    from typing import TypeAlias
-else:
-    from typing_extensions import TypeAlias
+from typing import Any, Optional, TypeAlias
 
 from ....import_utils import optional_import_block, require_optional_import
 from .document import Document, DocumentType
@@ -64,8 +58,8 @@ class Neo4jGraphQueryEngine:
         embedding: Optional["BaseEmbedding"] = None,
         entities: Optional["TypeAlias"] = None,
         relations: Optional["TypeAlias"] = None,
-        schema: Optional[Union[dict[str, str], list["Triple"]]] = None,
-        strict: Optional[bool] = False,
+        schema: dict[str, str] | list["Triple"] | None = None,
+        strict: bool | None = False,
     ):
         """Initialize a Neo4j Property graph.
         Please also refer to https://docs.llamaindex.ai/en/stable/examples/property_graph/graph_store/
@@ -96,7 +90,7 @@ class Neo4jGraphQueryEngine:
         self.schema = schema
         self.strict = strict
 
-    def init_db(self, input_doc: Optional[list[Document]] = None) -> None:
+    def init_db(self, input_doc: list[Document] | None = None) -> None:
         """Build the knowledge graph with input documents."""
         self.documents = self._load_doc(input_doc if input_doc is not None else [])
 
@@ -245,7 +239,7 @@ class Neo4jGraphQueryEngine:
         # To add more extractors, please refer to https://docs.llamaindex.ai/en/latest/module_guides/indexing/lpg_index_guide/#construction
         """
         #
-        kg_extractors: list["TransformComponent"] = [  # type: ignore[no-any-unimported]
+        kg_extractors: list[TransformComponent] = [  # type: ignore[no-any-unimported]
             SchemaLLMPathExtractor(
                 llm=self.llm,
                 possible_entities=self.entities,

autogen/agentchat/contrib/graph_rag/neo4j_graph_rag_capability.py

@@ -2,7 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from typing import Any, Optional, Union
+from typing import Any
 
 from .... import Agent, ConversableAgent, UserProxyAgent
 from .graph_query_engine import GraphStoreQueryResult
@@ -46,10 +46,10 @@ class Neo4jGraphCapability(GraphRagCapability):
     def _reply_using_neo4j_query(
         self,
         recipient: ConversableAgent,
-        messages: Optional[list[dict[str, Any]]] = None,
-        sender: Optional[Agent] = None,
-        config: Optional[Any] = None,
-    ) -> tuple[bool, Optional[Union[str, dict[str, Any]]]]:
+        messages: list[dict[str, Any]] | None = None,
+        sender: Agent | None = None,
+        config: Any | None = None,
+    ) -> tuple[bool, str | dict[str, Any] | None]:
         """Query neo4j and return the message. Internally, it queries the Property graph
         and returns the answer from the graph query engine.
         TODO: reply with a dictionary including both the answer and semantic source triplets.
@@ -74,7 +74,7 @@ class Neo4jGraphCapability(GraphRagCapability):
 
         return True, result.answer
 
-    def _get_last_question(self, message: Union[dict[str, Any], str]) -> Optional[Union[str, dict[str, Any]]]:
+    def _get_last_question(self, message: dict[str, Any] | str) -> str | dict[str, Any] | None:
         """Retrieves the last message from the conversation history."""
         if isinstance(message, str):
             return message