Require explicit azure auth settings when using AOI.

- Must set LanguageModel.azure_auth_type to either
"api_key" or "managed_identity" when using AOI.

Author: dworthen

.semversioner/next-release/patch-20250128230417263466.json

@@ -0,0 +1,4 @@
+{
+  "type": "patch",
+  "description": "Require explicit azure auth settings when using AOI."
+}

graphrag/config/errors.py

@@ -15,6 +15,15 @@ def __init__(self, llm_type: str, azure_auth_type: str | None = None) -> None:
         super().__init__(msg)
 
 
+class AzureAuthTypeMissingError(ValueError):
+    """Azure Auth type missing error."""
+
+    def __init__(self, llm_type: str) -> None:
+        """Init method definition."""
+        msg = f"azure_auth_type is required for {llm_type}. Please rerun `graphrag init` and set the azure_auth_type."
+        super().__init__(msg)
+
+
 class AzureApiBaseMissingError(ValueError):
     """Azure API Base missing error."""
 

graphrag/config/models/language_model_config.py

@@ -12,6 +12,7 @@
     ApiKeyMissingError,
     AzureApiBaseMissingError,
     AzureApiVersionMissingError,
+    AzureAuthTypeMissingError,
     AzureDeploymentNameMissingError,
     ConflictingSettingsError,
 )
@@ -61,6 +62,23 @@ def _validate_api_key(self) -> None:
         default=None,
     )
 
+    def _validate_azure_auth_type(self) -> None:
+        """Validate the Azure authentication type.
+
+        azure_auth_type is required when using Azure OpenAI
+        and explicitly defines the authentication type to use.
+
+        Raises
+        ------
+        AzureAuthTypeMissingError
+            If the Azure authentication type is missing when required.
+        """
+        if (
+            self.type == LLMType.AzureOpenAIChat
+            or self.type == LLMType.AzureOpenAIEmbedding
+        ) and self.azure_auth_type is None:
+            raise AzureAuthTypeMissingError(self.type.value)
+
     type: LLMType = Field(description="The type of LLM model to use.")
     model: str = Field(description="The LLM model to use.")
     encoding_model: str = Field(description="The encoding model to use", default="")
@@ -220,13 +238,16 @@ def _validate_azure_settings(self) -> None:
 
         Raises
         ------
+        AzureAuthTypeMissingError
+            If the Azure authentication type is missing and is required.
         AzureApiBaseMissingError
             If the API base is missing and is required.
         AzureApiVersionMissingError
             If the API version is missing and is required.
         AzureDeploymentNameMissingError
             If the deployment name is missing and is required.
         """
+        self._validate_azure_auth_type()
         self._validate_api_base()
         self._validate_api_version()
         self._validate_deployment_name()

graphrag/query/llm/get_client.py

@@ -5,7 +5,7 @@
 
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 
-from graphrag.config.enums import LLMType
+from graphrag.config.enums import AzureAuthType, LLMType
 from graphrag.config.models.graph_rag_config import GraphRagConfig
 from graphrag.query.llm.oai.chat_openai import ChatOpenAI
 from graphrag.query.llm.oai.embedding import OpenAIEmbedding
@@ -31,7 +31,8 @@ def get_llm(config: GraphRagConfig) -> ChatOpenAI:
         api_key=default_llm_settings.api_key,
         azure_ad_token_provider=(
             get_bearer_token_provider(DefaultAzureCredential(), audience)
-            if is_azure_client and not default_llm_settings.api_key
+            if is_azure_client
+            and default_llm_settings.azure_auth_type == AzureAuthType.ManagedIdentity
             else None
         ),
         api_base=default_llm_settings.api_base,
@@ -65,7 +66,8 @@ def get_text_embedder(config: GraphRagConfig) -> OpenAIEmbedding:
         api_key=embeddings_llm_settings.api_key,
         azure_ad_token_provider=(
             get_bearer_token_provider(DefaultAzureCredential(), audience)
-            if is_azure_client and not embeddings_llm_settings.api_key
+            if is_azure_client
+            and embeddings_llm_settings.azure_auth_type == AzureAuthType.ManagedIdentity
             else None
         ),
         api_base=embeddings_llm_settings.api_base,