From 8c27cf97599045b16ebac421eaf9a73d3db24a88 Mon Sep 17 00:00:00 2001 From: Young Date: Mon, 1 Jul 2024 05:11:59 +0000 Subject: [PATCH] Support managed_identity_client_id for DefaultAzureCredential --- rdagent/core/conf.py | 2 ++ rdagent/oai/llm_utils.py | 6 +++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/rdagent/core/conf.py b/rdagent/core/conf.py index 88476381..120cb8f8 100644 --- a/rdagent/core/conf.py +++ b/rdagent/core/conf.py @@ -15,8 +15,10 @@ class RDAgentSettings(BaseSettings): + # TODO: (xiao) I think most of the config should be in oai.config use_azure: bool = True use_azure_token_provider: bool = False + managed_identity_client_id: str | None = None max_retry: int = 10 retry_wait_seconds: int = 1 dump_chat_cache: bool = False diff --git a/rdagent/oai/llm_utils.py b/rdagent/oai/llm_utils.py index 7be21525..72ab564c 100644 --- a/rdagent/oai/llm_utils.py +++ b/rdagent/oai/llm_utils.py @@ -292,6 +292,7 @@ def __init__( # noqa: C901, PLR0912, PLR0915 else: self.use_azure = self.cfg.use_azure self.use_azure_token_provider = self.cfg.use_azure_token_provider + self.managed_identity_client_id = self.cfg.managed_identity_client_id self.chat_api_key = self.cfg.chat_openai_api_key if chat_api_key is None else chat_api_key self.chat_model = self.cfg.chat_model if chat_model is None else chat_model @@ -314,7 +315,10 @@ def __init__( # noqa: C901, PLR0912, PLR0915 if self.use_azure: if self.use_azure_token_provider: - credential = DefaultAzureCredential() + dac_kwargs = {} + if self.managed_identity_client_id is not None: + dac_kwargs["managed_identity_client_id"] = self.managed_identity_client_id + credential = DefaultAzureCredential(**dac_kwargs) token_provider = get_bearer_token_provider( credential, "https://cognitiveservices.azure.com/.default",