From 29c49a25717467001ba1f5b96489890899f82487 Mon Sep 17 00:00:00 2001
From: Minsoo Cheong <icycle0409@snu.ac.kr>
Date: Sun, 11 Jun 2023 19:27:51 +0900
Subject: [PATCH] fix: make `OpenAiProxyFactory` closable

---
 .../openai/DefaultOkHttpOpenAiClient.kt          | 16 +++++++++++++++-
 .../interfaice/openai/OpenAiApiAdapter.kt        |  2 ++
 .../interfaice/openai/OpenAiProxyFactory.kt      |  7 ++++++-
 3 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/DefaultOkHttpOpenAiClient.kt b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/DefaultOkHttpOpenAiClient.kt
index ed9e36a..7e2234f 100644
--- a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/DefaultOkHttpOpenAiClient.kt
+++ b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/DefaultOkHttpOpenAiClient.kt
@@ -51,18 +51,32 @@ class DefaultOkHttpOpenAiClient : OpenAiApiAdapter {
                 override fun onResponse(response: com.squareup.okhttp.Response) {
                     try {
                         if (!response.isSuccessful) {
-                            throw IOException("Error: ${response.code()} ${response.message()}")
+                            sink.error(IOException("Error: ${response.code()} ${response.message()}"))
                         } else {
                             sink.success(mapper.readValue(response.body().string()))
                         }
                     } catch (e: Exception) {
                         sink.error(e)
+                    } finally {
+                        try {
+                            if (response.body() != null) {
+                                response.body().close()
+                            }
+                        } catch (e: java.lang.Exception) {
+//                            logger.warn("Failed to quietly close Response", e)
+                        }
                     }
                 }
             })
         }
     }
 
+    override fun close() {
+        client.dispatcher?.executorService?.shutdown()
+        client.connectionPool?.evictAll()
+        client.cache?.close()
+    }
+
     companion object {
         val mapper = jacksonObjectMapper()
     }
diff --git a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiApiAdapter.kt b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiApiAdapter.kt
index f5cd4b5..fc4a976 100644
--- a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiApiAdapter.kt
+++ b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiApiAdapter.kt
@@ -19,4 +19,6 @@ interface OpenAiApiAdapter {
     fun getProperties(): OpenAiProperties
     fun setProperties(properties: OpenAiProperties)
     fun chat(request: ChatRequest): Publisher<ChatResponse>
+
+    fun close() {}
 }
diff --git a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiProxyFactory.kt b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiProxyFactory.kt
index d1913a7..af11929 100644
--- a/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiProxyFactory.kt
+++ b/core/src/main/kotlin/io/github/mscheong01/interfaice/openai/OpenAiProxyFactory.kt
@@ -15,12 +15,13 @@ package io.github.mscheong01.interfaice.openai
 
 import io.github.mscheong01.interfaice.AiProxyFactory
 import io.github.mscheong01.interfaice.TranscodingRules
+import java.io.Closeable
 import java.lang.reflect.Proxy
 import java.util.ServiceLoader
 
 class OpenAiProxyFactory(
     private val openAiApiAdapter: OpenAiApiAdapter
-) : AiProxyFactory {
+) : AiProxyFactory, Closeable {
     val customTranscodingRules = mutableListOf<TranscodingRules.CustomRule<*>>()
 
     constructor(
@@ -39,6 +40,10 @@ class OpenAiProxyFactory(
         this.customTranscodingRules.addAll(customTranscodingRules)
     }
 
+    override fun close() {
+        openAiApiAdapter.close()
+    }
+
     companion object {
         private fun getOpenAiApiAdapter(properties: OpenAiProperties): OpenAiApiAdapter {
             val adapters = ServiceLoader.load(OpenAiApiAdapter::class.java)