Mock AI (#180)

* Mock AI response

* Two ways to mock

* Fix NotImplementErrors

* Make mock embeddings return 0

* Format code

---------

Co-authored-by: Francisco Diaz <[email protected]>

Author: serras

core/src/commonMain/kotlin/com/xebia/functional/xef/auto/AI.kt

@@ -8,7 +8,9 @@ import com.xebia.functional.xef.embeddings.Embeddings
 import com.xebia.functional.xef.embeddings.OpenAIEmbeddings
 import com.xebia.functional.xef.env.OpenAIConfig
 import com.xebia.functional.xef.llm.openai.KtorOpenAIClient
+import com.xebia.functional.xef.llm.openai.MockOpenAIClient
 import com.xebia.functional.xef.llm.openai.OpenAIClient
+import com.xebia.functional.xef.llm.openai.simpleMockAIClient
 import com.xebia.functional.xef.vectorstores.CombinedVectorStore
 import com.xebia.functional.xef.vectorstores.LocalVectorStore
 import com.xebia.functional.xef.vectorstores.VectorStore
@@ -53,6 +55,21 @@ suspend fun <A> AIScope(block: suspend AIScope.() -> A, orElse: suspend (AIError
     orElse(e)
   }
 
+@OptIn(ExperimentalTime::class)
+suspend fun <A> MockAIScope(
+  mockClient: MockOpenAIClient,
+  block: suspend AIScope.() -> A,
+  orElse: suspend (AIError) -> A
+): A =
+  try {
+    val embeddings = OpenAIEmbeddings(OpenAIConfig(), mockClient)
+    val vectorStore = LocalVectorStore(embeddings)
+    val scope = AIScope(mockClient, vectorStore, embeddings)
+    block(scope)
+  } catch (e: AIError) {
+    orElse(e)
+  }
+
 /**
  * Run the [AI] value to produce _either_ an [AIError], or [A]. this method initialises all the
  * dependencies required to run the [AI] value and once it finishes it closes all the resources.
@@ -64,6 +81,20 @@ suspend fun <A> AIScope(block: suspend AIScope.() -> A, orElse: suspend (AIError
 suspend inline fun <reified A> AI<A>.toEither(): Either<AIError, A> =
   ai { invoke().right() }.getOrElse { it.left() }
 
+/**
+ * Run the [AI] value to produce _either_ an [AIError], or [A]. This method uses the [mockAI] to
+ * compute the different responses.
+ */
+suspend fun <A> AI<A>.mock(mockAI: MockOpenAIClient): Either<AIError, A> =
+  MockAIScope(mockAI, { invoke().right() }, { it.left() })
+
+/**
+ * Run the [AI] value to produce _either_ an [AIError], or [A]. This method uses the [mockAI] to
+ * compute the different responses.
+ */
+suspend fun <A> AI<A>.mock(mockAI: (String) -> String): Either<AIError, A> =
+  MockAIScope(simpleMockAIClient(mockAI), { invoke().right() }, { it.left() })
+
 /**
  * Run the [AI] value to produce [A]. this method initialises all the dependencies required to run
  * the [AI] value and once it finishes it closes all the resources.

core/src/commonMain/kotlin/com/xebia/functional/xef/llm/openai/MockAIClient.kt

@@ -0,0 +1,56 @@
+package com.xebia.functional.xef.llm.openai
+
+class MockOpenAIClient(
+  private val completion: (CompletionRequest) -> CompletionResult = {
+    throw NotImplementedError("completion not implemented")
+  },
+  private val chatCompletion: (ChatCompletionRequest) -> ChatCompletionResponse = {
+    throw NotImplementedError("chat completion not implemented")
+  },
+  private val embeddings: (EmbeddingRequest) -> EmbeddingResult = ::nullEmbeddings,
+  private val images: (ImagesGenerationRequest) -> ImagesGenerationResponse = {
+    throw NotImplementedError("images not implemented")
+  },
+) : OpenAIClient {
+  override suspend fun createCompletion(request: CompletionRequest): CompletionResult =
+    completion(request)
+
+  override suspend fun createChatCompletion(
+    request: ChatCompletionRequest
+  ): ChatCompletionResponse = chatCompletion(request)
+
+  override suspend fun createEmbeddings(request: EmbeddingRequest): EmbeddingResult =
+    embeddings(request)
+
+  override suspend fun createImages(request: ImagesGenerationRequest): ImagesGenerationResponse =
+    images(request)
+}
+
+fun nullEmbeddings(request: EmbeddingRequest): EmbeddingResult {
+  val results = request.input.mapIndexed { index, s -> Embedding(s, listOf(0F), index) }
+  return EmbeddingResult(request.model, "", results, Usage.ZERO)
+}
+
+fun simpleMockAIClient(execute: (String) -> String): MockOpenAIClient =
+  MockOpenAIClient(
+    completion = { req ->
+      val request = "${req.prompt.orEmpty()} ${req.suffix.orEmpty()}"
+      val response = execute(request)
+      val result = CompletionChoice(response, 0, null, "end")
+      val requestTokens = request.split(' ').size.toLong()
+      val responseTokens = response.split(' ').size.toLong()
+      val usage = Usage(requestTokens, responseTokens, requestTokens + responseTokens)
+      CompletionResult("FakeID123", "", 0, req.model, listOf(result), usage)
+    },
+    chatCompletion = { req ->
+      val responses =
+        req.messages.mapIndexed { ix, msg ->
+          val response = execute(msg.content)
+          Choice(Message(msg.role, response), "end", ix)
+        }
+      val requestTokens = req.messages.sumOf { it.content.split(' ').size.toLong() }
+      val responseTokens = responses.sumOf { it.message.content.split(' ').size.toLong() }
+      val usage = Usage(requestTokens, responseTokens, requestTokens + responseTokens)
+      ChatCompletionResponse("FakeID123", "", 0, req.model, usage, responses)
+    }
+  )

core/src/commonMain/kotlin/com/xebia/functional/xef/llm/openai/models.kt

@@ -114,7 +114,11 @@ data class Usage(
   @SerialName("prompt_tokens") val promptTokens: Long,
   @SerialName("completion_tokens") val completionTokens: Long? = null,
   @SerialName("total_tokens") val totalTokens: Long
-)
+) {
+  companion object {
+    val ZERO: Usage = Usage(0, 0, 0)
+  }
+}
 
 data class LLMModel(val name: String, val kind: Kind, val modelType: ModelType) {
   enum class Kind {