From e4742b8f6e32b4dccf9cecb39f498d734a98bfb8 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Thu, 8 Dec 2022 16:53:48 +0000 Subject: [PATCH 01/11] ChatGPT as conversation manager --- main.go | 23 +++-------------------- src/chatgpt/chatgpt.go | 37 ++++++++++++++++++++++++------------- 2 files changed, 27 insertions(+), 33 deletions(-) diff --git a/main.go b/main.go index c88eb68..ac3a825 100644 --- a/main.go +++ b/main.go @@ -68,18 +68,6 @@ func main() { log.Printf("Started Telegram bot! Message @%s to start.", bot.Self.UserName) - userConversations := make(map[int64]Conversation) - - editInterval := 1 * time.Second - if os.Getenv("EDIT_WAIT_SECONDS") != "" { - editSecond, err := strconv.ParseInt(os.Getenv("EDIT_WAIT_SECONDS"), 10, 64) - if err != nil { - log.Printf("Couldn't convert your edit seconds setting into int: %v", err) - editSecond = 1 - } - editInterval = time.Duration(editSecond) * time.Second - } - for update := range updates { if update.Message == nil { continue @@ -98,7 +86,7 @@ func main() { bot.Request(tgbotapi.NewChatAction(update.Message.Chat.ID, "typing")) if !update.Message.IsCommand() { - feed, err := chatGPT.SendMessage(update.Message.Text, userConversations[update.Message.Chat.ID].ConversationID, userConversations[update.Message.Chat.ID].LastMessageID) + feed, err := chatGPT.SendMessage(update.Message.Text, update.Message.Chat.ID) if err != nil { msg.Text = fmt.Sprintf("Error: %v", err) } @@ -109,8 +97,7 @@ func main() { debouncedType := ratelimit.Debounce((10 * time.Second), func() { bot.Request(tgbotapi.NewChatAction(update.Message.Chat.ID, "typing")) }) - - debouncedEdit := ratelimit.DebounceWithArgs(editInterval, func(text interface{}, messageId interface{}) { + debouncedEdit := ratelimit.DebounceWithArgs((1 * time.Second), func(text interface{}, messageId interface{}) { _, err = bot.Request(tgbotapi.EditMessageTextConfig{ BaseEdit: tgbotapi.BaseEdit{ ChatID: msg.ChatID, @@ -139,10 +126,6 @@ func main() { break pollResponse } - userConversations[update.Message.Chat.ID] = Conversation{ - LastMessageID: response.MessageId, - ConversationID: response.ConversationId, - } lastResp = markdown.EnsureFormatting(response.Message) msg.Text = lastResp @@ -183,7 +166,7 @@ func main() { case "start": msg.Text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." case "reload": - userConversations[update.Message.Chat.ID] = Conversation{} + chatGPT.ResetConversation(update.Message.Chat.ID) msg.Text = "Started a new conversation. Enjoy!" default: continue diff --git a/src/chatgpt/chatgpt.go b/src/chatgpt/chatgpt.go index 9bb34ce..0b4b88d 100644 --- a/src/chatgpt/chatgpt.go +++ b/src/chatgpt/chatgpt.go @@ -16,9 +16,15 @@ import ( const KEY_ACCESS_TOKEN = "accessToken" const USER_AGENT = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/107.0.0.0 Safari/537.36" +type Conversation struct { + ID string + LastMessageID string +} + type ChatGPT struct { SessionToken string AccessTokenMap expirymap.ExpiryMap + conversations map[int64]Conversation } type SessionResult struct { @@ -39,15 +45,14 @@ type MessageResponse struct { } type ChatResponse struct { - Message string - MessageId string - ConversationId string + Message string } -func Init(config config.Config) ChatGPT { - return ChatGPT{ +func Init(config config.Config) *ChatGPT { + return &ChatGPT{ AccessTokenMap: expirymap.New(), SessionToken: config.OpenAISession, + conversations: make(map[int64]Conversation), } } @@ -61,8 +66,11 @@ func (c *ChatGPT) EnsureAuth() error { return err } -func (c *ChatGPT) SendMessage(message string, conversationId string, messageId string) (chan ChatResponse, error) { - r := make(chan ChatResponse) +func (c *ChatGPT) ResetConversation(chatID int64) { + c.conversations[chatID] = Conversation{} +} + +func (c *ChatGPT) SendMessage(message string, tgChatID int64) (chan ChatResponse, error) { accessToken, err := c.refreshAccessToken() if err != nil { return nil, errors.New(fmt.Sprintf("Couldn't get access token: %v", err)) @@ -75,11 +83,14 @@ func (c *ChatGPT) SendMessage(message string, conversationId string, messageId s "Authorization": fmt.Sprintf("Bearer %s", accessToken), } - err = client.Connect(message, conversationId, messageId) + convo := c.conversations[tgChatID] + err = client.Connect(message, convo.ID, convo.LastMessageID) if err != nil { return nil, errors.New(fmt.Sprintf("Couldn't connect to ChatGPT: %v", err)) } + r := make(chan ChatResponse) + go func() { defer close(r) mainLoop: @@ -98,11 +109,11 @@ func (c *ChatGPT) SendMessage(message string, conversationId string, messageId s } if len(res.Message.Content.Parts) > 0 { - r <- ChatResponse{ - MessageId: res.Message.ID, - ConversationId: res.ConversationId, - Message: res.Message.Content.Parts[0], - } + convo.ID = res.ConversationId + convo.LastMessageID = res.Message.ID + c.conversations[tgChatID] = convo + + r <- ChatResponse{Message: res.Message.Content.Parts[0]} } } } From 11b48f4efc592413e5499b75499d23a1509f016d Mon Sep 17 00:00:00 2001 From: davidhilman Date: Thu, 8 Dec 2022 17:10:43 +0000 Subject: [PATCH 02/11] fix: sending of error from SendMessage. Separate bot into its own package --- main.go | 123 +++++++++-------------------------------------- src/tgbot/bot.go | 103 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 127 insertions(+), 99 deletions(-) create mode 100644 src/tgbot/bot.go diff --git a/main.go b/main.go index ac3a825..4fe1dc8 100644 --- a/main.go +++ b/main.go @@ -7,22 +7,14 @@ import ( "os/signal" "strconv" "syscall" - "time" - tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" "github.com/joho/godotenv" "github.com/m1guelpf/chatgpt-telegram/src/chatgpt" "github.com/m1guelpf/chatgpt-telegram/src/config" - "github.com/m1guelpf/chatgpt-telegram/src/markdown" - "github.com/m1guelpf/chatgpt-telegram/src/ratelimit" "github.com/m1guelpf/chatgpt-telegram/src/session" + "github.com/m1guelpf/chatgpt-telegram/src/tgbot" ) -type Conversation struct { - ConversationID string - LastMessageID string -} - func main() { config, err := config.Init() if err != nil { @@ -49,7 +41,7 @@ func main() { log.Fatalf("Couldn't load .env file: %v", err) } - bot, err := tgbotapi.NewBotAPI(os.Getenv("TELEGRAM_TOKEN")) + bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN")) if err != nil { log.Fatalf("Couldn't start Telegram bot: %v", err) } @@ -58,123 +50,56 @@ func main() { signal.Notify(c, os.Interrupt, syscall.SIGTERM) go func() { <-c - bot.StopReceivingUpdates() + bot.Stop() os.Exit(0) }() - updateConfig := tgbotapi.NewUpdate(0) - updateConfig.Timeout = 30 - updates := bot.GetUpdatesChan(updateConfig) - - log.Printf("Started Telegram bot! Message @%s to start.", bot.Self.UserName) + log.Printf("Started Telegram bot! Message @%s to start.", bot.Username) - for update := range updates { + for update := range bot.GetUpdatesChan() { if update.Message == nil { continue } - msg := tgbotapi.NewMessage(update.Message.Chat.ID, "") - msg.ReplyToMessageID = update.Message.MessageID - msg.ParseMode = "Markdown" + var ( + updateText = update.Message.Text + updateChatID = update.Message.Chat.ID + updateMessageID = update.Message.MessageID + ) userId := strconv.FormatInt(update.Message.Chat.ID, 10) if os.Getenv("TELEGRAM_ID") != "" && userId != os.Getenv("TELEGRAM_ID") { - msg.Text = "You are not authorized to use this bot." - bot.Send(msg) + bot.Send(updateChatID, updateMessageID, "You are not authorized to use this bot.") continue } - bot.Request(tgbotapi.NewChatAction(update.Message.Chat.ID, "typing")) if !update.Message.IsCommand() { - feed, err := chatGPT.SendMessage(update.Message.Text, update.Message.Chat.ID) - if err != nil { - msg.Text = fmt.Sprintf("Error: %v", err) - } - - var message tgbotapi.Message - var lastResp string - - debouncedType := ratelimit.Debounce((10 * time.Second), func() { - bot.Request(tgbotapi.NewChatAction(update.Message.Chat.ID, "typing")) - }) - debouncedEdit := ratelimit.DebounceWithArgs((1 * time.Second), func(text interface{}, messageId interface{}) { - _, err = bot.Request(tgbotapi.EditMessageTextConfig{ - BaseEdit: tgbotapi.BaseEdit{ - ChatID: msg.ChatID, - MessageID: messageId.(int), - }, - Text: text.(string), - ParseMode: "Markdown", - }) - - if err != nil { - if err.Error() == "Bad Request: message is not modified: specified new message content and reply markup are exactly the same as a current content and reply markup of the message" { - return - } - - log.Printf("Couldn't edit message: %v", err) - } - }) - - pollResponse: - for { - debouncedType() - - select { - case response, ok := <-feed: - if !ok { - break pollResponse - } - - lastResp = markdown.EnsureFormatting(response.Message) - msg.Text = lastResp - - if message.MessageID == 0 { - message, err = bot.Send(msg) - if err != nil { - log.Fatalf("Couldn't send message: %v", err) - } - } else { - debouncedEdit(lastResp, message.MessageID) - } - } - } - - _, err = bot.Request(tgbotapi.EditMessageTextConfig{ - BaseEdit: tgbotapi.BaseEdit{ - ChatID: msg.ChatID, - MessageID: message.MessageID, - }, - Text: lastResp, - ParseMode: "Markdown", - }) + bot.SendTyping(updateChatID) + feed, err := chatGPT.SendMessage(updateText, updateChatID) if err != nil { - if err.Error() == "Bad Request: message is not modified: specified new message content and reply markup are exactly the same as a current content and reply markup of the message" { - continue - } - - log.Printf("Couldn't perform final edit on message: %v", err) + bot.Send(updateChatID, updateMessageID, fmt.Sprintf("Error: %v", err)) + } else { + bot.SendAsLiveOutput(updateChatID, updateMessageID, feed) } - continue } + var text string switch update.Message.Command() { case "help": - msg.Text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." + text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." case "start": - msg.Text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." + text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." case "reload": - chatGPT.ResetConversation(update.Message.Chat.ID) - msg.Text = "Started a new conversation. Enjoy!" + chatGPT.ResetConversation(updateChatID) + text = "Started a new conversation. Enjoy!" default: - continue + text = "Unknown command. Send /help to see a list of commands." } - if _, err := bot.Send(msg); err != nil { - log.Printf("Couldn't send message: %v", err) - continue + if _, err := bot.Send(updateChatID, updateMessageID, text); err != nil { + log.Printf("Error sending message: %v", err) } } } diff --git a/src/tgbot/bot.go b/src/tgbot/bot.go new file mode 100644 index 0000000..c7c0eac --- /dev/null +++ b/src/tgbot/bot.go @@ -0,0 +1,103 @@ +package tgbot + +import ( + "log" + "time" + + tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api/v5" + "github.com/m1guelpf/chatgpt-telegram/src/chatgpt" + "github.com/m1guelpf/chatgpt-telegram/src/markdown" + "github.com/m1guelpf/chatgpt-telegram/src/ratelimit" +) + +type Bot struct { + Username string + api *tgbotapi.BotAPI +} + +func New(token string) (*Bot, error) { + api, err := tgbotapi.NewBotAPI(token) + if err != nil { + return nil, err + } + + return &Bot{ + Username: api.Self.UserName, + api: api, + }, nil +} + +func (b *Bot) GetUpdatesChan() tgbotapi.UpdatesChannel { + cfg := tgbotapi.NewUpdate(0) + cfg.Timeout = 30 + return b.api.GetUpdatesChan(cfg) +} + +func (b *Bot) Stop() { + b.api.StopReceivingUpdates() +} + +func (b *Bot) Send(chatID int64, replyTo int, text string) (tgbotapi.Message, error) { + text = markdown.EnsureFormatting(text) + msg := tgbotapi.NewMessage(chatID, text) + msg.ReplyToMessageID = replyTo + return b.api.Send(msg) +} + +func (b *Bot) SendEdit(chatID int64, messageID int, text string) error { + text = markdown.EnsureFormatting(text) + msg := tgbotapi.NewEditMessageText(chatID, messageID, text) + msg.ParseMode = "Markdown" + if _, err := b.api.Send(msg); err != nil { + if err.Error() == "Bad Request: message is not modified: specified new message content and reply markup are exactly the same as a current content and reply markup of the message" { + return nil + } + return err + } + return nil +} + +func (b *Bot) SendTyping(chatID int64) { + if _, err := b.api.Request(tgbotapi.NewChatAction(chatID, "typing")); err != nil { + log.Printf("Couldn't send typing action: %v", err) + } +} + +func (b *Bot) SendAsLiveOutput(chatID int64, replyTo int, feed chan chatgpt.ChatResponse) { + debouncedType := ratelimit.Debounce(10*time.Second, func() { b.SendTyping(chatID) }) + debouncedEdit := ratelimit.DebounceWithArgs(time.Second, func(text interface{}, messageId interface{}) { + if err := b.SendEdit(chatID, messageId.(int), text.(string)); err != nil { + log.Printf("Couldn't edit message: %v", err) + } + }) + + var message tgbotapi.Message + var lastResp string + +pollResponse: + for { + debouncedType() + + select { + case response, ok := <-feed: + if !ok { + break pollResponse + } + + lastResp = response.Message + + if message.MessageID == 0 { + var err error + if message, err = b.Send(chatID, replyTo, lastResp); err != nil { + log.Fatalf("Couldn't send message: %v", err) + } + } else { + debouncedEdit(lastResp, message.MessageID) + } + } + } + + if err := b.SendEdit(chatID, message.MessageID, lastResp); err != nil { + log.Printf("Couldn't perform final edit on message: %v", err) + } +} From 52d9d3d6cc67f9a67f5afabe332c1137df3f788c Mon Sep 17 00:00:00 2001 From: davidhilman Date: Fri, 9 Dec 2022 17:01:54 +0000 Subject: [PATCH 03/11] Resolve conflicts - add edit interval to tgbot.New --- main.go | 13 ++++++++++++- src/tgbot/bot.go | 14 ++++++++------ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/main.go b/main.go index 4fe1dc8..c7f0fa7 100644 --- a/main.go +++ b/main.go @@ -7,6 +7,7 @@ import ( "os/signal" "strconv" "syscall" + "time" "github.com/joho/godotenv" "github.com/m1guelpf/chatgpt-telegram/src/chatgpt" @@ -41,7 +42,17 @@ func main() { log.Fatalf("Couldn't load .env file: %v", err) } - bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN")) + editInterval := 1 * time.Second + if os.Getenv("EDIT_WAIT_SECONDS") != "" { + editSecond, err := strconv.ParseInt(os.Getenv("EDIT_WAIT_SECONDS"), 10, 64) + if err != nil { + log.Printf("Couldn't convert your edit seconds setting into int: %v", err) + editSecond = 1 + } + editInterval = time.Duration(editSecond) * time.Second + } + + bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN"), editInterval) if err != nil { log.Fatalf("Couldn't start Telegram bot: %v", err) } diff --git a/src/tgbot/bot.go b/src/tgbot/bot.go index c7c0eac..065a0fc 100644 --- a/src/tgbot/bot.go +++ b/src/tgbot/bot.go @@ -11,19 +11,21 @@ import ( ) type Bot struct { - Username string - api *tgbotapi.BotAPI + Username string + api *tgbotapi.BotAPI + editInterval time.Duration } -func New(token string) (*Bot, error) { +func New(token string, editInterval time.Duration) (*Bot, error) { api, err := tgbotapi.NewBotAPI(token) if err != nil { return nil, err } return &Bot{ - Username: api.Self.UserName, - api: api, + Username: api.Self.UserName, + api: api, + editInterval: editInterval, }, nil } @@ -65,7 +67,7 @@ func (b *Bot) SendTyping(chatID int64) { func (b *Bot) SendAsLiveOutput(chatID int64, replyTo int, feed chan chatgpt.ChatResponse) { debouncedType := ratelimit.Debounce(10*time.Second, func() { b.SendTyping(chatID) }) - debouncedEdit := ratelimit.DebounceWithArgs(time.Second, func(text interface{}, messageId interface{}) { + debouncedEdit := ratelimit.DebounceWithArgs(b.editInterval, func(text interface{}, messageId interface{}) { if err := b.SendEdit(chatID, messageId.(int), text.(string)); err != nil { log.Printf("Couldn't edit message: %v", err) } From 598864c0e18ffb8bf43cc63852a88569ab146926 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Fri, 9 Dec 2022 18:16:03 +0000 Subject: [PATCH 04/11] Add env config, load with viper --- go.mod | 4 ++- go.sum | 4 +-- main.go | 31 ++++++----------- src/config/config_test.go | 70 +++++++++++++++++++++++++++++++++++++++ src/config/env_config.go | 27 +++++++++++++++ 5 files changed, 113 insertions(+), 23 deletions(-) create mode 100644 src/config/config_test.go create mode 100644 src/config/env_config.go diff --git a/go.mod b/go.mod index 586700e..e6ef859 100644 --- a/go.mod +++ b/go.mod @@ -5,14 +5,15 @@ go 1.19 require ( github.com/go-telegram-bot-api/telegram-bot-api/v5 v5.5.1 github.com/google/uuid v1.3.0 - github.com/joho/godotenv v1.4.0 github.com/launchdarkly/eventsource v1.7.1 github.com/playwright-community/playwright-go v0.2000.1 github.com/spf13/viper v1.14.0 + github.com/stretchr/testify v1.8.1 ) require ( github.com/danwakefield/fnmatch v0.0.0-20160403171240-cbb64ac3d964 // indirect + github.com/davecgh/go-spew v1.1.1 // indirect github.com/fsnotify/fsnotify v1.6.0 // indirect github.com/go-stack/stack v1.8.1 // indirect github.com/hashicorp/hcl v1.0.0 // indirect @@ -20,6 +21,7 @@ require ( github.com/mitchellh/mapstructure v1.5.0 // indirect github.com/pelletier/go-toml v1.9.5 // indirect github.com/pelletier/go-toml/v2 v2.0.5 // indirect + github.com/pmezard/go-difflib v1.0.0 // indirect github.com/spf13/afero v1.9.2 // indirect github.com/spf13/cast v1.5.0 // indirect github.com/spf13/jwalterweatherman v1.1.0 // indirect diff --git a/go.sum b/go.sum index f45b846..895f011 100644 --- a/go.sum +++ b/go.sum @@ -132,8 +132,6 @@ github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/ianlancetaylor/demangle v0.0.0-20181102032728-5e5cf60278f6/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= github.com/ianlancetaylor/demangle v0.0.0-20200824232613-28f6c0f3b639/go.mod h1:aSSvb/t6k1mPoxDqO4vJh6VOCGPwU4O0C2/Eqndh1Sc= -github.com/joho/godotenv v1.4.0 h1:3l4+N6zfMWnkbPEXKng2o2/MR5mSwTrBih4ZEkkz1lg= -github.com/joho/godotenv v1.4.0/go.mod h1:f4LDr5Voq0i2e/R5DDNOoa2zzDfwtkZa6DnEwAbqwq4= github.com/jstemmer/go-junit-report v0.0.0-20190106144839-af01ea7f8024/go.mod h1:6v2b51hI/fHJwM22ozAgKL4VKDeJcHhJFhtBdhmNjmU= github.com/jstemmer/go-junit-report v0.9.1/go.mod h1:Brl9GWCQeLvo8nXZwPNNblvFj/XSXhF0NWZEnDohbsk= github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck= @@ -176,6 +174,7 @@ github.com/spf13/viper v1.14.0 h1:Rg7d3Lo706X9tHsJMUjdiwMpHB7W8WnSVOssIY+JElU= github.com/spf13/viper v1.14.0/go.mod h1:WT//axPky3FdvXHzGw33dNdXXXfFQqmEalje+egj8As= github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSSt89Yw= +github.com/stretchr/objx v0.5.0/go.mod h1:Yh+to48EsGEfYuaHDzXPcE3xhTkx73EhmCGUpEOglKo= github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4= github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5cxcmMvtA= @@ -184,6 +183,7 @@ github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU= github.com/stretchr/testify v1.8.1 h1:w7B6lhMri9wdJUVmEZPGGhZzrYTPvgJArz7wNPgYKsk= +github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4= github.com/subosito/gotenv v1.4.1 h1:jyEFiXpy21Wm81FBN71l9VoMMV8H8jG+qIK3GCpY6Qs= github.com/subosito/gotenv v1.4.1/go.mod h1:ayKnFf/c6rvx/2iiLrJUk1e6plDbT3edrFNGqEflhK0= github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74= diff --git a/main.go b/main.go index c7f0fa7..e26ce78 100644 --- a/main.go +++ b/main.go @@ -5,11 +5,9 @@ import ( "log" "os" "os/signal" - "strconv" "syscall" "time" - "github.com/joho/godotenv" "github.com/m1guelpf/chatgpt-telegram/src/chatgpt" "github.com/m1guelpf/chatgpt-telegram/src/config" "github.com/m1guelpf/chatgpt-telegram/src/session" @@ -17,42 +15,36 @@ import ( ) func main() { - config, err := config.Init() + persistentConfig, err := config.Init() if err != nil { log.Fatalf("Couldn't load config: %v", err) } - if config.OpenAISession == "" { + if persistentConfig.OpenAISession == "" { session, err := session.GetSession() if err != nil { log.Fatalf("Couldn't get OpenAI session: %v", err) } - err = config.Set("OpenAISession", session) + err = persistentConfig.Set("OpenAISession", session) if err != nil { log.Fatalf("Couldn't save OpenAI session: %v", err) } } - chatGPT := chatgpt.Init(config) + chatGPT := chatgpt.Init(persistentConfig) log.Println("Started ChatGPT") - err = godotenv.Load() + envConfig, err := config.LoadEnvConfig(".env") if err != nil { - log.Fatalf("Couldn't load .env file: %v", err) + log.Fatalf("Couldn't load .env config: %v", err) } - - editInterval := 1 * time.Second - if os.Getenv("EDIT_WAIT_SECONDS") != "" { - editSecond, err := strconv.ParseInt(os.Getenv("EDIT_WAIT_SECONDS"), 10, 64) - if err != nil { - log.Printf("Couldn't convert your edit seconds setting into int: %v", err) - editSecond = 1 - } - editInterval = time.Duration(editSecond) * time.Second + if envConfig.EditWaitSeconds == 0 { + log.Printf("EditWaitSeconds not set, defaulting to 1 second") + envConfig.EditWaitSeconds = 1 } - bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN"), editInterval) + bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN"), time.Duration(envConfig.EditWaitSeconds)) if err != nil { log.Fatalf("Couldn't start Telegram bot: %v", err) } @@ -78,8 +70,7 @@ func main() { updateMessageID = update.Message.MessageID ) - userId := strconv.FormatInt(update.Message.Chat.ID, 10) - if os.Getenv("TELEGRAM_ID") != "" && userId != os.Getenv("TELEGRAM_ID") { + if envConfig.TelegramID != 0 && envConfig.TelegramID != update.Message.Chat.ID { bot.Send(updateChatID, updateMessageID, "You are not authorized to use this bot.") continue } diff --git a/src/config/config_test.go b/src/config/config_test.go new file mode 100644 index 0000000..2f74e2c --- /dev/null +++ b/src/config/config_test.go @@ -0,0 +1,70 @@ +package config + +import ( + "fmt" + "github.com/stretchr/testify/require" + "os" + "testing" +) + +func createFile(name string, content string) (remove func(), err error) { + f, err := os.Create(name) + if err != nil { + return nil, err + } + defer f.Close() + + if _, err := f.WriteString(content); err != nil { + return nil, err + } + + return func() { + if err := os.Remove(name); err != nil { + panic(fmt.Sprintf("failed to remove file: %s", err)) + } + }, nil +} + +func setEnvVariables(vals map[string]string) func() { + for k, v := range vals { + os.Setenv(k, v) + } + return func() { + for k := range vals { + os.Unsetenv(k) + } + } +} + +func TestLoadEnvConfig(t *testing.T) { + const fileName = "test.env" + remove, err := createFile(fileName, `TELEGRAM_ID=123 +TELEGRAM_TOKEN=abc +EDIT_WAIT_SECONDS=10`, + ) + require.NoError(t, err) + defer remove() + + t.Run("should load all values from file", func(t *testing.T) { + cfg, err := LoadEnvConfig(fileName) + require.NoError(t, err) + require.Equal(t, int64(123), cfg.TelegramID) + require.Equal(t, "abc", cfg.TelegramToken) + require.Equal(t, 10, cfg.EditWaitSeconds) + }) + + t.Run("env variables should override file values", func(t *testing.T) { + unset := setEnvVariables(map[string]string{ + "TELEGRAM_ID": "456", + "TELEGRAM_TOKEN": "def", + "EDIT_WAIT_SECONDS": "20", + }) + + cfg, err := LoadEnvConfig(fileName) + unset() + require.NoError(t, err) + require.Equal(t, int64(456), cfg.TelegramID) + require.Equal(t, "def", cfg.TelegramToken) + require.Equal(t, 20, cfg.EditWaitSeconds) + }) +} diff --git a/src/config/env_config.go b/src/config/env_config.go new file mode 100644 index 0000000..780902d --- /dev/null +++ b/src/config/env_config.go @@ -0,0 +1,27 @@ +package config + +import "github.com/spf13/viper" + +type EnvConfig struct { + TelegramID int64 `mapstructure:"TELEGRAM_ID"` + TelegramToken string `mapstructure:"TELEGRAM_TOKEN"` + EditWaitSeconds int `mapstructure:"EDIT_WAIT_SECONDS"` +} + +// LoadEnvConfig loads config from .env file, variables from environment take precedence if provided. +func LoadEnvConfig(path string) (*EnvConfig, error) { + v := viper.New() + v.SetConfigFile(path) + v.SetConfigType("env") + v.AutomaticEnv() + + if err := v.ReadInConfig(); err != nil { + return nil, err + } + + var cfg EnvConfig + if err := v.Unmarshal(&cfg); err != nil { + return nil, err + } + return &cfg, nil +} From e6df94509a927939ac30de01b02976e102187f1b Mon Sep 17 00:00:00 2001 From: davidhilman Date: Fri, 9 Dec 2022 18:25:56 +0000 Subject: [PATCH 05/11] Update config.Init - use viper instance rather than global --- main.go | 9 ++++----- src/chatgpt/chatgpt.go | 2 +- src/config/config.go | 46 +++++++++++++++++++++--------------------- 3 files changed, 28 insertions(+), 29 deletions(-) diff --git a/main.go b/main.go index e26ce78..5eb0858 100644 --- a/main.go +++ b/main.go @@ -15,19 +15,18 @@ import ( ) func main() { - persistentConfig, err := config.Init() + persistentConfig, err := config.LoadOrCreatePersistentConfig() if err != nil { log.Fatalf("Couldn't load config: %v", err) } if persistentConfig.OpenAISession == "" { - session, err := session.GetSession() + token, err := session.GetSession() if err != nil { log.Fatalf("Couldn't get OpenAI session: %v", err) } - err = persistentConfig.Set("OpenAISession", session) - if err != nil { + if err = persistentConfig.SetSessionToken(token); err != nil { log.Fatalf("Couldn't save OpenAI session: %v", err) } } @@ -44,7 +43,7 @@ func main() { envConfig.EditWaitSeconds = 1 } - bot, err := tgbot.New(os.Getenv("TELEGRAM_TOKEN"), time.Duration(envConfig.EditWaitSeconds)) + bot, err := tgbot.New(envConfig.TelegramToken, time.Duration(envConfig.EditWaitSeconds)) if err != nil { log.Fatalf("Couldn't start Telegram bot: %v", err) } diff --git a/src/chatgpt/chatgpt.go b/src/chatgpt/chatgpt.go index 0b4b88d..20a7ad5 100644 --- a/src/chatgpt/chatgpt.go +++ b/src/chatgpt/chatgpt.go @@ -48,7 +48,7 @@ type ChatResponse struct { Message string } -func Init(config config.Config) *ChatGPT { +func Init(config *config.Config) *ChatGPT { return &ChatGPT{ AccessTokenMap: expirymap.New(), SessionToken: config.OpenAISession, diff --git a/src/config/config.go b/src/config/config.go index 3109651..b74e52b 100644 --- a/src/config/config.go +++ b/src/config/config.go @@ -9,46 +9,46 @@ import ( ) type Config struct { + v *viper.Viper + OpenAISession string } -// init tries to read the config from the file, and creates it if it doesn't exist. -func Init() (Config, error) { +// LoadOrCreatePersistentConfig uses the default config directory for the current OS +// to load or create a config file named "chatgpt.json" +func LoadOrCreatePersistentConfig() (*Config, error) { configPath, err := os.UserConfigDir() if err != nil { - return Config{}, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err)) + return nil, errors.New(fmt.Sprintf("Couldn't get user config dir: %v", err)) } - viper.SetConfigType("json") - viper.SetConfigName("chatgpt") - viper.AddConfigPath(configPath) + v := viper.New() + v.SetConfigType("json") + v.SetConfigName("chatgpt") + v.AddConfigPath(configPath) - if err := viper.ReadInConfig(); err != nil { + if err := v.ReadInConfig(); err != nil { if _, ok := err.(viper.ConfigFileNotFoundError); ok { - if err := viper.SafeWriteConfig(); err != nil { - return Config{}, errors.New(fmt.Sprintf("Couldn't create config file: %v", err)) + if err := v.SafeWriteConfig(); err != nil { + return nil, errors.New(fmt.Sprintf("Couldn't create config file: %v", err)) } } else { - return Config{}, errors.New(fmt.Sprintf("Couldn't read config file: %v", err)) + return nil, errors.New(fmt.Sprintf("Couldn't read config file: %v", err)) } } var cfg Config - err = viper.Unmarshal(&cfg) + err = v.Unmarshal(&cfg) if err != nil { - return Config{}, errors.New(fmt.Sprintf("Error parsing config: %v", err)) + return nil, errors.New(fmt.Sprintf("Error parsing config: %v", err)) } + cfg.v = v - return cfg, nil + return &cfg, nil } -// key should be part of the Config struct -func (cfg *Config) Set(key string, value interface{}) error { - viper.Set(key, value) - - err := viper.Unmarshal(&cfg) - if err != nil { - return errors.New(fmt.Sprintf("Error parsing config: %v", err)) - } - - return viper.WriteConfig() +func (cfg *Config) SetSessionToken(token string) error { + // key must match the struct field name + cfg.v.Set("OpenAISession", token) + cfg.OpenAISession = token + return cfg.v.WriteConfig() } From cec0daaf6221af463dc1569544278ea586f99fe9 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Fri, 9 Dec 2022 18:32:50 +0000 Subject: [PATCH 06/11] Allow multiple ids for TELEGRAM_ID env variable --- README.md | 1 + main.go | 7 ++++++- src/config/config_test.go | 6 +++--- src/config/env_config.go | 15 ++++++++++++--- 4 files changed, 22 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index e7d6302..3f63f6b 100644 --- a/README.md +++ b/README.md @@ -18,6 +18,7 @@ After you download the file, extract it into a folder and open the `env.example` - `TELEGRAM_ID` (Optional): Your Telegram User ID - If you set this, only you will be able to interact with the bot. - To get your ID, message `@userinfobot` on Telegram. + - Multiple IDs can be provided, separated by commas. - `EDIT_WAIT_SECONDS` (Optional): Amount of seconds to wait between edits - This is set to `1` by default, but you can increase if you start getting a lot of `Too Many Requests` errors. - Save the file, and rename it to `.env`. diff --git a/main.go b/main.go index 5eb0858..6fa976e 100644 --- a/main.go +++ b/main.go @@ -38,6 +38,9 @@ func main() { if err != nil { log.Fatalf("Couldn't load .env config: %v", err) } + if len(envConfig.TelegramID) == 0 { + log.Printf("Telegram ID not set, all users will be able to use the bot") + } if envConfig.EditWaitSeconds == 0 { log.Printf("EditWaitSeconds not set, defaulting to 1 second") envConfig.EditWaitSeconds = 1 @@ -67,9 +70,11 @@ func main() { updateText = update.Message.Text updateChatID = update.Message.Chat.ID updateMessageID = update.Message.MessageID + updateUserID = update.Message.From.ID ) - if envConfig.TelegramID != 0 && envConfig.TelegramID != update.Message.Chat.ID { + if len(envConfig.TelegramID) != 0 && !envConfig.HasTelegramID(updateUserID) { + log.Printf("User %d is not allowed to use this bot", updateUserID) bot.Send(updateChatID, updateMessageID, "You are not authorized to use this bot.") continue } diff --git a/src/config/config_test.go b/src/config/config_test.go index 2f74e2c..a799da1 100644 --- a/src/config/config_test.go +++ b/src/config/config_test.go @@ -48,14 +48,14 @@ EDIT_WAIT_SECONDS=10`, t.Run("should load all values from file", func(t *testing.T) { cfg, err := LoadEnvConfig(fileName) require.NoError(t, err) - require.Equal(t, int64(123), cfg.TelegramID) + require.Equal(t, []int64{123}, cfg.TelegramID) require.Equal(t, "abc", cfg.TelegramToken) require.Equal(t, 10, cfg.EditWaitSeconds) }) t.Run("env variables should override file values", func(t *testing.T) { unset := setEnvVariables(map[string]string{ - "TELEGRAM_ID": "456", + "TELEGRAM_ID": "456,789", "TELEGRAM_TOKEN": "def", "EDIT_WAIT_SECONDS": "20", }) @@ -63,7 +63,7 @@ EDIT_WAIT_SECONDS=10`, cfg, err := LoadEnvConfig(fileName) unset() require.NoError(t, err) - require.Equal(t, int64(456), cfg.TelegramID) + require.Equal(t, []int64{456, 789}, cfg.TelegramID) require.Equal(t, "def", cfg.TelegramToken) require.Equal(t, 20, cfg.EditWaitSeconds) }) diff --git a/src/config/env_config.go b/src/config/env_config.go index 780902d..15fa5f6 100644 --- a/src/config/env_config.go +++ b/src/config/env_config.go @@ -3,9 +3,18 @@ package config import "github.com/spf13/viper" type EnvConfig struct { - TelegramID int64 `mapstructure:"TELEGRAM_ID"` - TelegramToken string `mapstructure:"TELEGRAM_TOKEN"` - EditWaitSeconds int `mapstructure:"EDIT_WAIT_SECONDS"` + TelegramID []int64 `mapstructure:"TELEGRAM_ID"` + TelegramToken string `mapstructure:"TELEGRAM_TOKEN"` + EditWaitSeconds int `mapstructure:"EDIT_WAIT_SECONDS"` +} + +func (e *EnvConfig) HasTelegramID(id int64) bool { + for _, v := range e.TelegramID { + if v == id { + return true + } + } + return false } // LoadEnvConfig loads config from .env file, variables from environment take precedence if provided. From 9fe6bac22e0961e5f2f20747fa1ec82b3aef72a1 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Sat, 10 Dec 2022 11:02:40 +0000 Subject: [PATCH 07/11] Add tests for multiple ids being provided --- main.go | 8 +-- src/config/config_test.go | 126 ++++++++++++++++++++++++++++++-------- src/config/env_config.go | 34 +++++++++- 3 files changed, 136 insertions(+), 32 deletions(-) diff --git a/main.go b/main.go index 6fa976e..6ecf826 100644 --- a/main.go +++ b/main.go @@ -38,12 +38,8 @@ func main() { if err != nil { log.Fatalf("Couldn't load .env config: %v", err) } - if len(envConfig.TelegramID) == 0 { - log.Printf("Telegram ID not set, all users will be able to use the bot") - } - if envConfig.EditWaitSeconds == 0 { - log.Printf("EditWaitSeconds not set, defaulting to 1 second") - envConfig.EditWaitSeconds = 1 + if err := envConfig.Validate(); err != nil { + log.Fatalf("Invalid .env config: %v", err) } bot, err := tgbot.New(envConfig.TelegramToken, time.Duration(envConfig.EditWaitSeconds)) diff --git a/src/config/config_test.go b/src/config/config_test.go index a799da1..e2ef5ec 100644 --- a/src/config/config_test.go +++ b/src/config/config_test.go @@ -2,9 +2,10 @@ package config import ( "fmt" - "github.com/stretchr/testify/require" "os" "testing" + + "github.com/stretchr/testify/require" ) func createFile(name string, content string) (remove func(), err error) { @@ -37,34 +38,109 @@ func setEnvVariables(vals map[string]string) func() { } func TestLoadEnvConfig(t *testing.T) { - const fileName = "test.env" - remove, err := createFile(fileName, `TELEGRAM_ID=123 + for label, test := range map[string]struct { + fileContent string + envVars map[string]string + want *EnvConfig + }{ + "all values empty in file and env": { + fileContent: `TELEGRAM_ID= +TELEGRAM_TOKEN= +EDIT_WAIT_SECONDS=`, + want: &EnvConfig{ + TelegramID: []int64{}, + TelegramToken: "", + EditWaitSeconds: 0, + }, + }, + "all values provided in file, single TELEGRAM_ID": { + fileContent: `TELEGRAM_ID=123 +TELEGRAM_TOKEN=abc +EDIT_WAIT_SECONDS=10`, + want: &EnvConfig{ + TelegramID: []int64{123}, + TelegramToken: "abc", + EditWaitSeconds: 10, + }, + }, + "multiple TELEGRAM_IDs provided in file": { + fileContent: `TELEGRAM_ID=123,456 TELEGRAM_TOKEN=abc EDIT_WAIT_SECONDS=10`, - ) - require.NoError(t, err) - defer remove() + envVars: map[string]string{}, + want: &EnvConfig{ + TelegramID: []int64{123, 456}, + TelegramToken: "abc", + EditWaitSeconds: 10, + }, + }, + "env variables should override file values": { + fileContent: `TELEGRAM_ID=123 +TELEGRAM_TOKEN=abc +EDIT_WAIT_SECONDS=10`, + envVars: map[string]string{ + "TELEGRAM_ID": "456", + "TELEGRAM_TOKEN": "def", + "EDIT_WAIT_SECONDS": "20", + }, + want: &EnvConfig{ + TelegramID: []int64{456}, + TelegramToken: "def", + EditWaitSeconds: 20, + }, + }, + "multiple TELEGRAM_IDs provided in env": { + fileContent: `TELEGRAM_ID=123 +TELEGRAM_TOKEN=abc +EDIT_WAIT_SECONDS=10`, + envVars: map[string]string{ + "TELEGRAM_ID": "456,789", + }, + want: &EnvConfig{ + TelegramID: []int64{456, 789}, + TelegramToken: "abc", + EditWaitSeconds: 10, + }, + }, + } { + t.Run(label, func(t *testing.T) { + unset := setEnvVariables(test.envVars) + t.Cleanup(unset) - t.Run("should load all values from file", func(t *testing.T) { - cfg, err := LoadEnvConfig(fileName) - require.NoError(t, err) - require.Equal(t, []int64{123}, cfg.TelegramID) - require.Equal(t, "abc", cfg.TelegramToken) - require.Equal(t, 10, cfg.EditWaitSeconds) - }) + remove, err := createFile("test.env", test.fileContent) + require.NoError(t, err) + t.Cleanup(remove) - t.Run("env variables should override file values", func(t *testing.T) { - unset := setEnvVariables(map[string]string{ - "TELEGRAM_ID": "456,789", - "TELEGRAM_TOKEN": "def", - "EDIT_WAIT_SECONDS": "20", + cfg, err := LoadEnvConfig("test.env") + require.NoError(t, err) + require.Equal(t, test.want, cfg) }) + } - cfg, err := LoadEnvConfig(fileName) - unset() - require.NoError(t, err) - require.Equal(t, []int64{456, 789}, cfg.TelegramID) - require.Equal(t, "def", cfg.TelegramToken) - require.Equal(t, 20, cfg.EditWaitSeconds) - }) + //t.Run("all values empty in file", func(t *testing.T) { + // + //}) + // + //t.Run("should load all values from file", func(t *testing.T) { + // cfg, err := LoadEnvConfig(fileName) + // require.NoError(t, err) + // require.Equal(t, []int64{123}, cfg.TelegramID) + // require.Equal(t, "abc", cfg.TelegramToken) + // require.Equal(t, 10, cfg.EditWaitSeconds) + //}) + // + //t.Run("env variables should override file values", func(t *testing.T) { + // unset := setEnvVariables(map[string]string{ + // "TELEGRAM_ID": "456,789", + // "TELEGRAM_TOKEN": "def", + // "EDIT_WAIT_SECONDS": "20", + // }) + // + // cfg, err := LoadEnvConfig(fileName) + // unset() + // require.NoError(t, err) + // require.Equal(t, []int64{456, 789}, cfg.TelegramID) + // require.Equal(t, "def", cfg.TelegramToken) + // require.Equal(t, 20, cfg.EditWaitSeconds) + //}) } diff --git a/src/config/env_config.go b/src/config/env_config.go index 15fa5f6..4daacb2 100644 --- a/src/config/env_config.go +++ b/src/config/env_config.go @@ -1,6 +1,13 @@ package config -import "github.com/spf13/viper" +import ( + "errors" + "fmt" + "log" + "os" + + "github.com/spf13/viper" +) type EnvConfig struct { TelegramID []int64 `mapstructure:"TELEGRAM_ID"` @@ -19,6 +26,10 @@ func (e *EnvConfig) HasTelegramID(id int64) bool { // LoadEnvConfig loads config from .env file, variables from environment take precedence if provided. func LoadEnvConfig(path string) (*EnvConfig, error) { + if !fileExists(path) { + return nil, fmt.Errorf("config file '%s' does not exist", path) + } + v := viper.New() v.SetConfigFile(path) v.SetConfigType("env") @@ -34,3 +45,24 @@ func LoadEnvConfig(path string) (*EnvConfig, error) { } return &cfg, nil } + +func fileExists(path string) bool { + if _, err := os.Stat(path); err != nil { + return os.IsExist(err) + } + return true +} + +func (e *EnvConfig) Validate() error { + if e.TelegramToken == "" { + return errors.New("TELEGRAM_TOKEN is not set") + } + if len(e.TelegramID) == 0 { + log.Printf("TELEGRAM_ID is not set, all users will be able to use the bot") + } + if e.EditWaitSeconds < 0 { + log.Printf("EDIT_WAIT_SECONDS not set, defaulting to 1") + e.EditWaitSeconds = 1 + } + return nil +} From 583db0440a4f86670dd3da0d5608e9bb2a935928 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Sat, 10 Dec 2022 11:06:03 +0000 Subject: [PATCH 08/11] Test action --- .github/workflows/test.yaml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 .github/workflows/test.yaml diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml new file mode 100644 index 0000000..7eea50f --- /dev/null +++ b/.github/workflows/test.yaml @@ -0,0 +1,15 @@ +name: Test +on: [push, pull_request] +jobs: + go-test: + runs-on: ubuntu-latest + steps: + - name: Check out source code + uses: actions/checkout@v3 + - name: Setup + uses: actions/setup-go@v3 + with: + go-version-file: "go.mod" + cache: true + - name: Test + run: go test -v ./... \ No newline at end of file From 2b4018e006e8af61b2453c7ba9588b30a50714d8 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Sat, 10 Dec 2022 11:11:27 +0000 Subject: [PATCH 09/11] Remove commented tests --- src/config/config_test.go | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/src/config/config_test.go b/src/config/config_test.go index e2ef5ec..80ac03c 100644 --- a/src/config/config_test.go +++ b/src/config/config_test.go @@ -116,31 +116,4 @@ EDIT_WAIT_SECONDS=10`, require.Equal(t, test.want, cfg) }) } - - //t.Run("all values empty in file", func(t *testing.T) { - // - //}) - // - //t.Run("should load all values from file", func(t *testing.T) { - // cfg, err := LoadEnvConfig(fileName) - // require.NoError(t, err) - // require.Equal(t, []int64{123}, cfg.TelegramID) - // require.Equal(t, "abc", cfg.TelegramToken) - // require.Equal(t, 10, cfg.EditWaitSeconds) - //}) - // - //t.Run("env variables should override file values", func(t *testing.T) { - // unset := setEnvVariables(map[string]string{ - // "TELEGRAM_ID": "456,789", - // "TELEGRAM_TOKEN": "def", - // "EDIT_WAIT_SECONDS": "20", - // }) - // - // cfg, err := LoadEnvConfig(fileName) - // unset() - // require.NoError(t, err) - // require.Equal(t, []int64{456, 789}, cfg.TelegramID) - // require.Equal(t, "def", cfg.TelegramToken) - // require.Equal(t, 20, cfg.EditWaitSeconds) - //}) } From 8b58cf9fe7b6171491bd8f7e06c1d5f855cf1fd7 Mon Sep 17 00:00:00 2001 From: davidhilman Date: Sat, 10 Dec 2022 17:46:15 +0000 Subject: [PATCH 10/11] Update LoadEnvConfig to work when no file is provided --- main.go | 2 +- src/config/config_test.go | 20 +++++++++++++++++--- src/config/env_config.go | 26 +++++++++++++++++++------- 3 files changed, 37 insertions(+), 11 deletions(-) diff --git a/main.go b/main.go index 6ecf826..ccc1e77 100644 --- a/main.go +++ b/main.go @@ -38,7 +38,7 @@ func main() { if err != nil { log.Fatalf("Couldn't load .env config: %v", err) } - if err := envConfig.Validate(); err != nil { + if err := envConfig.ValidateWithDefaults(); err != nil { log.Fatalf("Invalid .env config: %v", err) } diff --git a/src/config/config_test.go b/src/config/config_test.go index 80ac03c..c6f5888 100644 --- a/src/config/config_test.go +++ b/src/config/config_test.go @@ -53,6 +53,18 @@ EDIT_WAIT_SECONDS=`, EditWaitSeconds: 0, }, }, + "no file, all values through env": { + envVars: map[string]string{ + "TELEGRAM_ID": "123,456", + "TELEGRAM_TOKEN": "token", + "EDIT_WAIT_SECONDS": "10", + }, + want: &EnvConfig{ + TelegramID: []int64{123, 456}, + TelegramToken: "token", + EditWaitSeconds: 10, + }, + }, "all values provided in file, single TELEGRAM_ID": { fileContent: `TELEGRAM_ID=123 TELEGRAM_TOKEN=abc @@ -107,9 +119,11 @@ EDIT_WAIT_SECONDS=10`, unset := setEnvVariables(test.envVars) t.Cleanup(unset) - remove, err := createFile("test.env", test.fileContent) - require.NoError(t, err) - t.Cleanup(remove) + if test.fileContent != "" { + remove, err := createFile("test.env", test.fileContent) + require.NoError(t, err) + t.Cleanup(remove) + } cfg, err := LoadEnvConfig("test.env") require.NoError(t, err) diff --git a/src/config/env_config.go b/src/config/env_config.go index 4daacb2..26a7485 100644 --- a/src/config/env_config.go +++ b/src/config/env_config.go @@ -1,8 +1,8 @@ package config import ( + "bytes" "errors" - "fmt" "log" "os" @@ -15,6 +15,12 @@ type EnvConfig struct { EditWaitSeconds int `mapstructure:"EDIT_WAIT_SECONDS"` } +// emptyConfig is used to initialize viper. +// It is required to register config keys with viper when in case no config file is provided. +const emptyConfig = `TELEGRAM_ID= +TELEGRAM_TOKEN= +EDIT_WAIT_SECONDS=` + func (e *EnvConfig) HasTelegramID(id int64) bool { for _, v := range e.TelegramID { if v == id { @@ -25,19 +31,25 @@ func (e *EnvConfig) HasTelegramID(id int64) bool { } // LoadEnvConfig loads config from .env file, variables from environment take precedence if provided. +// If no .env file is provided, config is loaded from environment variables. func LoadEnvConfig(path string) (*EnvConfig, error) { - if !fileExists(path) { - return nil, fmt.Errorf("config file '%s' does not exist", path) + fileExists := fileExists(path) + if !fileExists { + log.Printf("config file %s does not exist, using env variables", path) } v := viper.New() - v.SetConfigFile(path) v.SetConfigType("env") v.AutomaticEnv() - - if err := v.ReadInConfig(); err != nil { + if err := v.ReadConfig(bytes.NewBufferString(emptyConfig)); err != nil { return nil, err } + if fileExists { + v.SetConfigFile(path) + if err := v.ReadInConfig(); err != nil { + return nil, err + } + } var cfg EnvConfig if err := v.Unmarshal(&cfg); err != nil { @@ -53,7 +65,7 @@ func fileExists(path string) bool { return true } -func (e *EnvConfig) Validate() error { +func (e *EnvConfig) ValidateWithDefaults() error { if e.TelegramToken == "" { return errors.New("TELEGRAM_TOKEN is not set") } From e8b6cc02afd01e5ff6a851311ff4afb2e78bdb6c Mon Sep 17 00:00:00 2001 From: davidhilman Date: Sun, 11 Dec 2022 13:32:20 +0000 Subject: [PATCH 11/11] Update session token via command and manual_auth config variable --- README.md | 3 +++ env.example | 1 + main.go | 39 ++++++++++++++++++++++++++++----------- src/chatgpt/chatgpt.go | 18 +++++++++++++----- src/config/config.go | 4 ++++ src/config/config_test.go | 16 +++++++++++----- src/config/env_config.go | 8 +++++--- 7 files changed, 65 insertions(+), 24 deletions(-) diff --git a/README.md b/README.md index e8f287e..1ec915f 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,9 @@ After you download the file, extract it into a folder and open the `env.example` - Multiple IDs can be provided, separated by commas. - `EDIT_WAIT_SECONDS` (Optional): Amount of seconds to wait between edits - This is set to `1` by default, but you can increase if you start getting a lot of `Too Many Requests` errors. +- `MANUAL_AUTH` (Optional): Setting to true will disable the browser authentication + - Requires setting OpenAI Session Token manually, use `/setToken ` bot command + - See [auth section](#Authentication) for info on how to get your session token - Save the file, and rename it to `.env`. > **Note** Make sure you rename the file to _exactly_ `.env`! The program won't work otherwise. diff --git a/env.example b/env.example index 81e2546..6679680 100644 --- a/env.example +++ b/env.example @@ -1,3 +1,4 @@ TELEGRAM_ID= TELEGRAM_TOKEN= EDIT_WAIT_SECONDS=1 +MANUAL_AUTH=false \ No newline at end of file diff --git a/main.go b/main.go index ccc1e77..79449b3 100644 --- a/main.go +++ b/main.go @@ -14,13 +14,27 @@ import ( "github.com/m1guelpf/chatgpt-telegram/src/tgbot" ) +const ( + messageStart = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." + messageHelp = `/reload - clear chatGPT conversation history (Telegram messages will not be deleted) +/setToken - set the openAI session token` +) + func main() { + envConfig, err := config.LoadEnvConfig(".env") + if err != nil { + log.Fatalf("Couldn't load .env config: %v", err) + } + if err := envConfig.ValidateWithDefaults(); err != nil { + log.Fatalf("Invalid .env config: %v", err) + } + persistentConfig, err := config.LoadOrCreatePersistentConfig() if err != nil { log.Fatalf("Couldn't load config: %v", err) } - if persistentConfig.OpenAISession == "" { + if persistentConfig.OpenAISession == "" && !envConfig.ManualAuth { token, err := session.GetSession() if err != nil { log.Fatalf("Couldn't get OpenAI session: %v", err) @@ -34,14 +48,6 @@ func main() { chatGPT := chatgpt.Init(persistentConfig) log.Println("Started ChatGPT") - envConfig, err := config.LoadEnvConfig(".env") - if err != nil { - log.Fatalf("Couldn't load .env config: %v", err) - } - if err := envConfig.ValidateWithDefaults(); err != nil { - log.Fatalf("Invalid .env config: %v", err) - } - bot, err := tgbot.New(envConfig.TelegramToken, time.Duration(envConfig.EditWaitSeconds)) if err != nil { log.Fatalf("Couldn't start Telegram bot: %v", err) @@ -90,9 +96,20 @@ func main() { var text string switch update.Message.Command() { case "help": - text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." + text = messageHelp case "start": - text = "Send a message to start talking with ChatGPT. You can use /reload at any point to clear the conversation history and start from scratch (don't worry, it won't delete the Telegram messages)." + text = messageStart + case "setToken": + token := update.Message.CommandArguments() + if token == "" { + text = "Please provide a token. Example: /setToken eyJhB..." + break + } + if err := persistentConfig.SetSessionToken(token); err != nil { + text = fmt.Sprintf("Error: %v", err) + break + } + text = "Token set successfully." case "reload": chatGPT.ResetConversation(updateChatID) text = "Started a new conversation. Enjoy!" diff --git a/src/chatgpt/chatgpt.go b/src/chatgpt/chatgpt.go index 20a7ad5..6ca886b 100644 --- a/src/chatgpt/chatgpt.go +++ b/src/chatgpt/chatgpt.go @@ -8,7 +8,6 @@ import ( "net/http" "time" - "github.com/m1guelpf/chatgpt-telegram/src/config" "github.com/m1guelpf/chatgpt-telegram/src/expirymap" "github.com/m1guelpf/chatgpt-telegram/src/sse" ) @@ -21,8 +20,12 @@ type Conversation struct { LastMessageID string } +type Config interface { + GetSessionToken() string +} + type ChatGPT struct { - SessionToken string + cfg Config AccessTokenMap expirymap.ExpiryMap conversations map[int64]Conversation } @@ -48,10 +51,10 @@ type ChatResponse struct { Message string } -func Init(config *config.Config) *ChatGPT { +func Init(config Config) *ChatGPT { return &ChatGPT{ + cfg: config, AccessTokenMap: expirymap.New(), - SessionToken: config.OpenAISession, conversations: make(map[int64]Conversation), } } @@ -123,6 +126,11 @@ func (c *ChatGPT) SendMessage(message string, tgChatID int64) (chan ChatResponse } func (c *ChatGPT) refreshAccessToken() (string, error) { + sessionToken := c.cfg.GetSessionToken() + if sessionToken == "" { + return "", errors.New("no session token, use /setToken command to set") + } + cachedAccessToken, ok := c.AccessTokenMap.Get(KEY_ACCESS_TOKEN) if ok { return cachedAccessToken, nil @@ -134,7 +142,7 @@ func (c *ChatGPT) refreshAccessToken() (string, error) { } req.Header.Set("User-Agent", USER_AGENT) - req.Header.Set("Cookie", fmt.Sprintf("__Secure-next-auth.session-token=%s", c.SessionToken)) + req.Header.Set("Cookie", fmt.Sprintf("__Secure-next-auth.session-token=%s", sessionToken)) res, err := http.DefaultClient.Do(req) if err != nil { diff --git a/src/config/config.go b/src/config/config.go index b74e52b..bcae60a 100644 --- a/src/config/config.go +++ b/src/config/config.go @@ -14,6 +14,10 @@ type Config struct { OpenAISession string } +func (cfg *Config) GetSessionToken() string { + return cfg.OpenAISession +} + // LoadOrCreatePersistentConfig uses the default config directory for the current OS // to load or create a config file named "chatgpt.json" func LoadOrCreatePersistentConfig() (*Config, error) { diff --git a/src/config/config_test.go b/src/config/config_test.go index c6f5888..2772d68 100644 --- a/src/config/config_test.go +++ b/src/config/config_test.go @@ -44,13 +44,12 @@ func TestLoadEnvConfig(t *testing.T) { want *EnvConfig }{ "all values empty in file and env": { - fileContent: `TELEGRAM_ID= -TELEGRAM_TOKEN= -EDIT_WAIT_SECONDS=`, + fileContent: emptyConfig, want: &EnvConfig{ TelegramID: []int64{}, TelegramToken: "", EditWaitSeconds: 0, + ManualAuth: false, }, }, "no file, all values through env": { @@ -58,21 +57,25 @@ EDIT_WAIT_SECONDS=`, "TELEGRAM_ID": "123,456", "TELEGRAM_TOKEN": "token", "EDIT_WAIT_SECONDS": "10", + "MANUAL_AUTH": "true", }, want: &EnvConfig{ TelegramID: []int64{123, 456}, TelegramToken: "token", EditWaitSeconds: 10, + ManualAuth: true, }, }, "all values provided in file, single TELEGRAM_ID": { fileContent: `TELEGRAM_ID=123 TELEGRAM_TOKEN=abc -EDIT_WAIT_SECONDS=10`, +EDIT_WAIT_SECONDS=10 +MANUAL_AUTH=true`, want: &EnvConfig{ TelegramID: []int64{123}, TelegramToken: "abc", EditWaitSeconds: 10, + ManualAuth: true, }, }, "multiple TELEGRAM_IDs provided in file": { @@ -89,16 +92,19 @@ EDIT_WAIT_SECONDS=10`, "env variables should override file values": { fileContent: `TELEGRAM_ID=123 TELEGRAM_TOKEN=abc -EDIT_WAIT_SECONDS=10`, +EDIT_WAIT_SECONDS=10 +MANUAL_AUTH=false`, envVars: map[string]string{ "TELEGRAM_ID": "456", "TELEGRAM_TOKEN": "def", "EDIT_WAIT_SECONDS": "20", + "MANUAL_AUTH": "true", }, want: &EnvConfig{ TelegramID: []int64{456}, TelegramToken: "def", EditWaitSeconds: 20, + ManualAuth: true, }, }, "multiple TELEGRAM_IDs provided in env": { diff --git a/src/config/env_config.go b/src/config/env_config.go index 26a7485..c1c5c96 100644 --- a/src/config/env_config.go +++ b/src/config/env_config.go @@ -13,13 +13,15 @@ type EnvConfig struct { TelegramID []int64 `mapstructure:"TELEGRAM_ID"` TelegramToken string `mapstructure:"TELEGRAM_TOKEN"` EditWaitSeconds int `mapstructure:"EDIT_WAIT_SECONDS"` + ManualAuth bool `mapstructure:"MANUAL_AUTH"` } // emptyConfig is used to initialize viper. // It is required to register config keys with viper when in case no config file is provided. const emptyConfig = `TELEGRAM_ID= TELEGRAM_TOKEN= -EDIT_WAIT_SECONDS=` +EDIT_WAIT_SECONDS= +MANUAL_AUTH=` func (e *EnvConfig) HasTelegramID(id int64) bool { for _, v := range e.TelegramID { @@ -46,7 +48,7 @@ func LoadEnvConfig(path string) (*EnvConfig, error) { } if fileExists { v.SetConfigFile(path) - if err := v.ReadInConfig(); err != nil { + if err := v.MergeInConfig(); err != nil { return nil, err } } @@ -72,7 +74,7 @@ func (e *EnvConfig) ValidateWithDefaults() error { if len(e.TelegramID) == 0 { log.Printf("TELEGRAM_ID is not set, all users will be able to use the bot") } - if e.EditWaitSeconds < 0 { + if e.EditWaitSeconds <= 0 { log.Printf("EDIT_WAIT_SECONDS not set, defaulting to 1") e.EditWaitSeconds = 1 }