feat(BE-215): ollama fim api support

- OllamaOpenAIProvider support ollama generate api - support Array type to avoid array to String for AnySerializer
tddworks · Jun 19, 2024 · f1747de · f1747de
1 parent 4a33ccb
commit f1747de
Show file tree

Hide file tree

Showing 5 changed files with 27 additions and 4 deletions.
diff --git a/common/src/commonMain/kotlin/com/tddworks/common/network/api/ktor/api/AnySerializer.kt b/common/src/commonMain/kotlin/com/tddworks/common/network/api/ktor/api/AnySerializer.kt
@@ -27,6 +27,11 @@ object AnySerializer : KSerializer<Any> {
             JsonObject(mapContents)
         }
 
+        is Array<*> -> {
+            val arrayContents = value.map { arrayEntry -> serializeAny(arrayEntry) }
+            JsonArray(arrayContents)
+        }
+
         is List<*> -> {
             val arrayContents = value.map { listEntry -> serializeAny(listEntry) }
             JsonArray(arrayContents)
@@ -53,7 +58,7 @@ object AnySerializer : KSerializer<Any> {
             element.map { deserializeJsonElement(it) }
         }
 
-            is JsonPrimitive -> {
+        is JsonPrimitive -> {
             if (element.isString) element.contentOrNull!!
             else if (element.booleanOrNull != null) element.boolean
             else if (element.intOrNull != null) element.int

diff --git a/...ent-core/src/jvmTest/kotlin/com/tddworks/common/network/api/ktor/api/AnySerializerTest.kt b/...ent-core/src/jvmTest/kotlin/com/tddworks/common/network/api/ktor/api/AnySerializerTest.kt
@@ -97,6 +97,18 @@ class AnySerializerTest {
         assertEquals(AnySerialObject(mapOf("key" to "value")), deserialized)
     }
 
+    @Test
+    fun `should serialize array object`() {
+        // Given
+        val testObject = AnySerialObject(mapOf("key" to arrayOf("value1", "value2")))
+
+        // When
+        val serialized = Json.encodeToString(AnySerialObject.serializer(), testObject)
+
+        // Then
+        assertEquals("{\"any\":{\"key\":[\"value1\",\"value2\"]}}", serialized)
+    }
+
     @Test
     fun `should serialize list object`() {
         // Given

diff --git a/...ateway-core/src/commonMain/kotlin/com/tddworks/openai/gateway/api/OllamaOpenAIProvider.kt b/...ateway-core/src/commonMain/kotlin/com/tddworks/openai/gateway/api/OllamaOpenAIProvider.kt
@@ -11,7 +11,9 @@ import com.tddworks.openai.api.legacy.completions.api.Completion
 import com.tddworks.openai.api.legacy.completions.api.CompletionRequest
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.transform
+import kotlinx.serialization.ExperimentalSerializationApi
 
+@OptIn(ExperimentalSerializationApi::class)
 class OllamaOpenAIProvider(private val client: Ollama) : OpenAIProvider {
     /**
      * Check if the given OpenAIModel is supported by the available models.

diff --git a/...re/src/commonMain/kotlin/com/tddworks/openai/gateway/api/internal/DefaultOpenAIGateway.kt b/...re/src/commonMain/kotlin/com/tddworks/openai/gateway/api/internal/DefaultOpenAIGateway.kt
@@ -65,6 +65,8 @@ class DefaultOpenAIGateway(
      * @return A Completion object containing the completions for the provided request.
      */
     override suspend fun completions(request: CompletionRequest): Completion {
-        return openAI.completions(request)
+        return availableProviders.firstOrNull {
+            it.supports(request.model)
+        }?.completions(request) ?: openAI.completions(request)
     }
 }
diff --git a/...openai-gateway-core/src/jvmTest/kotlin/com/tddworks/ollama/api/chat/api/ExtensionsTest.kt b/...openai-gateway-core/src/jvmTest/kotlin/com/tddworks/ollama/api/chat/api/ExtensionsTest.kt
@@ -24,7 +24,7 @@ class ExtensionsTest {
             maxTokens = 10,
             temperature = 0.5,
             stream = false,
-            stop = "<EOT>",
+            stop = "<EOT>,<EOL>",
             streamOptions = mapOf(
                 "raw" to true
             )
@@ -41,7 +41,9 @@ class ExtensionsTest {
             assertNotNull(this)
             assertEquals(0.5, this?.get("temperature"))
             assertEquals(10, this?.get("num_predict"))
-            assertEquals("<EOT>", (this?.get("stop") as Array<*>)[0])
+            val stops = this?.get("stop") as Array<*>
+            assertEquals(2, stops.size)
+            assertEquals("<EOT>", stops[0])
         }
     }