Replcae OpenAI API calls with SDK API call

Author: NirajC-Microsoft

infra/scripts/index_scripts/02_process_data.py

@@ -1,13 +1,13 @@
 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 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'
@@ -60,16 +60,14 @@ 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
+    client = EmbeddingsClient(
+        endpoint=f"{openai_api_base}/openai/deployments/{model_id}",
+        credential=credential,
+        credential_scopes=["https://cognitiveservices.azure.com/.default"]
     )
 
-    embedding = client.embeddings.create(input=text, model=model_id).data[0].embedding
+    response = client.embed(input=[text])
+    embedding = response.data[0].embedding
     return embedding
 
 

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

scripts/data_utils.py

@@ -20,6 +20,7 @@
 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.storage.blob import ContainerClient
 from bs4 import BeautifulSoup
@@ -28,7 +29,6 @@
                                      PythonCodeTextSplitter,
                                      RecursiveCharacterTextSplitter,
                                      TextSplitter)
-from openai import AzureOpenAI
 from tqdm import tqdm
 
 # Configure environment variables
@@ -844,22 +844,26 @@ def get_embedding(
             api_version = "2024-02-01"
 
             if azure_credential is not None:
-                api_key = azure_credential.get_token(
-                    "https://cognitiveservices.azure.com/.default"
-                ).token
+                # 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 = AzureOpenAI(
-                api_version=api_version, azure_endpoint=endpoint, api_key=api_key
-            )
-            embeddings = client.embeddings.create(model=deployment_id, input=text)
-
-            return embeddings.model_dump()["data"][0]["embedding"]
+                client = EmbeddingsClient(
+                    endpoint=f"{endpoint}/openai/deployments/{deployment_id}",
+                    credential=AzureKeyCredential(api_key)
+                )
+            
+            response = client.embed(input=[text])
+            return response.data[0].embedding
 
     except Exception as e:
         raise Exception(

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