added a check to order ai and prompt content in memory (#446)

xebia-functional · Sep 22, 2023 · d2a6114 · d2a6114
1 parent db898c1
commit d2a6114
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 18 deletions.
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Chat.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/Chat.kt
@@ -46,7 +46,7 @@ interface Chat : LLM {
 
   @AiDsl
   suspend fun promptMessages(prompt: Prompt, scope: Conversation): List<String> {
-    val requestedMemories = prompt.messages.toMemory(this@Chat, scope)
+    val promptMemories = prompt.messages.toMemory(this@Chat, scope)
     val adaptedPrompt = PromptCalculator.adaptPromptToConversationAndModel(prompt, scope, this@Chat)
 
     val request =
@@ -60,7 +60,7 @@ interface Chat : LLM {
 
     return createChatCompletion(request)
       .choices
-      .addMessagesToMemory(this@Chat, scope, requestedMemories)
+      .addMessagesToMemory(this@Chat, scope, promptMemories)
       .mapNotNull { it.message?.content }
   }
 }
diff --git a/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/MemoryManagement.kt b/core/src/commonMain/kotlin/com/xebia/functional/xef/llm/MemoryManagement.kt
@@ -1,7 +1,11 @@
 package com.xebia.functional.xef.llm
 
 import com.xebia.functional.xef.conversation.Conversation
-import com.xebia.functional.xef.llm.models.chat.*
+import com.xebia.functional.xef.llm.models.chat.Choice
+import com.xebia.functional.xef.llm.models.chat.ChoiceWithFunctions
+import com.xebia.functional.xef.llm.models.chat.Message
+import com.xebia.functional.xef.llm.models.chat.Role
+import com.xebia.functional.xef.store.ConversationId
 import com.xebia.functional.xef.store.Memory
 import io.ktor.util.date.*
 
@@ -12,20 +16,18 @@ internal suspend fun List<Message>.addToMemory(chat: LLM, scope: Conversation) {
   }
 }
 
+internal fun Message.toMemory(cid: ConversationId, chat: LLM, delta: Int = 0): Memory =
+  Memory(
+    conversationId = cid,
+    content = this,
+    timestamp = getTimeMillis() + delta,
+    approxTokens = chat.tokensFromMessages(listOf(this))
+  )
+
 internal fun List<Message>.toMemory(chat: LLM, scope: Conversation): List<Memory> {
   val cid = scope.conversationId
   return if (cid != null) {
-    mapIndexed { delta, it ->
-      Memory(
-        conversationId = cid,
-        content = it,
-        // We are adding the delta to ensure that the timestamp is unique for every message.
-        // With this, we ensure that the order of the messages is preserved.
-        // We assume that the AI response time will be in the order of seconds.
-        timestamp = getTimeMillis() + delta,
-        approxTokens = chat.tokensFromMessages(listOf(it))
-      )
-    }
+    mapIndexed { index, it -> it.toMemory(cid, chat, index) }
   } else emptyList()
 }
 
@@ -47,7 +49,15 @@ internal suspend fun List<ChoiceWithFunctions>.addMessagesToMemory(
         name = role.name
       )
 
-    val newMessages = previousMemories + listOf(firstChoiceMessage).toMemory(chat, scope)
+    // Temporary solution to avoid duplicate timestamps when calling the AI.
+    val aiMemory =
+      firstChoiceMessage.toMemory(cid, chat).let {
+        if (previousMemories.isNotEmpty() && previousMemories.last().timestamp >= it.timestamp) {
+          it.copy(timestamp = previousMemories.last().timestamp + 1)
+        } else it
+      }
+
+    val newMessages = previousMemories + aiMemory
     scope.store.addMemories(newMessages)
   }
 }
@@ -65,7 +75,15 @@ internal suspend fun List<Choice>.addMessagesToMemory(
     val firstChoiceMessage =
       Message(role = role, content = firstChoice.message?.content ?: "", name = role.name)
 
-    val newMessages = previousMemories + listOf(firstChoiceMessage).toMemory(chat, scope)
+    // Temporary solution to avoid duplicate timestamps when calling the AI.
+    val aiMemory =
+      firstChoiceMessage.toMemory(cid, chat).let {
+        if (previousMemories.isNotEmpty() && previousMemories.last().timestamp >= it.timestamp) {
+          it.copy(timestamp = previousMemories.last().timestamp + 1)
+        } else it
+      }
+
+    val newMessages = previousMemories + aiMemory
     scope.store.addMemories(newMessages)
   }
 }
diff --git a/core/src/commonTest/kotlin/com/xebia/functional/xef/data/TestModel.kt b/core/src/commonTest/kotlin/com/xebia/functional/xef/data/TestModel.kt
@@ -7,7 +7,6 @@ import com.xebia.functional.xef.llm.models.chat.*
 import com.xebia.functional.xef.llm.models.embeddings.EmbeddingRequest
 import com.xebia.functional.xef.llm.models.embeddings.EmbeddingResult
 import com.xebia.functional.xef.llm.models.usage.Usage
-import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 
 class TestModel(
@@ -22,7 +21,6 @@ class TestModel(
     request: ChatCompletionRequest
   ): ChatCompletionResponse {
     requests.add(request)
-    delay(100) // Simulating a AI's delay response
     return ChatCompletionResponse(
       id = "fake-id",
       `object` = "fake-object",