Enhanced Type Hints for Improved Readability and Maintainability: Addressing Type Mismatch in query_quadrant Return Statement

autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py

@@ -103,7 +103,7 @@ def __init__(
         self._hnsw_config = self._retrieve_config.get("hnsw_config", None)
         self._payload_indexing = self._retrieve_config.get("payload_indexing", False)
 
-    def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
+    def retrieve_docs(self, problem: str, n_results: Optional[int] = 20, search_string: Optional[str] = ""):
         """
         Args:
             problem (str): the problem to be solved.
@@ -150,21 +150,21 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =
 
 def create_qdrant_from_dir(
     dir_path: str,
-    max_tokens: int = 4000,
-    client: QdrantClient = None,
-    collection_name: str = "all-my-documents",
-    chunk_mode: str = "multi_lines",
-    must_break_at_empty_line: bool = True,
-    embedding_model: str = "BAAI/bge-small-en-v1.5",
-    custom_text_split_function: Callable = None,
-    custom_text_types: List[str] = TEXT_FORMATS,
-    recursive: bool = True,
-    extra_docs: bool = False,
-    parallel: int = 0,
-    on_disk: bool = False,
+    max_tokens: Optional[int] = 4000,
+    client: Optional[QdrantClient] = None,
+    collection_name: Optional[str] = "all-my-documents",
+    chunk_mode: Optional[str] = "multi_lines",
+    must_break_at_empty_line: Optional[bool] = True,
+    embedding_model: Optional[str] = "BAAI/bge-small-en-v1.5",
+    custom_text_split_function: Optional[Callable] = None,
+    custom_text_types: Optional[List[str]] = TEXT_FORMATS,
+    recursive: Optional[bool] = True,
+    extra_docs: Optional[bool] = False,
+    parallel: Optional[int] = 0,
+    on_disk: Optional[bool] = False,
     quantization_config: Optional[models.QuantizationConfig] = None,
     hnsw_config: Optional[models.HnswConfigDiff] = None,
-    payload_indexing: bool = False,
+    payload_indexing: Optional[bool] = False,
     qdrant_client_options: Optional[Dict] = {},
 ):
     """Create a Qdrant collection from all the files in a given directory, the directory can also be a single file or a
@@ -255,11 +255,11 @@ def create_qdrant_from_dir(
 
 def query_qdrant(
     query_texts: List[str],
-    n_results: int = 10,
-    client: QdrantClient = None,
-    collection_name: str = "all-my-documents",
-    search_string: str = "",
-    embedding_model: str = "BAAI/bge-small-en-v1.5",
+    n_results: Optional[int] = 10,
+    client: Optional[QdrantClient] = None,
+    collection_name: Optional[str] = "all-my-documents",
+    search_string: Optional[str] = "",
+    embedding_model: Optional[str] = "BAAI/bge-small-en-v1.5",
     qdrant_client_options: Optional[Dict] = {},
 ) -> List[List[QueryResponse]]:
     """Perform a similarity search with filters on a Qdrant collection
@@ -304,10 +304,18 @@ class QueryResponse(BaseModel, extra="forbid"):  # type: ignore
         ),
     )
 
-    data = {
-        "ids": [[result.id for result in sublist] for sublist in results],
-        "documents": [[result.document for result in sublist] for sublist in results],
-        "distances": [[result.score for result in sublist] for sublist in results],
-        "metadatas": [[result.metadata for result in sublist] for sublist in results],
-    }
+    data = [
+        [
+            QueryResponse(
+                id=result.id,
+                embedding=result.embedding,
+                metadata=result.metadata,
+                document=result.document,
+                score=result.score,
+            )
+            for result in sublist
+        ]
+        for sublist in results
+    ]
+
     return data

autogen/math_utils.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional, Tuple, Union
 
 from autogen import DEFAULT_MODEL, oai
 
@@ -9,7 +9,7 @@
 }
 
 
-def solve_problem(problem: str, **config) -> str:
+def solve_problem(problem: str, **config) -> Tuple[Optional[str], int]:
     """(openai<1) Solve the math problem.
 
     Args:
@@ -25,7 +25,7 @@ def solve_problem(problem: str, **config) -> str:
     return results.get("voted_answer"), response["cost"]
 
 
-def remove_boxed(string: str) -> Optional[str]:
+def remove_boxed(string: str) -> Union[str, None]:
     """Source: https://github.com/hendrycks/math
     Extract the text within a \\boxed{...} environment.
     Example:
@@ -36,15 +36,15 @@ def remove_boxed(string: str) -> Optional[str]:
     """
     left = "\\boxed{"
     try:
-        if not all((string[: len(left)] == left, string[-1] == "}")):
+        if not ((string.startswith(left) and string[-1] == "}")):
             raise AssertionError
 
         return string[len(left) : -1]
     except Exception:
         return None
 
 
-def last_boxed_only_string(string: str) -> Optional[str]:
+def last_boxed_only_string(string: str) -> Union[str, None]:
     """Source: https://github.com/hendrycks/math
     Extract the last \\boxed{...} or \\fbox{...} element from a string.
     """
@@ -96,7 +96,7 @@ def _fix_fracs(string: str) -> str:
                 new_str += substr
             else:
                 try:
-                    if not len(substr) >= 2:
+                    if len(substr) < 2:
                         raise AssertionError
                 except Exception:
                     return string
@@ -132,7 +132,7 @@ def _fix_a_slash_b(string: str) -> str:
     try:
         a = int(a_str)
         b = int(b_str)
-        if not string == "{}/{}".format(a, b):
+        if string != "{}/{}".format(a, b):
             raise AssertionError
         new_string = "\\frac{" + str(a) + "}{" + str(b) + "}"
         return new_string
@@ -147,7 +147,7 @@ def _remove_right_units(string: str) -> str:
     """
     if "\\text{ " in string:
         splits = string.split("\\text{ ")
-        if not len(splits) == 2:
+        if len(splits) != 2:
             raise AssertionError
         return splits[0]
     else:
@@ -161,16 +161,17 @@ def _fix_sqrt(string: str) -> str:
     >>> _fix_sqrt("\\sqrt3")
     \\sqrt{3}
     """
-    if "\\sqrt" not in string:
+    SQRT_LITERAL = "\\sqrt"  # Define a constant for the repeated literal
+    if SQRT_LITERAL not in string:
         return string
-    splits = string.split("\\sqrt")
+    splits = string.split(SQRT_LITERAL)
     new_string = splits[0]
     for split in splits[1:]:
         if split[0] != "{":
             a = split[0]
-            new_substr = "\\sqrt{" + a + "}" + split[1:]
+            new_substr = SQRT_LITERAL + "{" + a + "}" + split[1:]
         else:
-            new_substr = "\\sqrt" + split
+            new_substr = SQRT_LITERAL + split
         new_string += new_substr
     return new_string
 
@@ -310,7 +311,7 @@ def voting_counts(responses):
     return answers
 
 
-def eval_math_responses(responses, solution=None, **args):
+def eval_math_responses(responses, solution=None, **args) -> Dict:
     """Select a response for a math problem using voting, and check if the response is correct if the solution is provided.
 
     Args:

autogen/retrieve_utils.py

@@ -2,7 +2,7 @@
 import hashlib
 import os
 import re
-from typing import Callable, List, Tuple, Union
+from typing import Callable, List, Optional, Tuple, Union
 from urllib.parse import urlparse
 
 import chromadb
@@ -195,7 +195,9 @@ def split_files_to_chunks(
     return chunks, sources
 
 
-def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMATS, recursive: bool = True):
+def get_files_from_dir(
+    dir_path: Union[str, List[str]], types: Optional[list] = TEXT_FORMATS, recursive: Optional[bool] = True
+):
     """Return a list of all the files in a given directory, a url, a file path or a list of them."""
     if len(types) == 0:
         raise ValueError("types cannot be empty.")
@@ -245,7 +247,7 @@ def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMA
     return files
 
 
-def parse_html_to_markdown(html: str, url: str = None) -> str:
+def parse_html_to_markdown(html: str, url: Optional[str] = None) -> str:
     """Parse HTML to markdown."""
     soup = BeautifulSoup(html, "html.parser")
     title = soup.title.string
@@ -278,14 +280,16 @@ def parse_html_to_markdown(html: str, url: str = None) -> str:
 
 def _generate_file_name_from_url(url: str, max_length=255) -> str:
     url_bytes = url.encode("utf-8")
-    hash = hashlib.blake2b(url_bytes).hexdigest()
+    url_hash = hashlib.blake2b(url_bytes).hexdigest()
     parsed_url = urlparse(url)
     file_name = os.path.basename(url)
-    file_name = f"{parsed_url.netloc}_{file_name}_{hash[:min(8, max_length-len(parsed_url.netloc)-len(file_name)-1)]}"
+    file_name = (
+        f"{parsed_url.netloc}_{file_name}_{url_hash[:min(8, max_length-len(parsed_url.netloc)-len(file_name)-1)]}"
+    )
     return file_name
 
 
-def get_file_from_url(url: str, save_path: str = None) -> Tuple[str, str]:
+def get_file_from_url(url: str, save_path: str = None) -> Union[Tuple[str, str], None]:
     """Download a file from a URL."""
     if save_path is None:
         save_path = "tmp/chromadb"
@@ -333,19 +337,19 @@ def is_url(string: str):
 
 def create_vector_db_from_dir(
     dir_path: Union[str, List[str]],
-    max_tokens: int = 4000,
-    client: API = None,
-    db_path: str = "tmp/chromadb.db",
-    collection_name: str = "all-my-documents",
-    get_or_create: bool = False,
-    chunk_mode: str = "multi_lines",
-    must_break_at_empty_line: bool = True,
-    embedding_model: str = "all-MiniLM-L6-v2",
-    embedding_function: Callable = None,
-    custom_text_split_function: Callable = None,
-    custom_text_types: List[str] = TEXT_FORMATS,
-    recursive: bool = True,
-    extra_docs: bool = False,
+    max_tokens: Optional[int] = 4000,
+    client: Optional[API] = None,
+    db_path: Optional[str] = "tmp/chromadb.db",
+    collection_name: Optional[str] = "all-my-documents",
+    get_or_create: Optional[bool] = False,
+    chunk_mode: Optional[str] = "multi_lines",
+    must_break_at_empty_line: Optional[bool] = True,
+    embedding_model: Optional[str] = "all-MiniLM-L6-v2",
+    embedding_function: Optional[Callable] = None,
+    custom_text_split_function: Optional[Callable] = None,
+    custom_text_types: Optional[List[str]] = TEXT_FORMATS,
+    recursive: Optional[bool] = True,
+    extra_docs: Optional[bool] = False,
 ) -> API:
     """Create a vector db from all the files in a given directory, the directory can also be a single file or a url to
         a single file. We support chromadb compatible APIs to create the vector db, this function is not required if
@@ -426,13 +430,13 @@ def create_vector_db_from_dir(
 
 def query_vector_db(
     query_texts: List[str],
-    n_results: int = 10,
-    client: API = None,
-    db_path: str = "tmp/chromadb.db",
-    collection_name: str = "all-my-documents",
-    search_string: str = "",
-    embedding_model: str = "all-MiniLM-L6-v2",
-    embedding_function: Callable = None,
+    n_results: Optional[int] = 10,
+    client: Optional[API] = None,
+    db_path: Optional[str] = "tmp/chromadb.db",
+    collection_name: Optional[str] = "all-my-documents",
+    search_string: Optional[str] = "",
+    embedding_model: Optional[str] = "all-MiniLM-L6-v2",
+    embedding_function: Optional[Callable] = None,
 ) -> QueryResult:
     """Query a vector db. We support chromadb compatible APIs, it's not required if you prepared your own vector db
         and query function.

autogen/runtime_logging.py

@@ -31,8 +31,7 @@ def start(logger_type: str = "sqlite", config: Optional[Dict[str, Any]] = None)
         is_logging = True
     except Exception as e:
         logger.error(f"[runtime logging] Failed to start logging: {e}")
-    finally:
-        return session_id
+    return session_id
 
 
 def log_chat_completion(