Updated the azure client to support AAD auth. (#2879)

microsoft · Jun 6, 2024 · f904b4f · f904b4f
1 parent 84c7c24
commit f904b4f
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 4 deletions.
diff --git a/autogen/logger/sqlite_logger.py b/autogen/logger/sqlite_logger.py
@@ -236,7 +236,16 @@ def log_new_agent(self, agent: ConversableAgent, init_args: Dict[str, Any]) -> N
 
         args = to_dict(
             init_args,
-            exclude=("self", "__class__", "api_key", "organization", "base_url", "azure_endpoint"),
+            exclude=(
+                "self",
+                "__class__",
+                "api_key",
+                "organization",
+                "base_url",
+                "azure_endpoint",
+                "azure_ad_token",
+                "azure_ad_token_provider",
+            ),
             no_recursive=(Agent,),
         )
 
@@ -301,7 +310,17 @@ def log_new_wrapper(self, wrapper: OpenAIWrapper, init_args: Dict[str, Union[LLM
             return
 
         args = to_dict(
-            init_args, exclude=("self", "__class__", "api_key", "organization", "base_url", "azure_endpoint")
+            init_args,
+            exclude=(
+                "self",
+                "__class__",
+                "api_key",
+                "organization",
+                "base_url",
+                "azure_endpoint",
+                "azure_ad_token",
+                "azure_ad_token_provider",
+            ),
         )
 
         query = """
@@ -323,7 +342,17 @@ def log_new_client(
             return
 
         args = to_dict(
-            init_args, exclude=("self", "__class__", "api_key", "organization", "base_url", "azure_endpoint")
+            init_args,
+            exclude=(
+                "self",
+                "__class__",
+                "api_key",
+                "organization",
+                "base_url",
+                "azure_endpoint",
+                "azure_ad_token",
+                "azure_ad_token_provider",
+            ),
         )
 
         query = """

diff --git a/autogen/oai/client.py b/autogen/oai/client.py
@@ -407,6 +407,14 @@ def _configure_azure_openai(self, config: Dict[str, Any], openai_config: Dict[st
             openai_config["azure_deployment"] = openai_config["azure_deployment"].replace(".", "")
         openai_config["azure_endpoint"] = openai_config.get("azure_endpoint", openai_config.pop("base_url", None))
 
+        # Create a default Azure token provider if requested
+        if openai_config.get("azure_ad_token_provider") == "DEFAULT":
+            import azure.identity
+
+            openai_config["azure_ad_token_provider"] = azure.identity.get_bearer_token_provider(
+                azure.identity.DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+
     def _register_default_client(self, config: Dict[str, Any], openai_config: Dict[str, Any]) -> None:
         """Create a client with the given config to override openai_config,
         after removing extra kwargs.

diff --git a/autogen/oai/openai_utils.py b/autogen/oai/openai_utils.py
@@ -13,7 +13,7 @@
 from openai.types.beta.assistant import Assistant
 from packaging.version import parse
 
-NON_CACHE_KEY = ["api_key", "base_url", "api_type", "api_version"]
+NON_CACHE_KEY = ["api_key", "base_url", "api_type", "api_version", "azure_ad_token", "azure_ad_token_provider"]
 DEFAULT_AZURE_API_VERSION = "2024-02-15-preview"
 OAI_PRICE1K = {
     # https://openai.com/api/pricing/