chatgpt context

kukume · Sep 6, 2024 · 66a4fd2 · 66a4fd2
1 parent 75c8171
commit 66a4fd2
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 18 deletions.
diff --git a/build.gradle.kts b/build.gradle.kts
@@ -19,7 +19,6 @@ repositories {
 
 dependencies {
     implementation("org.jetbrains.kotlinx:kotlinx-coroutines-reactor")
-    implementation("io.projectreactor.kotlin:reactor-kotlin-extensions")
     implementation("org.springframework.boot:spring-boot-starter-data-mongodb-reactive")
     implementation("org.springframework.boot:spring-boot-starter-mail")
     implementation("com.github.pengrad:java-telegram-bot-api:7.9.1")

diff --git a/src/main/kotlin/me/kuku/telegram/extension/OpenaiExtension.kt b/src/main/kotlin/me/kuku/telegram/extension/OpenaiExtension.kt
@@ -3,13 +3,10 @@ package me.kuku.telegram.extension
 import com.aallam.openai.api.chat.*
 import com.aallam.openai.api.model.ModelId
 import com.aallam.openai.client.OpenAI
-import com.fasterxml.jackson.databind.JsonNode
 import com.pengrad.telegrambot.model.PhotoSize
 import com.pengrad.telegrambot.model.request.ParseMode
 import com.pengrad.telegrambot.request.EditMessageText
 import com.pengrad.telegrambot.request.GetFile
-import io.ktor.client.call.*
-import io.ktor.client.request.*
 import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onEach
@@ -19,17 +16,19 @@ import me.kuku.telegram.context.Locality
 import me.kuku.telegram.context.asyncExecute
 import me.kuku.telegram.context.byteArray
 import me.kuku.telegram.entity.BotConfigService
-import me.kuku.utils.Jackson
+import me.kuku.telegram.utils.CacheManager
 import me.kuku.utils.base64Encode
-import me.kuku.utils.client
-import me.kuku.utils.setJsonBody
 import org.springframework.stereotype.Component
+import java.time.Duration
 
 @Component
 class OpenaiExtension(
     private val botConfigService: BotConfigService
 ) {
 
+    private val cache = CacheManager.getCache<String, MutableList<ChatMessage>>("gpt-chat-context", Duration.ofMinutes(2))
+
+
     fun AbilitySubscriber.openai() {
 
         sub(name = "chat", locality = Locality.ALL) {
@@ -51,24 +50,31 @@ class OpenaiExtension(
                 val base64 = getFileResponse.file().byteArray().base64Encode()
                 photoList.add(base64)
             }
+
+            val key = chatId.toString() + message.from().id()
+
+            val cacheBody = cache[key] ?: mutableListOf()
+
             val botConfigEntity = botConfigService.find()
             if (botConfigEntity.openaiToken.ifEmpty { "" }.isEmpty()) error("not setting openai token")
 
             val openai = OpenAI(botConfigEntity.openaiToken)
 
+            val chatMessage = ChatMessage(
+                role = ChatRole.User,
+                content = ContentPartBuilder().also {
+                    it.text(text)
+                    for (photo in photoList) {
+                        it.image("data:image/jpeg;base64,$photo")
+                    }
+                }.build()
+            )
+
+            cacheBody.add(chatMessage)
+
             val request = ChatCompletionRequest(
                 model = ModelId("gpt-4o-mini"),
-                messages = listOf(
-                    ChatMessage(
-                        role = ChatRole.User,
-                        content = ContentPartBuilder().also {
-                            it.text(text)
-                            for (photo in photoList) {
-                                it.image("data:image/jpeg;base64,$photo")
-                            }
-                        }.build()
-                    )
-                ),
+                messages = cacheBody,
                 streamOptions = streamOptions {
                     includeUsage = true
                 }
@@ -100,6 +106,8 @@ class OpenaiExtension(
                     }
                 }.onCompletion {
                     val sendText = "$prefix\n```text\n$openaiText```"
+                    cacheBody.add(ChatMessage(ChatRole.Assistant, openaiText))
+                    cache[key] = cacheBody
                     if (alreadySendText != sendText) {
                         alreadySendText = sendText
                         val editMessageText = EditMessageText(chatId, sendMessageId, sendText)

diff --git a/src/main/kotlin/me/kuku/telegram/utils/CacheManager.kt b/src/main/kotlin/me/kuku/telegram/utils/CacheManager.kt
@@ -127,6 +127,8 @@ open class Cache<K, V>: Serializable {
         }
     }
 
+    operator fun set(key: K, value: V) = put(key, value)
+
     fun put(key: K, value: V) {
         if (expire != null && expire!! > 0) {
             put(key, value, expire!!)