apply suggestions

Author: Jonas-Isr

foundation-models/openai/src/main/java/com/sap/ai/sdk/foundationmodels/openai/OpenAiClient.java

@@ -28,6 +28,8 @@
 import com.sap.cloud.sdk.cloudplatform.connectivity.Header;
 import com.sap.cloud.sdk.cloudplatform.connectivity.HttpDestination;
 import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.stream.Stream;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
@@ -50,6 +52,7 @@ public final class OpenAiClient {
   @Nullable private String systemPrompt = null;
 
   @Nonnull private final Destination destination;
+  @Nonnull private final List<Header> customHeaders = new ArrayList<>();
 
   /**
    * Create a new OpenAI client for the given foundation model, using the default resource group.
@@ -139,11 +142,10 @@ public OpenAiClient withSystemPrompt(@Nonnull final String systemPrompt) {
   @Beta
   @Nonnull
   public OpenAiClient withHeader(@Nonnull final String key, @Nonnull final String value) {
-    final var newDestination =
-        DefaultHttpDestination.fromDestination(this.destination)
-            .header(new Header(key, value))
-            .build();
-    return new OpenAiClient(newDestination);
+    final var newClient = new OpenAiClient(this.destination);
+    newClient.customHeaders.addAll(this.customHeaders);
+    newClient.customHeaders.add(new Header(key, value));
+    return newClient;
   }
 
   /**
@@ -414,7 +416,7 @@ private <T> T execute(
       @Nonnull final Object payload,
       @Nonnull final Class<T> responseType) {
     final var request = new HttpPost(path);
-    serializeAndSetHttpEntity(request, payload);
+    serializeAndSetHttpEntity(request, payload, this.customHeaders);
     return executeRequest(request, responseType);
   }
 
@@ -424,15 +426,18 @@ private <D extends StreamedDelta> Stream<D> executeStream(
       @Nonnull final Object payload,
       @Nonnull final Class<D> deltaType) {
     final var request = new HttpPost(path);
-    serializeAndSetHttpEntity(request, payload);
+    serializeAndSetHttpEntity(request, payload, this.customHeaders);
     return streamRequest(request, deltaType);
   }
 
   private static void serializeAndSetHttpEntity(
-      @Nonnull final BasicClassicHttpRequest request, @Nonnull final Object payload) {
+      @Nonnull final BasicClassicHttpRequest request,
+      @Nonnull final Object payload,
+      @Nonnull final List<Header> customHeaders) {
     try {
       final var json = JACKSON.writeValueAsString(payload);
       request.setEntity(new StringEntity(json, ContentType.APPLICATION_JSON));
+      customHeaders.forEach(h -> request.addHeader(h.getName(), h.getValue()));
     } catch (final JsonProcessingException e) {
       throw new OpenAiClientException("Failed to serialize request parameters", e)
           .setHttpRequest(request);

foundation-models/openai/src/test/java/com/sap/ai/sdk/foundationmodels/openai/OpenAiClientTest.java

@@ -487,17 +487,24 @@ void testCustomHeaders() {
     stubForChatCompletion();
     final var request =
         new OpenAiChatCompletionRequest("Hello World! Why is this phrase so famous?");
+    final var clientWithHeader = client.withHeader("Header-For-Both", "value");
 
-    final var result = client.withHeader("foo", "bar").chatCompletion(request);
+    final var result = clientWithHeader.withHeader("foo", "bar").chatCompletion(request);
     assertThat(result).isNotNull();
 
     var streamResult =
-        client
+        clientWithHeader
             .withHeader("foot", "baz")
             .streamChatCompletion("Hello World! Why is this phrase so famous?");
     assertThat(streamResult).isNotNull();
 
-    verify(postRequestedFor(anyUrl()).withHeader("foo", equalTo("bar")));
-    verify(postRequestedFor(anyUrl()).withHeader("foot", equalTo("baz")));
+    verify(
+        postRequestedFor(anyUrl())
+            .withHeader("Header-For-Both", equalTo("value"))
+            .withHeader("foo", equalTo("bar")));
+    verify(
+        postRequestedFor(anyUrl())
+            .withHeader("Header-For-Both", equalTo("value"))
+            .withHeader("foot", equalTo("baz")));
   }
 }

orchestration/src/test/java/com/sap/ai/sdk/orchestration/OrchestrationUnitTest.java

@@ -178,14 +178,22 @@ void testCustomHeaders() {
                     .withBodyFile("templatingResponse.json")
                     .withHeader("Content-Type", "application/json")));
 
-    final var result = client.withHeader("foo", "bar").chatCompletion(prompt, config);
+    final var clientWithHeader = client.withHeader("Header-For-Both", "value");
+    final var result = clientWithHeader.withHeader("foo", "bar").chatCompletion(prompt, config);
     assertThat(result).isNotNull();
 
-    var streamResult = client.withHeader("foot", "baz").streamChatCompletion(prompt, config);
+    var streamResult =
+        clientWithHeader.withHeader("foot", "baz").streamChatCompletion(prompt, config);
     assertThat(streamResult).isNotNull();
 
-    verify(postRequestedFor(urlPathEqualTo("/v2/completion")).withHeader("foo", equalTo("bar")));
-    verify(postRequestedFor(urlPathEqualTo("/v2/completion")).withHeader("foot", equalTo("baz")));
+    verify(
+        postRequestedFor(urlPathEqualTo("/v2/completion"))
+            .withHeader("Header-For-Both", equalTo("value"))
+            .withHeader("foo", equalTo("bar")));
+    verify(
+        postRequestedFor(urlPathEqualTo("/v2/completion"))
+            .withHeader("Header-For-Both", equalTo("value"))
+            .withHeader("foot", equalTo("baz")));
   }
 
   @Test