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 index b6bb3b90a..ee269e973 100644 --- 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 { - 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 index fa8c8792a..2bd7c395a 100644 --- 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.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.toMemory(chat: LLM, scope: Conversation): List { 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.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.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 index 4d1217740..ad1d22b7c 100644 --- 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",