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

autogen/agentchat/contrib/vectordb/pgvectordb.py

@@ -7,7 +7,8 @@
 import os
 import re
 import urllib.parse
-from typing import Any, Callable, Optional, Union
+from collections.abc import Callable
+from typing import Any, Optional
 
 from ....import_utils import optional_import_block, require_optional_import
 from .base import Document, ItemID, QueryResults, VectorDB
@@ -41,11 +42,11 @@ class Collection:
 
     def __init__(
         self,
-        client: Optional[Any] = None,
+        client: Any | None = None,
         collection_name: str = "ag2-docs",
-        embedding_function: Optional[Callable[..., Any]] = None,
-        metadata: Optional[Any] = None,
-        get_or_create: Optional[Any] = None,
+        embedding_function: Callable[..., Any] | None = None,
+        metadata: Any | None = None,
+        get_or_create: Any | None = None,
     ):
         """Initialize the Collection object.
 
@@ -85,9 +86,9 @@ class Collection:
     def add(
         self,
         ids: list[ItemID],
-        documents: Optional[list[Document]],
-        embeddings: Optional[list[Any]] = None,
-        metadatas: Optional[list[Any]] = None,
+        documents: list[Document] | None,
+        embeddings: list[Any] | None = None,
+        metadatas: list[Any] | None = None,
     ) -> None:
         """Add documents to the collection.
 
@@ -130,8 +131,8 @@ class Collection:
         self,
         ids: list[ItemID],
         documents: list[Document],
-        embeddings: Optional[list[Any]] = None,
-        metadatas: Optional[list[Any]] = None,
+        embeddings: list[Any] | None = None,
+        metadatas: list[Any] | None = None,
     ) -> None:
         """Upsert documents into the collection.
 
@@ -232,11 +233,11 @@ class Collection:
 
     def get(
         self,
-        ids: Optional[str] = None,
-        include: Optional[str] = None,
-        where: Optional[str] = None,
-        limit: Optional[Union[int, str]] = None,
-        offset: Optional[Union[int, str]] = None,
+        ids: str | None = None,
+        include: str | None = None,
+        where: str | None = None,
+        limit: int | str | None = None,
+        offset: int | str | None = None,
     ) -> list[Document]:
         """Retrieve documents from the collection.
 
@@ -380,11 +381,11 @@ class Collection:
     def query(
         self,
         query_texts: list[str],
-        collection_name: Optional[str] = None,
-        n_results: Optional[int] = 10,
-        distance_type: Optional[str] = "euclidean",
-        distance_threshold: Optional[float] = -1,
-        include_embedding: Optional[bool] = False,
+        collection_name: str | None = None,
+        n_results: int | None = 10,
+        distance_type: str | None = "euclidean",
+        distance_threshold: float | None = -1,
+        include_embedding: bool | None = False,
     ) -> QueryResults:
         """Query documents in the collection.
 
@@ -465,7 +466,7 @@ class Collection:
         array = [float(num) for num in array_string.split()]
         return array
 
-    def modify(self, metadata, collection_name: Optional[str] = None) -> None:
+    def modify(self, metadata, collection_name: str | None = None) -> None:
         """Modify metadata for the collection.
 
         Args:
@@ -481,7 +482,7 @@ class Collection:
         cursor.execute("UPDATE collectionsSET metadata = '%s'WHERE collection_name = '%s';", (metadata, self.name))
         cursor.close()
 
-    def delete(self, ids: list[ItemID], collection_name: Optional[str] = None) -> None:
+    def delete(self, ids: list[ItemID], collection_name: str | None = None) -> None:
         """Delete documents from the collection.
 
         Args:
@@ -498,7 +499,7 @@ class Collection:
         cursor.execute(f"DELETE FROM {self.name} WHERE id IN ({id_placeholders});", ids)
         cursor.close()
 
-    def delete_collection(self, collection_name: Optional[str] = None) -> None:
+    def delete_collection(self, collection_name: str | None = None) -> None:
         """Delete the entire collection.
 
         Args:
@@ -513,9 +514,7 @@ class Collection:
         cursor.execute(f"DROP TABLE IF EXISTS {self.name}")
         cursor.close()
 
-    def create_collection(
-        self, collection_name: Optional[str] = None, dimension: Optional[Union[str, int]] = None
-    ) -> None:
+    def create_collection(self, collection_name: str | None = None, dimension: str | int | None = None) -> None:
         """Create a new collection.
 
         Args:
@@ -558,15 +557,15 @@ class PGVectorDB(VectorDB):
         self,
         *,
         conn: Optional["psycopg.Connection"] = None,
-        connection_string: Optional[str] = None,
-        host: Optional[str] = None,
-        port: Optional[Union[int, str]] = None,
-        dbname: Optional[str] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        connect_timeout: Optional[int] = 10,
+        connection_string: str | None = None,
+        host: str | None = None,
+        port: int | str | None = None,
+        dbname: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        connect_timeout: int | None = 10,
         embedding_function: Callable = None,
-        metadata: Optional[dict[str, Any]] = None,
+        metadata: dict[str, Any] | None = None,
     ) -> None:
         """Initialize the vector database.
 
@@ -615,13 +614,13 @@ class PGVectorDB(VectorDB):
     def establish_connection(
         self,
         conn: Optional["psycopg.Connection"] = None,
-        connection_string: Optional[str] = None,
-        host: Optional[str] = None,
-        port: Optional[Union[int, str]] = None,
-        dbname: Optional[str] = None,
-        username: Optional[str] = None,
-        password: Optional[str] = None,
-        connect_timeout: Optional[int] = 10,
+        connection_string: str | None = None,
+        host: str | None = None,
+        port: int | str | None = None,
+        dbname: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+        connect_timeout: int | None = 10,
     ) -> "psycopg.Connection":
         """Establishes a connection to a PostgreSQL database using psycopg.
 

autogen/agentchat/contrib/vectordb/qdrant.py

@@ -7,7 +7,7 @@
 import abc
 import logging
 from collections.abc import Sequence
-from typing import Any, Optional, Union
+from typing import Any
 
 from ....import_utils import optional_import_block, require_optional_import
 from .base import Document, ItemID, QueryResults, VectorDB
@@ -20,7 +20,7 @@ with optional_import_block():
 
 logger = get_logger(__name__)
 
-Embeddings = Union[Sequence[float], Sequence[int]]
+Embeddings = Sequence[float] | Sequence[int]
 
 
 class EmbeddingFunction(abc.ABC):
@@ -37,9 +37,9 @@ class FastEmbedEmbeddingFunction(EmbeddingFunction):
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
         batch_size: int = 256,
-        cache_dir: Optional[str] = None,
-        threads: Optional[int] = None,
-        parallel: Optional[int] = None,
+        cache_dir: str | None = None,
+        threads: int | None = None,
+        parallel: int | None = None,
         **kwargs: Any,
     ):
         """Initialize fastembed.TextEmbedding.

autogen/agentchat/contrib/web_surfer.py

@@ -7,8 +7,9 @@
 import copy
 import logging
 import re
+from collections.abc import Callable
 from datetime import datetime
-from typing import Annotated, Any, Callable, Literal, Optional, Union
+from typing import Annotated, Any, Literal
 
 from ... import Agent, AssistantAgent, ConversableAgent, OpenAIWrapper, UserProxyAgent
 from ...browser_utils import SimpleTextBrowser
@@ -32,17 +33,17 @@ class WebSurferAgent(ConversableAgent):
     def __init__(
         self,
         name: str,
-        system_message: Optional[Union[str, list[str]]] = DEFAULT_PROMPT,
-        description: Optional[str] = DEFAULT_DESCRIPTION,
-        is_termination_msg: Optional[Callable[[dict[str, Any]], bool]] = None,
-        max_consecutive_auto_reply: Optional[int] = None,
+        system_message: str | list[str] | None = DEFAULT_PROMPT,
+        description: str | None = DEFAULT_DESCRIPTION,
+        is_termination_msg: Callable[[dict[str, Any]], bool] | None = None,
+        max_consecutive_auto_reply: int | None = None,
         human_input_mode: Literal["ALWAYS", "NEVER", "TERMINATE"] = "TERMINATE",
-        function_map: Optional[dict[str, Callable[..., Any]]] = None,
-        code_execution_config: Union[dict[str, Any], Literal[False]] = False,
-        llm_config: Optional[Union[LLMConfig, dict[str, Any], Literal[False]]] = None,
-        summarizer_llm_config: Optional[Union[LLMConfig, dict[str, Any], Literal[False]]] = None,
-        default_auto_reply: Optional[Union[str, dict[str, Any]]] = "",
-        browser_config: Optional[dict[str, Any]] = None,
+        function_map: dict[str, Callable[..., Any]] | None = None,
+        code_execution_config: dict[str, Any] | Literal[False] = False,
+        llm_config: LLMConfig | dict[str, Any] | Literal[False] | None = None,
+        summarizer_llm_config: LLMConfig | dict[str, Any] | Literal[False] | None = None,
+        default_auto_reply: str | dict[str, Any] | None = "",
+        browser_config: dict[str, Any] | None = None,
         **kwargs: Any,
     ):
         super().__init__(
@@ -91,7 +92,7 @@ class WebSurferAgent(ConversableAgent):
         self.register_reply([Agent, None], ConversableAgent.check_termination_and_human_reply)
 
     def _create_summarizer_client(
-        self, summarizer_llm_config: Union[LLMConfig, dict[str, Any]], llm_config: Union[LLMConfig, dict[str, Any]]
+        self, summarizer_llm_config: LLMConfig | dict[str, Any], llm_config: LLMConfig | dict[str, Any]
     ) -> None:
         # If the summarizer_llm_config is None, we copy it from the llm_config
         if summarizer_llm_config is None:
@@ -200,8 +201,8 @@ class WebSurferAgent(ConversableAgent):
                 description="Uses AI to read the page and directly answer a given question based on the content.",
             )
             def _answer_from_page(
-                question: Annotated[Optional[str], "The question to directly answer."],
-                url: Annotated[Optional[str], "[Optional] The url of the page. (Defaults to the current page)"] = None,
+                question: Annotated[str | None, "The question to directly answer."],
+                url: Annotated[str | None, "[Optional] The url of the page. (Defaults to the current page)"] = None,
             ) -> str:
                 if url is not None and url != self.browser.address:
                     self.browser.visit_page(url)
@@ -257,17 +258,17 @@ class WebSurferAgent(ConversableAgent):
             )
             def _summarize_page(
                 url: Annotated[
-                    Optional[str], "[Optional] The url of the page to summarize. (Defaults to current page)"
+                    str | None, "[Optional] The url of the page to summarize. (Defaults to current page)"
                 ] = None,
             ) -> str:
                 return _answer_from_page(url=url, question=None)
 
     def generate_surfer_reply(
         self,
-        messages: Optional[list[dict[str, str]]] = None,
-        sender: Optional[Agent] = None,
-        config: Optional[OpenAIWrapper] = None,
-    ) -> tuple[bool, Optional[Union[str, dict[str, str]]]]:
+        messages: list[dict[str, str]] | None = None,
+        sender: Agent | None = None,
+        config: OpenAIWrapper | None = None,
+    ) -> tuple[bool, str | dict[str, str] | None]:
         """Generate a reply using autogen.oai."""
         if messages is None:
             messages = self._oai_messages[sender]