custom openai model and host

kukume · Sep 24, 2024 · 651fbbf · 651fbbf
1 parent 66a4fd2
commit 651fbbf
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 6 deletions.
diff --git a/src/main/kotlin/me/kuku/telegram/entity/BotConfigEntity.kt b/src/main/kotlin/me/kuku/telegram/entity/BotConfigEntity.kt
@@ -21,6 +21,8 @@ class BotConfigEntity {
     var updatePush: Status = Status.OFF
 
     var openaiToken: String = ""
+    var openaiUrl: String = ""
+    var openaiModel: String = ""
 
     fun twoCaptchaKey() = twoCaptchaKey.ifEmpty { null }
 }

diff --git a/src/main/kotlin/me/kuku/telegram/extension/OpenaiExtension.kt b/src/main/kotlin/me/kuku/telegram/extension/OpenaiExtension.kt
@@ -3,6 +3,7 @@ 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.aallam.openai.client.OpenAIHost
 import com.pengrad.telegrambot.model.PhotoSize
 import com.pengrad.telegrambot.model.request.ParseMode
 import com.pengrad.telegrambot.request.EditMessageText
@@ -57,8 +58,10 @@ class OpenaiExtension(
 
             val botConfigEntity = botConfigService.find()
             if (botConfigEntity.openaiToken.ifEmpty { "" }.isEmpty()) error("not setting openai token")
+            val openaiHost = if (botConfigEntity.openaiUrl.isEmpty()) OpenAIHost.OpenAI else OpenAIHost(botConfigEntity.openaiUrl)
+            val openaiModel = botConfigEntity.openaiModel.ifEmpty { "gpt-4o-mini" }
 
-            val openai = OpenAI(botConfigEntity.openaiToken)
+            val openai = OpenAI(botConfigEntity.openaiToken, host = openaiHost)
 
             val chatMessage = ChatMessage(
                 role = ChatRole.User,
@@ -73,7 +76,7 @@ class OpenaiExtension(
             cacheBody.add(chatMessage)
 
             val request = ChatCompletionRequest(
-                model = ModelId("gpt-4o-mini"),
+                model = ModelId(openaiModel),
                 messages = cacheBody,
                 streamOptions = streamOptions {
                     includeUsage = true
@@ -85,7 +88,7 @@ class OpenaiExtension(
                 val sendMessageObject = response.message()
                 val sendMessageId = sendMessageObject.messageId()
                 var openaiText = ""
-                var prefix = ">model: gpt\\-4o\\-mini\n"
+                var prefix = ">model: ${openaiModel.replace("-", "\\-")}\n"
                 var alreadySendText = ""
                 var i = 5
                 openai.chatCompletions(request).onEach {

diff --git a/src/main/kotlin/me/kuku/telegram/extension/PushExtension.kt b/src/main/kotlin/me/kuku/telegram/extension/PushExtension.kt
@@ -166,8 +166,7 @@ class PushBody {
     }
 
     enum class Type {
-        TEXT, IMAGE, VIDEO, FILE
-    }
+        TEXT, IMAGE, VIDEO, FILE    }
 
 
 }
diff --git a/src/main/kotlin/me/kuku/telegram/extension/SettingExtension.kt b/src/main/kotlin/me/kuku/telegram/extension/SettingExtension.kt
@@ -24,13 +24,15 @@ class SettingExtension(
         val clearLog = inlineKeyboardButton("清空日志", "settingsClearLog")
         val updatePush = inlineKeyboardButton("${init().updatePush}github更新推送", "githubUpdatePushSwitch")
         val openai = inlineKeyboardButton("设置openai的token", "openaiToken")
+        val openaiUrl = inlineKeyboardButton("设置openai的url", "openaiUrl")
+        val openaiModel = inlineKeyboardButton("设置openai的model", "openaiModel")
         return InlineKeyboardMarkup(
             arrayOf(blackSetting, adminSetting),
             arrayOf(url),
             arrayOf(twoCaptcha),
             arrayOf(sendLog, clearLog),
             arrayOf(updatePush),
-            arrayOf(openai)
+            arrayOf(openai, openaiUrl, openaiModel)
         )
     }
 
@@ -157,6 +159,22 @@ class SettingExtension(
             botConfigService.save(entity)
             editMessageText("设置openai的token成功")
         }
+        callback("openaiUrl") {
+            editMessageText("请发送openai的url")
+            val url = nextMessage().text()
+            val entity = init()
+            entity.openaiUrl = url
+            botConfigService.save(entity)
+            editMessageText("设置openai的url成功")
+        }
+        callback("openaiModel") {
+            editMessageText("请发送openai的model")
+            val model = nextMessage().text()
+            val entity = init()
+            entity.openaiModel = model
+            botConfigService.save(entity)
+            editMessageText("设置openai的model成功")
+        }
     }
 
     fun TelegramSubscribe.log() {