Add AI Search audience parameter (#10)

backend/graphrag-wheel/note.txt

@@ -2,4 +2,4 @@ This graphrag wheel file was built from the following repo
 
 https://github.com/microsoft/graphrag
 
-on commit hash 8cb189635e90d49231f3f09b54e69d4daae1371d
+on commit hash b860d08a907e834166edf03e617d9cfeac946a64

backend/src/api/index.py

@@ -58,6 +58,7 @@
 )  # TODO: update API to use the AzureStorageClientManager
 
 ai_search_url = os.environ["AI_SEARCH_URL"]
+ai_search_audience = os.environ["AI_SEARCH_AUDIENCE"]
 
 index_route = APIRouter(
     prefix="/index",
@@ -544,9 +545,10 @@ async def delete_index(index_name: str):
         except Exception:
             pass
 
-        url = ai_search_url
         index_client = SearchIndexClient(
-            endpoint=url, credential=DefaultAzureCredential()
+            endpoint=ai_search_url, 
+            credential=DefaultAzureCredential(), 
+            audience=ai_search_audience
         )
         ai_search_index_name = f"{sanitized_index_name}_description_embedding"
         if ai_search_index_name in index_client.list_index_names():

backend/src/api/pipeline-settings.yaml

@@ -72,6 +72,7 @@ embeddings:
     title_column: name
     overwrite: True
     url: $AI_SEARCH_URL
+    audience: $AI_SEARCH_AUDIENCE
 
 # entity_extraction:
 #   prompt: PLACEHOLDER

backend/src/meta_agent/community/retrieve.py

@@ -322,8 +322,9 @@ def similarity_search_by_vector(
         docs = []
         for collection_idx, collection_name in enumerate(self.collections):
             add_on = "-" + str(collection_idx)
+            audience = os.environ["AI_SEARCH_AUDIENCE"]
             db_connection = SearchClient(
-                self.url, collection_name, DefaultAzureCredential()
+                self.url, collection_name, DefaultAzureCredential(), audience=audience,
             )
             response = db_connection.search(
                 vector_queries=[vectorized_query],

docs/DEPLOYMENT-GUIDE.md

@@ -74,10 +74,11 @@ In the `deploy.parameters.json` file, provide values for the following required
 `GRAPHRAG_LLM_DEPLOYMENT_NAME`         |                                    | Yes | Deployment name of the gpt-4 turbo model.
 `GRAPHRAG_EMBEDDING_MODEL`             | text-embedding-ada-002             | Yes | Name of the Azure OpenAI embedding model.
 `GRAPHRAG_EMBEDDING_DEPLOYMENT_NAME`   |                                    | Yes | Deployment name of the Azure OpenAI embedding model.
-`GRAPHRAG_COGNITIVE_SERVICES_ENDPOINT` |                                  | No  | Endpoint for cognitive services identity authorization. Will default to `https://cognitiveservices.azure.com/.default` for Azure Commercial cloud but should be defined for deployments in other Azure clouds.
+`GRAPHRAG_COGNITIVE_SERVICES_ENDPOINT` |                                    | No  | Endpoint for cognitive services identity authorization. Will default to `https://cognitiveservices.azure.com/.default` for Azure Commercial cloud but should be defined for deployments in other Azure clouds.
 `APIM_NAME`                            |                                    | No  | Hostname of the API. Must be a globally unique name. The API will be accessible at `https://<APIM_NAME>.azure-api.net`. If not provided a unique name will be generated.
 `RESOURCE_BASE_NAME`                   |                                    | No  | Suffix to apply to all azure resource names. If not provided a unique suffix will be generated.
 `AISEARCH_ENDPOINT_SUFFIX`             |                                    | No  | Suffix to apply to AI search endpoint. Will default to `search.windows.net` for Azure Commercial cloud but should be overriden for deployments in other Azure clouds.
+`AISEARCH_AUDIENCE`                    |                                    | No  | Audience for AAD for AI Search. Will default to `https://search.azure.com/` for Azure Commercial cloud but should be overriden for deployments in other Azure clouds.D
 `REPORTERS`                            |                                    | No  | The type of logging to enable. If not provided, logging will be saved to a file in Azure Storage and to the console in AKS.
 
 ## 5. Deploy the solution accelerator

infra/deploy.sh

@@ -7,6 +7,7 @@
 aksNamespace="graphrag"
 
 # OPTIONAL PARAMS
+AISEARCH_AUDIENCE=""
 AISEARCH_ENDPOINT_SUFFIX=""
 APIM_NAME=""
 RESOURCE_BASE_NAME=""
@@ -129,6 +130,10 @@ populateOptionalParams () {
         AISEARCH_ENDPOINT_SUFFIX="search.windows.net"
         printf "\tsetting AISEARCH_ENDPOINT_SUFFIX=$AISEARCH_ENDPOINT_SUFFIX\n"
     fi
+    if [ -z "$AISEARCH_AUDIENCE" ]; then
+        AISEARCH_AUDIENCE="https://search.azure.com"
+        printf "\tsetting AISEARCH_AUDIENCE=$AISEARCH_AUDIENCE\n"
+    fi
     if [ -z "$PUBLISHER_NAME" ]; then
         PUBLISHER_NAME="publisher"
         printf "\tsetting PUBLISHER_NAME=$PUBLISHER_NAME\n"
@@ -359,6 +364,7 @@ deployHelmChart () {
         --set "ingress.host=$graphragHostname" \
         --set "graphragConfig.APIM_GATEWAY_URL=$apimGatewayUrl" \
         --set "graphragConfig.AI_SEARCH_URL=https://$aiSearchName.$AISEARCH_ENDPOINT_SUFFIX" \
+        --set "graphragConfig.AI_SEARCH_AUDIENCE=$AISEARCH_AUDIENCE" \
         --set "graphragConfig.COSMOS_URI_ENDPOINT=$cosmosEndpoint" \
         --set "graphragConfig.GRAPHRAG_API_BASE=$GRAPHRAG_API_BASE" \
         --set "graphragConfig.GRAPHRAG_API_VERSION=$GRAPHRAG_API_VERSION" \