feat(azure): Add HttpRequest.Builder extension methods (#9)

* Added HttpRequest.Builder extension methods * remove unused imports --------- Co-authored-by: Robert Craigie <[email protected]>
openai · Nov 21, 2024 · 097c7c9 · 097c7c9
1 parent 656d3b5
commit 097c7c9
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 30 deletions.
diff --git a/openai-java-core/src/main/kotlin/com/openai/azure/HttpRequestBuilderExtensions.kt b/openai-java-core/src/main/kotlin/com/openai/azure/HttpRequestBuilderExtensions.kt
@@ -0,0 +1,27 @@
+package com.openai.azure
+
+import com.openai.core.ClientOptions
+import com.openai.core.http.HttpRequest
+import com.openai.core.isAzureEndpoint
+import com.openai.credential.BearerTokenCredential
+
+@JvmSynthetic
+internal fun HttpRequest.Builder.addPathSegmentsForAzure(
+    clientOptions: ClientOptions,
+    deploymentModel: String
+): HttpRequest.Builder = apply {
+    if (isAzureEndpoint(clientOptions.baseUrl)) {
+        addPathSegments("openai", "deployments", deploymentModel)
+    }
+}
+
+@JvmSynthetic
+internal fun HttpRequest.Builder.replaceBearerTokenForAzure(
+    clientOptions: ClientOptions
+): HttpRequest.Builder = apply {
+    if (
+        isAzureEndpoint(clientOptions.baseUrl) && clientOptions.credential is BearerTokenCredential
+    ) {
+        replaceHeaders("Authorization", "Bearer ${clientOptions.credential.token()}")
+    }
+}
diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/CompletionServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/CompletionServiceAsyncImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.async
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.JsonValue
 import com.openai.core.RequestOptions
@@ -41,10 +43,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("completions")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/EmbeddingServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/EmbeddingServiceAsyncImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.async
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.RequestOptions
 import com.openai.core.handlers.errorHandler
@@ -35,10 +37,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("embeddings")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/async/ImageServiceAsyncImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/async/ImageServiceAsyncImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.async
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.RequestOptions
 import com.openai.core.handlers.errorHandler
@@ -34,10 +36,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().get().toString())
                 .addPathSegments("images", "generations")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/...ai-java-core/src/main/kotlin/com/openai/services/async/chat/CompletionServiceAsyncImpl.kt b/...ai-java-core/src/main/kotlin/com/openai/services/async/chat/CompletionServiceAsyncImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.async.chat
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.JsonValue
 import com.openai.core.RequestOptions
@@ -17,9 +19,7 @@ import com.openai.core.http.HttpRequest
 import com.openai.core.http.HttpResponse.Handler
 import com.openai.core.http.StreamResponse
 import com.openai.core.http.toAsync
-import com.openai.core.isAzureEndpoint
 import com.openai.core.json
-import com.openai.credential.BearerTokenCredential
 import com.openai.errors.OpenAIError
 import com.openai.models.ChatCompletion
 import com.openai.models.ChatCompletionChunk
@@ -49,23 +49,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
-                .apply {
-                    if (isAzureEndpoint(clientOptions.baseUrl)) {
-                        addPathSegments("openai", "deployments", params.model().toString())
-                    }
-                }
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("chat", "completions")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
-                .apply {
-                    if (
-                        isAzureEndpoint(clientOptions.baseUrl) &&
-                            clientOptions.credential is BearerTokenCredential
-                    ) {
-                        putHeader("Authorization", "Bearer ${clientOptions.credential.token()}")
-                    }
-                }
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/CompletionServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/CompletionServiceImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.blocking
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.JsonValue
 import com.openai.core.RequestOptions
@@ -38,10 +40,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("completions")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/EmbeddingServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/EmbeddingServiceImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.blocking
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.RequestOptions
 import com.openai.core.handlers.errorHandler
@@ -34,10 +36,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("embeddings")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/ImageServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/ImageServiceImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.blocking
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.RequestOptions
 import com.openai.core.handlers.errorHandler
@@ -33,10 +35,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
+                .addPathSegmentsForAzure(clientOptions, params.model().get().toString())
                 .addPathSegments("images", "generations")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()

diff --git a/openai-java-core/src/main/kotlin/com/openai/services/blocking/chat/CompletionServiceImpl.kt b/openai-java-core/src/main/kotlin/com/openai/services/blocking/chat/CompletionServiceImpl.kt
@@ -2,6 +2,8 @@
 
 package com.openai.services.blocking.chat
 
+import com.openai.azure.addPathSegmentsForAzure
+import com.openai.azure.replaceBearerTokenForAzure
 import com.openai.core.ClientOptions
 import com.openai.core.JsonValue
 import com.openai.core.RequestOptions
@@ -15,9 +17,7 @@ import com.openai.core.http.HttpMethod
 import com.openai.core.http.HttpRequest
 import com.openai.core.http.HttpResponse.Handler
 import com.openai.core.http.StreamResponse
-import com.openai.core.isAzureEndpoint
 import com.openai.core.json
-import com.openai.credential.BearerTokenCredential
 import com.openai.errors.OpenAIError
 import com.openai.models.ChatCompletion
 import com.openai.models.ChatCompletionChunk
@@ -46,23 +46,12 @@ constructor(
         val request =
             HttpRequest.builder()
                 .method(HttpMethod.POST)
-                .apply {
-                    if (isAzureEndpoint(clientOptions.baseUrl)) {
-                        addPathSegments("openai", "deployments", params.model().toString())
-                    }
-                }
+                .addPathSegmentsForAzure(clientOptions, params.model().toString())
                 .addPathSegments("chat", "completions")
                 .putAllQueryParams(clientOptions.queryParams)
                 .replaceAllQueryParams(params.getQueryParams())
                 .putAllHeaders(clientOptions.headers)
-                .apply {
-                    if (
-                        isAzureEndpoint(clientOptions.baseUrl) &&
-                            clientOptions.credential is BearerTokenCredential
-                    ) {
-                        putHeader("Authorization", "Bearer ${clientOptions.credential.token()}")
-                    }
-                }
+                .replaceBearerTokenForAzure(clientOptions)
                 .replaceAllHeaders(params.getHeaders())
                 .body(json(clientOptions.jsonMapper, params.getBody()))
                 .build()