Azure OpenAI On Your Data pre-GA API Example

pregaaoaioyd.py

# Setup the role assignments from Azure OpenAI system assigned managed identity to Azure search service. 
# Required roles: `Search Index Data Reader` and `Search Service Contributor`.

import os
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential, get_bearer_token_provider

endpoint = os.environ.get("AOAIEndpoint")
deployment = os.environ.get("AOAIDeploymentId")
search_endpoint = os.environ.get("SearchEndpoint")
search_index = os.environ.get("SearchIndex")

token_provider = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default")

client = AzureOpenAI(
    azure_endpoint=endpoint,
    azure_ad_token_provider=token_provider,
    api_version="2023-12-01-preview",
)

completion = client.chat.completions.create(
    model=deployment,
    messages=[
        {
            "role": "user",
            "content": "Who is DRI?",
        },
    ],
    extra_body={
        "data_sources": [
            {
                "type": "AzureCognitiveSearch",
                "parameters": {
                    "endpoint": search_endpoint,
                    "index_name": search_index,
                    "authentication": {
                        "type": "SystemAssignedManagedIdentity"
                    }
                }
            }
        ]
    }
)

print(completion.model_dump_json(indent=2))