From 46a4d4f466e541a39f0e3218e74b02640c594176 Mon Sep 17 00:00:00 2001
From: David Vega Lichacz <7826728+realdavidvega@users.noreply.github.com>
Date: Thu, 22 Aug 2024 18:37:46 +0200
Subject: [PATCH] feat: expose and add default Json serialization to ToolConfig
 (#778)

Co-authored-by: Francisco Diaz <francisco.d@47deg.com>
---
 .../xef/llm/assistants/Assistant.kt           | 20 ++++++++++++++++---
 .../functional/xef/llm/assistants/Tool.kt     |  4 +++-
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Assistant.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Assistant.kt
index 55e6ae86e..484e1492b 100644
--- a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Assistant.kt
+++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Assistant.kt
@@ -1,7 +1,21 @@
 package com.xebia.functional.xef.llm.assistants
 
 import com.xebia.functional.openai.generated.api.Assistants
-import com.xebia.functional.openai.generated.model.*
+import com.xebia.functional.openai.generated.model.AssistantObject
+import com.xebia.functional.openai.generated.model.AssistantObjectToolsInner
+import com.xebia.functional.openai.generated.model.AssistantToolsCode
+import com.xebia.functional.openai.generated.model.AssistantToolsFileSearch
+import com.xebia.functional.openai.generated.model.AssistantToolsFunction
+import com.xebia.functional.openai.generated.model.CreateAssistantRequest
+import com.xebia.functional.openai.generated.model.CreateAssistantRequestModel
+import com.xebia.functional.openai.generated.model.CreateAssistantRequestToolResources
+import com.xebia.functional.openai.generated.model.CreateAssistantRequestToolResourcesCodeInterpreter
+import com.xebia.functional.openai.generated.model.CreateAssistantRequestToolResourcesFileSearch
+import com.xebia.functional.openai.generated.model.FunctionObject
+import com.xebia.functional.openai.generated.model.ModifyAssistantRequest
+import com.xebia.functional.openai.generated.model.ModifyAssistantRequestToolResources
+import com.xebia.functional.openai.generated.model.ModifyAssistantRequestToolResourcesCodeInterpreter
+import com.xebia.functional.openai.generated.model.ModifyAssistantRequestToolResourcesFileSearch
 import com.xebia.functional.xef.Config
 import com.xebia.functional.xef.OpenAI
 import com.xebia.functional.xef.llm.assistants.AssistantThread.Companion.defaultConfig
@@ -51,14 +65,14 @@ class Assistant(
       val toolConfig = toolsConfig.firstOrNull { it.functionObject.name == name }
 
       val toolSerializer = toolConfig?.serializers ?: error("Function $name not registered")
-      val input = config.json.decodeFromString(toolSerializer.inputSerializer, args)
+      val input = toolConfig.json.decodeFromString(toolSerializer.inputSerializer, args)
 
       val tool: Tool<Any?, Any?> = toolConfig.tool as Tool<Any?, Any?>
 
       val schema = buildJsonSchema(toolSerializer.outputSerializer.descriptor)
       val output: Any? = tool(input)
       val result =
-        config.json.encodeToJsonElement(
+        toolConfig.json.encodeToJsonElement(
           toolSerializer.outputSerializer as KSerializer<Any?>,
           output
         )
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Tool.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Tool.kt
index 0fa92b3b5..7373026e7 100644
--- a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Tool.kt
+++ b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/assistants/Tool.kt
@@ -4,6 +4,7 @@ import com.xebia.functional.openai.generated.model.FunctionObject
 import com.xebia.functional.xef.llm.chatFunction
 import com.xebia.functional.xef.llm.defaultFunctionDescription
 import kotlinx.serialization.KSerializer
+import kotlinx.serialization.json.Json
 import kotlinx.serialization.serializer
 
 fun interface Tool<Input, out Output> {
@@ -14,7 +15,8 @@ fun interface Tool<Input, out Output> {
     data class ToolConfig<Input, out Output>(
       val functionObject: FunctionObject,
       val serializers: ToolSerializer,
-      val tool: Tool<Input, Output>
+      val tool: Tool<Input, Output>,
+      val json: Json = Json.Default
     )
 
     data class ToolSerializer(