From f06f8334396c12f8cdd99c2aa3c7123580361671 Mon Sep 17 00:00:00 2001
From: Antoine Rey <antoine.rey@gmail.com>
Date: Thu, 7 Nov 2024 15:14:07 +0100
Subject: [PATCH] fix #2038 Add custom-headers to the Azure Open AI Spring Boot
 Starter (#50)

Fix issue https://github.com/langchain4j/langchain4j/issues/2038

Is there any markdown documentation file to complete?
Do I need to create an integration test? If so, how?
---
 .../java/dev/langchain4j/azure/openai/spring/AutoConfig.java  | 4 ++++
 .../langchain4j/azure/openai/spring/ChatModelProperties.java  | 2 ++
 .../azure/openai/spring/EmbeddingModelProperties.java         | 3 +++
 .../langchain4j/azure/openai/spring/ImageModelProperties.java | 3 +++
 4 files changed, 12 insertions(+)

diff --git a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/AutoConfig.java b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/AutoConfig.java
index 9a056498..74281776 100644
--- a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/AutoConfig.java
+++ b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/AutoConfig.java
@@ -47,6 +47,7 @@ AzureOpenAiChatModel openAiChatModel(Properties properties) {
                 .timeout(Duration.ofSeconds(chatModelProperties.getTimeout() == null ? 0 : chatModelProperties.getTimeout()))
                 .maxRetries(chatModelProperties.getMaxRetries())
                 .proxyOptions(ProxyOptions.fromConfiguration(Configuration.getGlobalConfiguration()))
+                .customHeaders(chatModelProperties.getCustomHeaders())
                 .logRequestsAndResponses(chatModelProperties.getLogRequestsAndResponses() != null && chatModelProperties.getLogRequestsAndResponses());
         if (chatModelProperties.getNonAzureApiKey() != null) {
             builder.nonAzureApiKey(chatModelProperties.getNonAzureApiKey());
@@ -81,6 +82,7 @@ AzureOpenAiStreamingChatModel openAiStreamingChatModel(Properties properties) {
                 .frequencyPenalty(chatModelProperties.getFrequencyPenalty())
                 .timeout(Duration.ofSeconds(chatModelProperties.getTimeout() == null ? 0 : chatModelProperties.getTimeout()))
                 .proxyOptions(ProxyOptions.fromConfiguration(Configuration.getGlobalConfiguration()))
+                .customHeaders(chatModelProperties.getCustomHeaders())
                 .logRequestsAndResponses(chatModelProperties.getLogRequestsAndResponses() != null && chatModelProperties.getLogRequestsAndResponses());
         if (chatModelProperties.getNonAzureApiKey() != null) {
             builder.nonAzureApiKey(chatModelProperties.getNonAzureApiKey());
@@ -110,6 +112,7 @@ AzureOpenAiEmbeddingModel openAiEmbeddingModel(Properties properties, Tokenizer
                 .tokenizer(tokenizer)
                 .timeout(Duration.ofSeconds(embeddingModelProperties.getTimeout() == null ? 0 : embeddingModelProperties.getTimeout()))
                 .proxyOptions(ProxyOptions.fromConfiguration(Configuration.getGlobalConfiguration()))
+                .customHeaders(embeddingModelProperties.getCustomHeaders())
                 .logRequestsAndResponses(embeddingModelProperties.getLogRequestsAndResponses() != null && embeddingModelProperties.getLogRequestsAndResponses());
 
         if (embeddingModelProperties.getNonAzureApiKey() != null) {
@@ -144,6 +147,7 @@ AzureOpenAiImageModel openAiImageModel(Properties properties) {
                 .timeout(imageModelProperties.getTimeout() == null ? null : Duration.ofSeconds(imageModelProperties.getTimeout()))
                 .maxRetries(imageModelProperties.getMaxRetries())
                 .proxyOptions(ProxyOptions.fromConfiguration(Configuration.getGlobalConfiguration()))
+                .customHeaders(imageModelProperties.getCustomHeaders())
                 .logRequestsAndResponses(imageModelProperties.getLogRequestsAndResponses() != null && imageModelProperties.getLogRequestsAndResponses());
         if (imageModelProperties.getNonAzureApiKey() != null) {
             builder.nonAzureApiKey(imageModelProperties.getNonAzureApiKey());
diff --git a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ChatModelProperties.java b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ChatModelProperties.java
index 2058afd4..5b27eb5c 100644
--- a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ChatModelProperties.java
+++ b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ChatModelProperties.java
@@ -4,6 +4,7 @@
 import lombok.Setter;
 
 import java.util.List;
+import java.util.Map;
 
 @Getter
 @Setter
@@ -25,4 +26,5 @@ class ChatModelProperties {
     Integer timeout;
     Integer maxRetries;
     Boolean logRequestsAndResponses;
+    Map<String, String> customHeaders;
 }
\ No newline at end of file
diff --git a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/EmbeddingModelProperties.java b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/EmbeddingModelProperties.java
index c65cccc4..e4140a31 100644
--- a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/EmbeddingModelProperties.java
+++ b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/EmbeddingModelProperties.java
@@ -3,6 +3,8 @@
 import lombok.Getter;
 import lombok.Setter;
 
+import java.util.Map;
+
 @Getter
 @Setter
 class EmbeddingModelProperties {
@@ -15,4 +17,5 @@ class EmbeddingModelProperties {
     Integer timeout;
     Integer maxRetries;
     Boolean logRequestsAndResponses;
+    Map<String, String> customHeaders;
 }
\ No newline at end of file
diff --git a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ImageModelProperties.java b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ImageModelProperties.java
index 215cd64d..75d2bd28 100644
--- a/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ImageModelProperties.java
+++ b/langchain4j-azure-open-ai-spring-boot-starter/src/main/java/dev/langchain4j/azure/openai/spring/ImageModelProperties.java
@@ -3,6 +3,8 @@
 import lombok.Getter;
 import lombok.Setter;
 
+import java.util.Map;
+
 @Getter
 @Setter
 class ImageModelProperties {
@@ -19,4 +21,5 @@ class ImageModelProperties {
     Integer timeout;
     Integer maxRetries;
     Boolean logRequestsAndResponses;
+    Map<String, String> customHeaders;
 }
\ No newline at end of file