[Azure OpenAI] Migrate to Entra ID authentication, for improved security

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAIResponsibleAIIT.java

@@ -1,5 +1,6 @@
 package dev.langchain4j.model.azure;
 
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.data.image.Image;
 import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.data.message.UserMessage;
@@ -40,7 +41,7 @@ void chat_message_should_trigger_content_filter_for_violence(String deploymentNa
 
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -63,7 +64,7 @@ void chat_message_should_trigger_content_filter_for_self_harm(String deploymentN
 
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -82,7 +83,7 @@ void image_should_trigger_content_filter_for_sexual_content() {
 
  AzureOpenAiImageModel model = AzureOpenAiImageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("dall-e-3")
  .logRequestsAndResponses(true)
  .build();
@@ -100,7 +101,7 @@ void language_model_should_trigger_content_filter_for_violence() {
 
  LanguageModel model = AzureOpenAiLanguageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("gpt-35-turbo-instruct")
  .tokenizer(new OpenAiTokenizer(GPT_3_5_TURBO_INSTRUCT))
  .temperature(0.0)
@@ -128,7 +129,7 @@ void streaming_chat_message_should_trigger_content_filter_for_violence(String de
 
  StreamingChatLanguageModel model = AzureOpenAiStreamingChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -174,7 +175,7 @@ void streaming_language_should_trigger_content_filter_for_violence(String deploy
 
  StreamingLanguageModel model = AzureOpenAiStreamingLanguageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("gpt-35-turbo-instruct")
  .tokenizer(new OpenAiTokenizer(GPT_3_5_TURBO_INSTRUCT))
  .temperature(0.0)

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiChatModelIT.java

@@ -3,6 +3,7 @@
 import com.azure.ai.openai.models.ChatCompletionsJsonResponseFormat;
 import com.azure.ai.openai.models.ChatCompletionsResponseFormat;
 import com.azure.core.util.BinaryData;
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
@@ -39,7 +40,7 @@ void should_generate_answer_and_return_token_usage_and_finish_reason_stop(String
 
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -69,7 +70,7 @@ void should_generate_answer_and_return_token_usage_and_finish_reason_length(Stri
 
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .maxTokens(3)
@@ -100,7 +101,7 @@ void should_call_function_with_argument(String deploymentName, String gptVersion
 
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -166,7 +167,7 @@ void should_call_function_with_argument(String deploymentName, String gptVersion
  void should_call_function_with_no_argument(String deploymentName, String gptVersion) {
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -201,7 +202,7 @@ void should_call_function_with_no_argument(String deploymentName, String gptVers
  void should_use_json_format(String deploymentName, String gptVersion) {
  ChatLanguageModel model = AzureOpenAiChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .responseFormat(new ChatCompletionsJsonResponseFormat())

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiEmbeddingModelIT.java

@@ -1,5 +1,6 @@
 package dev.langchain4j.model.azure;
 
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.embedding.EmbeddingModel;
@@ -22,7 +23,7 @@ public class AzureOpenAiEmbeddingModelIT {
 
  EmbeddingModel model = AzureOpenAiEmbeddingModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("text-embedding-ada-002")
  .tokenizer(new OpenAiTokenizer(TEXT_EMBEDDING_ADA_002))
  .logRequestsAndResponses(true)

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiImageModelIT.java

@@ -1,6 +1,7 @@
 package dev.langchain4j.model.azure;
 
 import com.azure.ai.openai.models.ImageGenerationResponseFormat;
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.data.image.Image;
 import dev.langchain4j.model.output.Response;
 import org.junit.jupiter.api.Test;
@@ -23,7 +24,7 @@ void should_generate_image_with_url() {
 
  AzureOpenAiImageModel model = AzureOpenAiImageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("dall-e-3")
  .logRequestsAndResponses(true)
  .build();
@@ -47,7 +48,7 @@ void should_generate_image_with_url() {
  void should_generate_image_in_base64() throws IOException {
  AzureOpenAiImageModel model = AzureOpenAiImageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("dall-e-3")
  .logRequestsAndResponses(false) // The image is big, so we don't want to log it by default
  .responseFormat(ImageGenerationResponseFormat.BASE64.toString())

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiLanguageModelIT.java

@@ -1,5 +1,6 @@
 package dev.langchain4j.model.azure;
 
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.model.language.LanguageModel;
 import dev.langchain4j.model.openai.OpenAiTokenizer;
 import dev.langchain4j.model.output.Response;
@@ -19,7 +20,7 @@ class AzureOpenAiLanguageModelIT {
 
  LanguageModel model = AzureOpenAiLanguageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("gpt-35-turbo-instruct")
  .tokenizer(new OpenAiTokenizer(GPT_3_5_TURBO_INSTRUCT))
  .temperature(0.0)

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiStreamingChatModelIT.java

@@ -1,6 +1,7 @@
 package dev.langchain4j.model.azure;
 
 import com.azure.ai.openai.models.ChatCompletionsJsonResponseFormat;
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
 import dev.langchain4j.agent.tool.ToolSpecification;
 import dev.langchain4j.data.message.AiMessage;
@@ -46,7 +47,7 @@ void should_stream_answer(String deploymentName, String gptVersion) throws Excep
 
  StreamingChatLanguageModel model = AzureOpenAiStreamingChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)
@@ -102,7 +103,7 @@ void should_use_json_format(String deploymentName, String gptVersion) throws Exc
 
  StreamingChatLanguageModel model = AzureOpenAiStreamingChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .responseFormat(new ChatCompletionsJsonResponseFormat())
@@ -164,7 +165,7 @@ void should_return_tool_execution_request(String deploymentName, String gptVersi
 
  StreamingChatLanguageModel model = AzureOpenAiStreamingChatModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName(deploymentName)
  .tokenizer(new OpenAiTokenizer(gptVersion))
  .logRequestsAndResponses(true)

langchain4j-azure-open-ai/src/test/java/dev/langchain4j/model/azure/AzureOpenAiStreamingLanguageModelIT.java

@@ -1,5 +1,6 @@
 package dev.langchain4j.model.azure;
 
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import dev.langchain4j.model.StreamingResponseHandler;
 import dev.langchain4j.model.language.StreamingLanguageModel;
 import dev.langchain4j.model.openai.OpenAiTokenizer;
@@ -22,7 +23,7 @@ class AzureOpenAiStreamingLanguageModelIT {
 
  StreamingLanguageModel model = AzureOpenAiStreamingLanguageModel.builder()
  .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
- .apiKey(System.getenv("AZURE_OPENAI_KEY"))
+ .tokenCredential(new DefaultAzureCredentialBuilder().build())
  .deploymentName("gpt-35-turbo-instruct")
  .tokenizer(new OpenAiTokenizer(GPT_3_5_TURBO_INSTRUCT))
  .temperature(0.0)

langchain4j-azure-open-ai/src/test/script/deploy-azure-openai-models.sh

@@ -5,7 +5,7 @@
 
 echo "Setting up environment variables..."
 echo "----------------------------------"
-PROJECT="langchain4j"
+PROJECT="langchain4j-$RANDOM"
 RESOURCE_GROUP="rg-$PROJECT"
 LOCATION="swedencentral"
 TAG="$PROJECT"
@@ -20,15 +20,36 @@ az group create \
 
 echo "Creating the Cognitive Service..."
 echo "---------------------------------"
-az cognitiveservices account create \
+COGNITIVE_SERVICE_ID=$(az cognitiveservices account create \
  --name "$AI_SERVICE" \
  --resource-group "$RESOURCE_GROUP" \
  --location "$LOCATION" \
  --custom-domain "$AI_SERVICE" \
  --tags system="$TAG" \
  --kind "OpenAI" \
- --sku "S0"
-
+ --sku "S0" \
+ | jq -r ".id")
+
+
+# Security
+# - Disable API Key authentication
+# - Assign a system Managed Identity to the Cognitive Service -> this is for using from other Azure services, like Azure Container Apps
+# - Assign the Contributor role on this resource group to the current user, so he can use the models from the CLI (this is how tests would be normally executed)
+az resource update --ids $COGNITIVE_SERVICE_ID --set properties.disableLocalAuth=true --latest-include-preview
+
+az cognitiveservices account identity assign \
+ --name "$AI_SERVICE" \
+ --resource-group "$RESOURCE_GROUP"
+
+PRINCIPAL_ID=$(az ad signed-in-user show --query id -o tsv)
+SUBSCRIPTION_ID=$(az account show --query id -o tsv)
+
+az role assignment create \
+ --role "5e0bd9bd-7b93-4f28-af87-19fc36ad61bd" \
+ --assignee-object-id "$PRINCIPAL_ID" \
+ --scope /subscriptions/"$SUBSCRIPTION_ID"/resourceGroups/"$RESOURCE_GROUP" \
+ --assignee-principal-type User
+
 # If you want to know the available models, run the following Azure CLI command:
 # az cognitiveservices account list-models --resource-group "$RESOURCE_GROUP" --name "$AI_SERVICE" -o table 
 
@@ -92,20 +113,13 @@ az cognitiveservices account deployment create \
  --sku-capacity 1 \
  --sku-name "Standard"
 
-echo "Storing the key and endpoint in environment variables..."
+echo "Storing endpoint in an environment variable..."
 echo "--------------------------------------------------------"
-AZURE_OPENAI_KEY=$(
- az cognitiveservices account keys list \
- --name "$AI_SERVICE" \
- --resource-group "$RESOURCE_GROUP" \
- | jq -r .key1
- )
 AZURE_OPENAI_ENDPOINT=$(
  az cognitiveservices account show \
  --name "$AI_SERVICE" \
  --resource-group "$RESOURCE_GROUP" \
  | jq -r .properties.endpoint
  )
 
-echo "AZURE_OPENAI_KEY=$AZURE_OPENAI_KEY"
 echo "AZURE_OPENAI_ENDPOINT=$AZURE_OPENAI_ENDPOINT"