Fix scope for Azure auth with LiteLLM

graphrag/index/validate_config.py

@@ -33,6 +33,7 @@ def validate_config_names(parameters: GraphRagConfig) -> None:
         logger.info("LLM Config Params Validated")
     except Exception as e:  # noqa: BLE001
         logger.error(f"LLM configuration error detected. Exiting...\n{e}")  # noqa
+        print("Failed to validate language model params", e)  # noqa: T201
         sys.exit(1)
 
     # Validate Embeddings LLM configs
@@ -53,4 +54,5 @@ def validate_config_names(parameters: GraphRagConfig) -> None:
         logger.info("Embedding LLM Config Params Validated")
     except Exception as e:  # noqa: BLE001
         logger.error(f"Embedding LLM configuration error detected. Exiting...\n{e}")  # noqa
+        print("Failed to validate embedding model params", e)  # noqa: T201
         sys.exit(1)

graphrag/language_model/providers/litellm/chat_model.py

@@ -86,9 +86,10 @@ def _create_base_completions(
             msg = "Azure Managed Identity authentication is only supported for Azure models."
             raise ValueError(msg)
 
+        base_args["azure_scope"] = base_args.pop("audience")
         base_args["azure_ad_token_provider"] = get_bearer_token_provider(
             DefaultAzureCredential(),
-            COGNITIVE_SERVICES_AUDIENCE,
+            model_config.audience or COGNITIVE_SERVICES_AUDIENCE,
         )
 
     def _base_completion(**kwargs: Any) -> ModelResponse | CustomStreamWrapper:

graphrag/language_model/providers/litellm/embedding_model.py

@@ -72,9 +72,10 @@ def _create_base_embeddings(
             msg = "Azure Managed Identity authentication is only supported for Azure models."
             raise ValueError(msg)
 
+        base_args["azure_scope"] = base_args.pop("audience")
         base_args["azure_ad_token_provider"] = get_bearer_token_provider(
             DefaultAzureCredential(),
-            COGNITIVE_SERVICES_AUDIENCE,
+            model_config.audience or COGNITIVE_SERVICES_AUDIENCE,
         )
 
     def _base_embedding(**kwargs: Any) -> EmbeddingResponse: