refactor: Replcae OpenAI API calls with SDK API call

Author: Avijit-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

@@ -1,13 +1,14 @@
 from azure.keyvault.secrets import SecretClient
-from openai import AzureOpenAI
+from azure.ai.inference import EmbeddingsClient
 import re
 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
-from azure.identity import (AzureCliCredential, get_bearer_token_provider)
+from azure.identity import AzureCliCredential
 
 
 key_vault_name = 'kv_to-be-replaced'
@@ -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,18 +57,19 @@ 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"
-    ad_token_provider = get_bearer_token_provider(
-        credential, "https://cognitiveservices.azure.com/.default"
-    )
-    client = AzureOpenAI(
-        api_version=openai_api_version,
-        azure_endpoint=openai_api_base,
-        azure_ad_token_provider=ad_token_provider
+def get_embeddings(text: str, ai_project_endpoint: str):
+    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"]
     )
 
-    embedding = client.embeddings.create(input=text, model=model_id).data[0].embedding
+    response = embeddings_client.embed(model=embedding_model, input=[text])
+    embedding = response.data[0].embedding
     return embedding
 
 
@@ -126,12 +126,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/index_scripts/requirements.txt

@@ -1,5 +1,5 @@
 azure-storage-file-datalake==12.20.0
-openai==1.84.0
+azure-ai-inference==1.0.0b9
 pypdf==5.6.0
 # pyodbc
 tiktoken==0.9.0

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,27 +13,51 @@
 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
 import requests
 import tiktoken
 from azure.ai.documentintelligence import DocumentIntelligenceClient
 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
 from langchain.text_splitter import (MarkdownTextSplitter,
                                      PythonCodeTextSplitter,
                                      RecursiveCharacterTextSplitter,
                                      TextSplitter)
-from openai import AzureOpenAI
 from tqdm import tqdm
 
 # 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,45 +849,27 @@ 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")
-    )
+    # Get AI Project endpoint from Key Vault
+    ai_project_endpoint = get_secrets_from_kv("AZURE-AI-AGENT-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"
-        )
+    # 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"
-            api_version = "2024-02-01"
-
-            if azure_credential is not None:
-                api_key = azure_credential.get_token(
-                    "https://cognitiveservices.azure.com/.default"
-                ).token
-            else:
-                api_key = (
-                    embedding_model_key
-                    if embedding_model_key
-                    else os.getenv("AZURE_OPENAI_API_KEY")
-                )
-
-            client = AzureOpenAI(
-                api_version=api_version, azure_endpoint=endpoint, api_key=api_key
-            )
-            embeddings = client.embeddings.create(model=deployment_id, input=text)
+        credential = azure_credential if azure_credential is not None else AzureCliCredential()
+        embeddings_client = EmbeddingsClient(
+            endpoint=inference_endpoint,
+            credential=credential,
+            credential_scopes=["https://cognitiveservices.azure.com/.default"]
+        )
 
-            return embeddings.model_dump()["data"][0]["embedding"]
+        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}"
         )
 
 

src/requirements.txt

@@ -1,6 +1,5 @@
 azure-identity==1.25.0
 # Flask[async]==2.3.2
-openai==2.0.1
 azure-search-documents==11.7.0b1
 azure-storage-blob==12.26.0
 python-dotenv==1.1.1