Make OpenAiAiAssertionModel httpClient customizable and add timeout s…

…etting
takahirom · Nov 5, 2024 · 80b36ab · 80b36ab
1 parent 36285d7
commit 80b36ab
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/.../takahirom/roborazzi/OpenAiRoborazziAi.kt → ...hirom/roborazzi/OpenAiAiAssertionModel.kt b/.../takahirom/roborazzi/OpenAiRoborazziAi.kt → ...hirom/roborazzi/OpenAiAiAssertionModel.kt
@@ -4,6 +4,8 @@ import com.github.takahirom.roborazzi.AiAssertionOptions.AiAssertionModel.Compan
 import com.github.takahirom.roborazzi.AiAssertionOptions.AiAssertionModel.Companion.DefaultTemperature
 import com.github.takahirom.roborazzi.CaptureResults.Companion.json
 import io.ktor.client.HttpClient
+import io.ktor.client.plugins.HttpTimeout
+import io.ktor.client.plugins.HttpTimeout.Plugin.INFINITE_TIMEOUT_MS
 import io.ktor.client.plugins.contentnegotiation.ContentNegotiation
 import io.ktor.client.plugins.logging.LogLevel
 import io.ktor.client.plugins.logging.Logger
@@ -42,22 +44,26 @@ class OpenAiAiAssertionModel(
   private val seed: Int = 1566,
   private val requestBuilderModifier: (HttpRequestBuilder.() -> Unit) = {
     header("Authorization", "Bearer $apiKey")
-  }
-) : AiAssertionOptions.AiAssertionModel {
+  },
   private val httpClient: HttpClient = HttpClient {
     install(ContentNegotiation) {
       json(
         json = json
       )
     }
+    install(HttpTimeout) {
+      requestTimeoutMillis = INFINITE_TIMEOUT_MS
+      socketTimeoutMillis = 80_000
+    }
     // log
     if (loggingEnabled) {
       install(Logging) {
         logger = Logger.SIMPLE
         level = LogLevel.ALL
       }
     }
-  }
+  },
+) : AiAssertionOptions.AiAssertionModel {
 
   override fun assert(
     referenceImageFilePath: String,
@@ -180,7 +186,7 @@ private fun parseOpenAiResponse(
   val openAiResult = try {
     val element = json.parseToJsonElement(responseText)
     val resultsElement = element.jsonObject["results"]
-    val results = if (resultsElement!=null) {
+    val results = if (resultsElement != null) {
       json.decodeFromJsonElement<List<OpenAiConditionResult>>(resultsElement)
     } else {
       emptyList()