Fix threading issue for logging (#1901)

* fix mypy errors for logging * cleanup * formatting * fix threading issue in logging * remove kwarg --------- Co-authored-by: Eric Zhu <[email protected]>
microsoft · Mar 8, 2024 · 811fd69 · 811fd69
1 parent 2a62ffc
commit 811fd69
Show file tree

Hide file tree

Showing 5 changed files with 162 additions and 130 deletions.
diff --git a/autogen/logger/base_logger.py b/autogen/logger/base_logger.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Dict, TYPE_CHECKING, Union
+from typing import Any, Dict, List, TYPE_CHECKING, Union
 import sqlite3
 import uuid
 
@@ -11,6 +11,9 @@
 if TYPE_CHECKING:
     from autogen import ConversableAgent, OpenAIWrapper
 
+ConfigItem = Dict[str, Union[str, List[str]]]
+LLMConfig = Dict[str, Union[None, float, int, ConfigItem, List[ConfigItem]]]
+
 
 class BaseLogger(ABC):
     @abstractmethod
@@ -25,10 +28,11 @@ def start(self) -> str:
 
     @abstractmethod
     def log_chat_completion(
+        self,
         invocation_id: uuid.UUID,
         client_id: int,
         wrapper_id: int,
-        request: Dict,
+        request: Dict[str, Union[float, str, List[Dict[str, str]]]],
         response: Union[str, ChatCompletion],
         is_cached: int,
         cost: float,
@@ -54,7 +58,7 @@ def log_chat_completion(
         ...
 
     @abstractmethod
-    def log_new_agent(agent: ConversableAgent, init_args: Dict) -> None:
+    def log_new_agent(self, agent: ConversableAgent, init_args: Dict[str, Any]) -> None:
         """
         Log the birth of a new agent.
 
@@ -65,7 +69,7 @@ def log_new_agent(agent: ConversableAgent, init_args: Dict) -> None:
         ...
 
     @abstractmethod
-    def log_new_wrapper(wrapper: OpenAIWrapper, init_args: Dict) -> None:
+    def log_new_wrapper(self, wrapper: OpenAIWrapper, init_args: Dict[str, Union[LLMConfig, List[LLMConfig]]]) -> None:
         """
         Log the birth of a new OpenAIWrapper.
 
@@ -76,7 +80,9 @@ def log_new_wrapper(wrapper: OpenAIWrapper, init_args: Dict) -> None:
         ...
 
     @abstractmethod
-    def log_new_client(client: Union[AzureOpenAI, OpenAI], wrapper: OpenAIWrapper, init_args: Dict) -> None:
+    def log_new_client(
+        self, client: Union[AzureOpenAI, OpenAI], wrapper: OpenAIWrapper, init_args: Dict[str, Any]
+    ) -> None:
         """
         Log the birth of a new OpenAIWrapper.
 
@@ -87,14 +93,14 @@ def log_new_client(client: Union[AzureOpenAI, OpenAI], wrapper: OpenAIWrapper, i
         ...
 
     @abstractmethod
-    def stop() -> None:
+    def stop(self) -> None:
         """
         Close the connection to the logging database, and stop logging.
         """
         ...
 
     @abstractmethod
-    def get_connection() -> Union[sqlite3.Connection]:
+    def get_connection(self) -> Union[None, sqlite3.Connection]:
         """
         Return a connection to the logging database.
         """

diff --git a/autogen/logger/logger_utils.py b/autogen/logger/logger_utils.py
@@ -5,14 +5,14 @@
 __all__ = ("get_current_ts", "to_dict")
 
 
-def get_current_ts():
+def get_current_ts() -> str:
     return datetime.datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S.%f")
 
 
 def to_dict(
     obj: Union[int, float, str, bool, Dict[Any, Any], List[Any], Tuple[Any, ...], Any],
-    exclude: Tuple[str] = (),
-    no_recursive: Tuple[str] = (),
+    exclude: Tuple[str, ...] = (),
+    no_recursive: Tuple[Any, ...] = (),
 ) -> Any:
     if isinstance(obj, (int, float, str, bool)):
         return obj