feat: Allow for custom Ktor HttpClientPlugin installations (#239)

Author: rasharab

openai-client/build.gradle.kts

@@ -31,7 +31,7 @@ kotlin {
                 api(libs.coroutines.core)
                 api(libs.okio)
                 implementation(libs.serialization.json)
-                implementation(libs.ktor.client.core)
+                api(libs.ktor.client.core)
                 implementation(libs.ktor.client.logging)
                 implementation(libs.ktor.client.auth)
                 implementation(libs.ktor.client.content.negotiation)

openai-client/src/commonMain/kotlin/com.aallam.openai.client/OpenAI.kt

@@ -4,6 +4,7 @@ import com.aallam.openai.api.http.Timeout
 import com.aallam.openai.client.internal.OpenAIApi
 import com.aallam.openai.client.internal.createHttpClient
 import com.aallam.openai.client.internal.http.HttpTransport
+import io.ktor.client.*
 import kotlin.time.Duration.Companion.seconds
 
 /**
@@ -24,6 +25,7 @@ public interface OpenAI : Completions, Files, Edits, Embeddings, Models, Moderat
  * @param proxy HTTP proxy url
  * @param host OpenAI host configuration.
  * @param retry rate limit retry configuration
+ * @param httpClientConfig additional custom client configuration
  */
 public fun OpenAI(
     token: String,
@@ -34,6 +36,7 @@ public fun OpenAI(
     host: OpenAIHost = OpenAIHost.OpenAI,
     proxy: ProxyConfig? = null,
     retry: RetryStrategy = RetryStrategy(),
+    httpClientConfig: HttpClientConfig<*>.() -> Unit = {}
 ): OpenAI = OpenAI(
     config = OpenAIConfig(
         token = token,
@@ -44,6 +47,7 @@ public fun OpenAI(
         host = host,
         proxy = proxy,
         retry = retry,
+        httpClientConfig = httpClientConfig,
     )
 )
 

openai-client/src/commonMain/kotlin/com.aallam.openai.client/OpenAIConfig.kt

@@ -3,6 +3,7 @@ package com.aallam.openai.client
 import com.aallam.openai.api.http.Timeout
 import com.aallam.openai.api.logging.LogLevel
 import com.aallam.openai.api.logging.Logger
+import io.ktor.client.HttpClientConfig
 import kotlin.time.Duration
 import kotlin.time.Duration.Companion.seconds
 
@@ -18,6 +19,7 @@ import kotlin.time.Duration.Companion.seconds
  * @param proxy HTTP proxy url
  * @param host OpenAI host configuration.
  * @param retry rate limit retry configuration
+ * @param httpClientConfig additional custom client configuration
  */
 public class OpenAIConfig(
     public val token: String,
@@ -28,6 +30,7 @@ public class OpenAIConfig(
     public val host: OpenAIHost = OpenAIHost.OpenAI,
     public val proxy: ProxyConfig? = null,
     public val retry: RetryStrategy = RetryStrategy(),
+    public val httpClientConfig: HttpClientConfig<*>.() -> Unit = {}
 ) {
 
     @Deprecated(
@@ -47,6 +50,7 @@ public class OpenAIConfig(
         host: OpenAIHost = OpenAIHost.OpenAI,
         proxy: ProxyConfig? = null,
         retry: RetryStrategy = RetryStrategy(),
+        httpClientConfig: HttpClientConfig<*>.() -> Unit = {}
     ) : this(
         token = token,
         logging = LoggingConfig(
@@ -59,6 +63,7 @@ public class OpenAIConfig(
         host = host,
         proxy = proxy,
         retry = retry,
+        httpClientConfig = httpClientConfig,
     )
 }
 

openai-client/src/commonMain/kotlin/com.aallam.openai.client/internal/HttpClient.kt

@@ -79,6 +79,8 @@ internal fun createHttpClient(config: OpenAIConfig): HttpClient {
         }
 
         expectSuccess = true
+
+        config.httpClientConfig(this)
     }
 }
 

openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestConfigure.kt

@@ -0,0 +1,34 @@
+package com.aallam.openai.client
+
+import com.aallam.openai.api.http.Timeout
+import io.ktor.client.plugins.api.createClientPlugin
+import kotlin.test.Test
+import kotlinx.coroutines.test.runTest
+import kotlin.test.assertTrue
+import kotlin.time.Duration.Companion.minutes
+
+class TestConfigure : TestOpenAI() {
+    @Test
+    fun configureClientPlugin() = runTest {
+        val responseHeaders = mutableListOf<String>()
+        val plugin = createClientPlugin("CustomHeaderPlugin") {
+            onResponse { response ->
+                response.headers.entries().forEach { entry ->
+                    responseHeaders.add(entry.key)
+                }
+            }
+        }
+        val openAI = generateOpenAI(
+            OpenAIConfig(
+                token = token,
+                timeout = Timeout(socket = 1.minutes)
+            ) {
+                install(plugin)
+            }
+        )
+
+        val resModels = openAI.models()
+        assertTrue { resModels.isNotEmpty() }
+        assertTrue { responseHeaders.isNotEmpty() }
+    }
+}

openai-client/src/commonTest/kotlin/com/aallam/openai/client/TestOpenAI.kt

@@ -8,25 +8,34 @@ import com.aallam.openai.client.internal.env
 import com.aallam.openai.client.internal.http.HttpTransport
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
-import okio.Path.Companion.toPath
-import okio.Source
 import kotlin.time.Duration.Companion.minutes
 
 internal val token: String
     get() = requireNotNull(env("OPENAI_API_KEY")) { "OPENAI_API_KEY environment variable must be set." }
 
-internal val transport = HttpTransport(
-    createHttpClient(
-        OpenAIConfig(
-            token = token,
-            logging = LoggingConfig(logLevel = LogLevel.All),
-            timeout = Timeout(socket = 1.minutes),
+
+internal val openAIConfig: OpenAIConfig = OpenAIConfig(
+    token = token,
+    logging = LoggingConfig(logLevel = LogLevel.All),
+    timeout = Timeout(socket = 1.minutes),
+)
+
+private fun transport(config: OpenAIConfig? = null): HttpTransport {
+    return HttpTransport(
+        createHttpClient(
+            config ?: openAIConfig
         )
     )
-)
+}
 
 abstract class TestOpenAI {
-    internal val openAI = OpenAIApi(transport)
+    internal val openAI = OpenAIApi(transport())
+
+    internal fun generateOpenAI(
+        config: OpenAIConfig
+    ): OpenAIApi {
+        return OpenAIApi(transport(config))
+    }
 
     fun test(testBody: suspend TestScope.() -> Unit) = runTest(timeout = 1.minutes, testBody = testBody)
 }