Mock AI

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,50 @@
+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("completion not implemented")
+  },
+  private val embeddings: (EmbeddingRequest) -> EmbeddingResult = {
+    throw NotImplementedError("completion not implemented")
+  },
+  private val images: (ImagesGenerationRequest) -> ImagesGenerationResponse = {
+    throw NotImplementedError("completion 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 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)
+    }
+  )