Debug: gemini client was not logged

microsoft · May 21, 2024 · 2cdf538 · 2cdf538
1 parent 9f33724
commit 2cdf538
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -435,7 +435,8 @@ def _register_default_client(self, config: Dict[str, Any], openai_config: Dict[s
             elif api_type is not None and api_type.startswith("google"):
                 if gemini_import_exception:
                     raise ImportError("Please install `google-generativeai` to use Google OpenAI API.")
-                self._clients.append(GeminiClient(**openai_config))
+                client = GeminiClient(**openai_config)
+                self._clients.append(client)
             else:
                 client = OpenAI(**openai_config)
                 self._clients.append(OpenAIClient(client))

diff --git a/autogen/runtime_logging.py b/autogen/runtime_logging.py
@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     from autogen import Agent, ConversableAgent, OpenAIWrapper
+    from autogen.oai.gemini import GeminiClient
 
 logger = logging.getLogger(__name__)
 
@@ -94,7 +95,9 @@ def log_new_wrapper(wrapper: OpenAIWrapper, init_args: Dict[str, Union[LLMConfig
     autogen_logger.log_new_wrapper(wrapper, init_args)
 
 
-def log_new_client(client: Union[AzureOpenAI, OpenAI], wrapper: OpenAIWrapper, init_args: Dict[str, Any]) -> None:
+def log_new_client(
+    client: Union[AzureOpenAI, OpenAI, GeminiClient], wrapper: OpenAIWrapper, init_args: Dict[str, Any]
+) -> None:
     if autogen_logger is None:
         logger.error("[runtime logging] log_new_client: autogen logger is None")
         return