Add Foundry changes

Author: NirajC-Microsoft

infra/main.bicep

@@ -1105,6 +1105,11 @@ module keyvault 'br/public:avm/res/key-vault/vault:0.12.1' = {
       {name: 'AZURE-OPENAI-PREVIEW-API-VERSION', value: azureOpenaiAPIVersion}
       {name: 'AZURE-OPEN-AI-DEPLOYMENT-MODEL', value: gptModelName}
       {name: 'TENANT-ID', value: subscription().tenantId}
+      {
+        name: 'AZURE-AI-AGENT-ENDPOINT'
+        value: aiFoundryAiProjectEndpoint
+      }
+      
     ]
   }
   dependsOn:[

infra/scripts/index_scripts/02_process_data.py

@@ -4,6 +4,7 @@
 import time
 import pypdf
 from io import BytesIO
+from urllib.parse import urlparse
 from azure.search.documents import SearchClient
 from azure.storage.filedatalake import DataLakeServiceClient
 from azure.search.documents.indexes import SearchIndexClient
@@ -36,9 +37,7 @@ def get_secrets_from_kv(secret_name: str) -> str:
 
 # Retrieve secrets from Key Vault
 search_endpoint = get_secrets_from_kv("AZURE-SEARCH-ENDPOINT")
-openai_api_base = get_secrets_from_kv("AZURE-OPENAI-ENDPOINT")
-openai_api_version = get_secrets_from_kv("AZURE-OPENAI-PREVIEW-API-VERSION")
-deployment = get_secrets_from_kv("AZURE-OPEN-AI-DEPLOYMENT-MODEL")
+ai_project_endpoint = get_secrets_from_kv("AZURE-AI-AGENT-ENDPOINT")
 account_name = get_secrets_from_kv("ADLS-ACCOUNT-NAME")
 print("Secrets retrieved from Key Vault.")
 
@@ -58,15 +57,29 @@ def get_secrets_from_kv(secret_name: str) -> str:
 
 
 # Function: Get Embeddings
-def get_embeddings(text: str, openai_api_base, openai_api_version):
-    model_id = "text-embedding-ada-002"
-    client = EmbeddingsClient(
-        endpoint=f"{openai_api_base}/openai/deployments/{model_id}",
+def get_embeddings(text: str, ai_project_endpoint: str):
+    """Get embeddings using Azure AI Foundry SDK.
+    
+    Args:
+        text: Text to embed
+        ai_project_endpoint: Azure AI Project endpoint from Key Vault
+                            (e.g., https://aif-xyz.services.ai.azure.com/api/projects/proj-xyz)
+    
+    Returns:
+        List of embedding values
+    """
+    embedding_model = "text-embedding-ada-002"
+    
+    # Construct inference endpoint with /models path
+    inference_endpoint = f"https://{urlparse(ai_project_endpoint).netloc}/models"
+    
+    embeddings_client = EmbeddingsClient(
+        endpoint=inference_endpoint,
         credential=credential,
         credential_scopes=["https://cognitiveservices.azure.com/.default"]
     )
 
-    response = client.embed(input=[text])
+    response = embeddings_client.embed(model=embedding_model, input=[text])
     embedding = response.data[0].embedding
     return embedding
 
@@ -124,12 +137,12 @@ def prepare_search_doc(content, document_id):
         chunk_id = f"{document_id}_{str(idx).zfill(2)}"
 
         try:
-            v_contentVector = get_embeddings(str(chunk), openai_api_base, openai_api_version)
+            v_contentVector = get_embeddings(str(chunk), ai_project_endpoint)
         except Exception as e:
             print(f"Error occurred: {e}. Retrying after 30 seconds...")
             time.sleep(30)
             try:
-                v_contentVector = get_embeddings(str(chunk), openai_api_base, openai_api_version)
+                v_contentVector = get_embeddings(str(chunk), ai_project_endpoint)
             except Exception as e:
                 print(f"Retry failed: {e}. Setting v_contentVector to an empty list.")
                 v_contentVector = []

infra/scripts/run_create_index_scripts.sh

@@ -119,6 +119,7 @@ fi
 #Replace key vault name 
 sed -i "s/kv_to-be-replaced/${keyvaultName}/g" "infra/scripts/index_scripts/01_create_search_index.py"
 sed -i "s/kv_to-be-replaced/${keyvaultName}/g" "infra/scripts/index_scripts/02_process_data.py"
+sed -i "s/kv_to-be-replaced/${keyvaultName}/g" "scripts/data_utils.py"
 if [ -n "$managedIdentityClientId" ]; then
     sed -i "s/mici_to-be-replaced/${managedIdentityClientId}/g" "infra/scripts/index_scripts/01_create_search_index.py"
     sed -i "s/mici_to-be-replaced/${managedIdentityClientId}/g" "infra/scripts/index_scripts/02_process_data.py"
@@ -181,6 +182,7 @@ fi
 # revert the key vault name and managed identity client id in the python files
 sed -i "s/${keyvaultName}/kv_to-be-replaced/g" "infra/scripts/index_scripts/01_create_search_index.py"
 sed -i "s/${keyvaultName}/kv_to-be-replaced/g" "infra/scripts/index_scripts/02_process_data.py"
+sed -i "s/${keyvaultName}/kv_to-be-replaced/g" "scripts/data_utils.py"
 if [ -n "$managedIdentityClientId" ]; then
     sed -i "s/${managedIdentityClientId}/mici_to-be-replaced/g" "infra/scripts/index_scripts/01_create_search_index.py"
     sed -i "s/${managedIdentityClientId}/mici_to-be-replaced/g" "infra/scripts/index_scripts/02_process_data.py"

scripts/data_utils.py

@@ -13,6 +13,7 @@
 from dataclasses import dataclass
 from functools import partial
 from typing import Any, Callable, Dict, Generator, List, Optional, Tuple, Union
+from urllib.parse import urlparse
 
 import fitz
 import markdown
@@ -22,6 +23,8 @@
 from azure.ai.documentintelligence.models import AnalyzeDocumentRequest
 from azure.ai.inference import EmbeddingsClient
 from azure.core.credentials import AzureKeyCredential
+from azure.identity import AzureCliCredential
+from azure.keyvault.secrets import SecretClient
 from azure.storage.blob import ContainerClient
 from bs4 import BeautifulSoup
 from dotenv import load_dotenv
@@ -34,6 +37,27 @@
 # Configure environment variables
 load_dotenv()  # take environment variables from .env.
 
+# Key Vault name - replaced during deployment
+key_vault_name = 'kv_to-be-replaced'
+
+
+def get_secrets_from_kv(secret_name: str) -> str:
+    """Retrieves a secret value from Azure Key Vault.
+    
+    Args:
+        secret_name: Name of the secret
+    
+    Returns:
+        The secret value
+    """
+    kv_credential = AzureCliCredential()
+    secret_client = SecretClient(
+        vault_url=f"https://{key_vault_name}.vault.azure.net/",
+        credential=kv_credential
+    )
+    return secret_client.get_secret(secret_name).value
+
+
 FILE_FORMAT_DICT = {
     "md": "markdown",
     "txt": "text",
@@ -825,47 +849,33 @@ def get_payload_and_headers_cohere(text, aad_token) -> Tuple[Dict, Dict]:
 def get_embedding(
     text, embedding_model_endpoint=None, embedding_model_key=None, azure_credential=None
 ):
-    endpoint = (
-        embedding_model_endpoint
-        if embedding_model_endpoint
-        else os.environ.get("EMBEDDING_MODEL_ENDPOINT")
-    )
-
-    FLAG_EMBEDDING_MODEL = os.getenv("FLAG_EMBEDDING_MODEL", "AOAI")
-
-    if azure_credential is None and (endpoint is None):
-        raise Exception(
-            "EMBEDDING_MODEL_ENDPOINT and EMBEDDING_MODEL_KEY are required for embedding"
-        )
+    # Get AI Project endpoint from Key Vault
+    ai_project_endpoint = get_secrets_from_kv("AZURE-AI-AGENT-ENDPOINT")
+    
+    # Construct inference endpoint: https://aif-xyz.services.ai.azure.com/models
+    inference_endpoint = f"https://{urlparse(ai_project_endpoint).netloc}/models"
+    embedding_model = "text-embedding-ada-002"
 
     try:
-        if FLAG_EMBEDDING_MODEL == "AOAI":
-            deployment_id = "embedding"
-
-            if azure_credential is not None:
-                # Use managed identity credential with credential_scopes parameter
-                client = EmbeddingsClient(
-                    endpoint=f"{endpoint}/openai/deployments/{deployment_id}",
-                    credential=azure_credential,
-                    credential_scopes=["https://cognitiveservices.azure.com/.default"]
-                )
-            else:
-                # Use API key credential
-                api_key = (
-                    embedding_model_key
-                    if embedding_model_key
-                    else os.getenv("AZURE_OPENAI_API_KEY")
-                )
-                client = EmbeddingsClient(
-                    endpoint=f"{endpoint}/openai/deployments/{deployment_id}",
-                    credential=AzureKeyCredential(api_key)
-                )
-            response = client.embed(input=[text])
-            return response.data[0].embedding
+        if azure_credential is not None:
+            embeddings_client = EmbeddingsClient(
+                endpoint=inference_endpoint,
+                credential=azure_credential,
+                credential_scopes=["https://cognitiveservices.azure.com/.default"]
+            )
+        else:
+            api_key = embedding_model_key or os.getenv("AZURE_OPENAI_API_KEY")
+            embeddings_client = EmbeddingsClient(
+                endpoint=inference_endpoint,
+                credential=AzureKeyCredential(api_key)
+            )
+        
+        response = embeddings_client.embed(model=embedding_model, input=[text])
+        return response.data[0].embedding
 
     except Exception as e:
         raise Exception(
-            f"Error getting embeddings with endpoint={endpoint} with error={e}"
+            f"Error getting embeddings with endpoint={inference_endpoint} with error={e}"
         )