feat: telegram add openai config

everpcpc · Apr 22, 2024 · e47eb55 · e47eb55
1 parent 72400a1
commit e47eb55
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 15 deletions.
diff --git a/telegram/checks.go b/telegram/checks.go
@@ -15,7 +15,6 @@ import (
 )
 
 func checkRepeat(b *Bot, message *tgbotapi.Message) {
-
 	key := "tg_last_" + strconv.FormatInt(message.Chat.ID, 10)
 	flattendMsg := strings.TrimSpace(message.Text)
 	defer b.redis.LTrim(key, 0, 10)
@@ -136,7 +135,7 @@ func checkOpenAI(b *Bot, message *tgbotapi.Message) {
 		submessage = submessage.ReplyToMessage
 	}
 	resp, err := b.ai.CreateChatCompletion(context.TODO(), openai.ChatCompletionRequest{
-		Model:       openai.GPT3Dot5Turbo,
+		Model:       b.aiModel,
 		Messages:    chatMessages,
 		Temperature: 0.0,
 	})

diff --git a/telegram/telegram.go b/telegram/telegram.go
@@ -32,14 +32,19 @@ var (
 )
 
 type Config struct {
-	Token          string  `json:"token"`
-	SelfID         int64   `json:"selfID"`
-	AdmissionID    int64   `json:"admissionID"`
-	WhitelistChats []int64 `json:"whitelistChats"`
-	ComicPath      string  `json:"comicPath"`
-	DeleteDelay    string  `json:"deleteDelay"`
-	OpenAIBaseURL  string  `json:"openAIBaseURL"`
-	OpenAIKey      string  `json:"openAIKey"`
+	Token          string        `json:"token"`
+	SelfID         int64         `json:"selfID"`
+	AdmissionID    int64         `json:"admissionID"`
+	WhitelistChats []int64       `json:"whitelistChats"`
+	ComicPath      string        `json:"comicPath"`
+	DeleteDelay    string        `json:"deleteDelay"`
+	OpenAI         *OpenAIConfig `json:"openAIConfig"`
+}
+
+type OpenAIConfig struct {
+	BaseURL string `json:"baseURL"`
+	Key     string `json:"key"`
+	Model   string `json:"model"`
 }
 
 type DownloadPixiv struct {
@@ -65,6 +70,7 @@ type Bot struct {
 
 	pixivBot *pixiv.Bot
 	ai       *openai.Client
+	aiModel  string
 }
 
 func NewBot(cfg *Config) (b *Bot, err error) {
@@ -114,10 +120,11 @@ func (b *Bot) WithMeilisearch(meili *meilisearch.Client) *Bot {
 	return b
 }
 
-func (b *Bot) WithOpenAI(baseURL, key string) *Bot {
-	config := openai.DefaultConfig(key)
-	config.BaseURL = baseURL
+func (b *Bot) WithOpenAI(cfg *OpenAIConfig) *Bot {
+	config := openai.DefaultConfig(cfg.Key)
+	config.BaseURL = cfg.BaseURL
 	b.ai = openai.NewClientWithConfig(config)
+	b.aiModel = cfg.Model
 	return b
 }
 

diff --git a/yubari.go b/yubari.go
@@ -86,9 +86,9 @@ func main() {
 	}
 	telegramBot = telegramBot.WithLogger(logger).WithRedis(redisClient).WithQueue(queue).WithMeilisearch(meili)
 	telegramBot = telegramBot.WithPixiv(pixivBot)
-	if cfg.Telegram.OpenAIKey != "" {
+	if cfg.Telegram.OpenAI != nil {
 		logger.Debug("bot: openai: enabled")
-		telegramBot = telegramBot.WithOpenAI(cfg.Telegram.OpenAIBaseURL, cfg.Telegram.OpenAIKey)
+		telegramBot = telegramBot.WithOpenAI(cfg.Telegram.OpenAI)
 	}
 
 	rssUpdate := make(chan string)