From 6cfa0c21b21dbcc667fde0d86865bf5e7abbd328 Mon Sep 17 00:00:00 2001 From: Sid Sun Date: Sun, 26 May 2024 22:26:09 +0530 Subject: [PATCH] Add edit support and enhance models command Signed-off-by: Sid Sun --- pkg/bot/handlers/completion/completion.go | 20 ++++++++------------ pkg/bot/handlers/completion/util.go | 9 +++++++++ pkg/bot/handlers/models/models.go | 8 ++------ pkg/bot/handlers/models/utils.go | 19 ++++++++++--------- pkg/bot/router/router.go | 5 +++-- 5 files changed, 32 insertions(+), 29 deletions(-) create mode 100644 pkg/bot/handlers/completion/util.go diff --git a/pkg/bot/handlers/completion/completion.go b/pkg/bot/handlers/completion/completion.go index 087d223..d96a554 100644 --- a/pkg/bot/handlers/completion/completion.go +++ b/pkg/bot/handlers/completion/completion.go @@ -15,13 +15,13 @@ import ( var logger = slog.Default().With(slog.String("package", "Completion")) // Handler handles all repeat requests -func Handler(b *tele.Bot, isResend bool) tele.HandlerFunc { +func Handler(b *tele.Bot, handlerMode HandlerMode) tele.HandlerFunc { return func(c tele.Context) error { - logger.Info("[Completion] [Attempt]", slog.Bool("is_resend", isResend)) + logger.Info("[Completion] [Attempt]", slog.Any("handlerMode", handlerMode)) // Check if this is a resend attempt - promptID := c.Message().ID - if isResend { + message := c.Message() + if handlerMode == RegenerateCompletion { if c.Message().ReplyTo == nil { c.Send("Reply to the message you want to regenerate from") return nil @@ -35,15 +35,14 @@ func Handler(b *tele.Bot, isResend bool) tele.HandlerFunc { c.Send("Last message for resend must be a user message") return nil } - promptID = c.Message().ReplyTo.ID + message = c.Message().ReplyTo } + promptID := message.ID // notify user we are processing c.Notify(tele.Typing) - if !isResend { - addMessageToChain(c.Message()) - } + addMessageToChain(message) updatesChan := make(chan contract.CompletionUpdate, 100) go func() { @@ -55,10 +54,7 @@ func Handler(b *tele.Bot, isResend bool) tele.HandlerFunc { firstCompletion := <-updatesChan - replyTo := c.Message() - if isResend { - replyTo = c.Message().ReplyTo - } + replyTo := message botMessage, err := b.Send(c.Chat(), firstCompletion.Message, &tele.SendOptions{ ReplyTo: replyTo, diff --git a/pkg/bot/handlers/completion/util.go b/pkg/bot/handlers/completion/util.go new file mode 100644 index 0000000..dc66126 --- /dev/null +++ b/pkg/bot/handlers/completion/util.go @@ -0,0 +1,9 @@ +package completion + +type HandlerMode int + +const ( + NewCompletion HandlerMode = iota // Initialize the first value with iota + EditCompletion + RegenerateCompletion +) diff --git a/pkg/bot/handlers/models/models.go b/pkg/bot/handlers/models/models.go index 0be7a00..880c669 100644 --- a/pkg/bot/handlers/models/models.go +++ b/pkg/bot/handlers/models/models.go @@ -17,9 +17,7 @@ func GetModelsHandler(b *tele.Bot) tele.HandlerFunc { modelInfoMessage, modelOptions := getInlineKeyboardMarkup(store.ModelStore[c.Chat().ID]) mkp := b.NewMarkup() - mkp.InlineKeyboard = [][]tele.InlineButton{ - modelOptions, - } + mkp.InlineKeyboard = modelOptions b.Send(c.Chat(), modelInfoMessage, mkp) @@ -40,9 +38,7 @@ func CallbackHandler(b *tele.Bot) tele.HandlerFunc { store.ModelStore[c.Chat().ID] = model modelInfoMessage, modelOptions := getInlineKeyboardMarkup(store.ModelStore[c.Chat().ID]) mkp := b.NewMarkup() - mkp.InlineKeyboard = [][]tele.InlineButton{ - modelOptions, - } + mkp.InlineKeyboard = modelOptions _, err := b.Edit(c.Callback().Message, modelInfoMessage, mkp) if err != nil { logger.Error("[Models] [Callback] [Error]", slog.String("error", err.Error())) diff --git a/pkg/bot/handlers/models/utils.go b/pkg/bot/handlers/models/utils.go index a52681d..067d744 100644 --- a/pkg/bot/handlers/models/utils.go +++ b/pkg/bot/handlers/models/utils.go @@ -7,22 +7,23 @@ import ( tele "gopkg.in/telebot.v3" ) -func getInlineKeyboardMarkup(currentModel string) (string, []tele.InlineButton) { - var modelOptions []tele.InlineButton - modelInfoMessage := "Here are the available models: \n\n" +func getInlineKeyboardMarkup(currentModel string) (string, [][]tele.InlineButton) { + var modelOptions [][]tele.InlineButton + modelInfoMessage := "Here are the available models: \n" if currentModel == "" { currentModel = "default" } for _, modelName := range config.GlobalConfig.ModelNames { options := config.GlobalConfig.Models[modelName] - text := modelName + text := fmt.Sprintf("%s (%s) - %d", modelName, options.Model, options.Tweaks.ContextLength) if currentModel == modelName { - text = fmt.Sprintf("*%s*", modelName) + text = fmt.Sprintf("*%s* (%s) - %d", modelName, options.Model, options.Tweaks.ContextLength) } - modelInfoMessage += fmt.Sprintf("%s (%s) - %d\n", modelName, options.Model, options.Tweaks.ContextLength) - modelOptions = append(modelOptions, tele.InlineButton{ - Data: "model_" + modelName, - Text: text, + modelOptions = append(modelOptions, []tele.InlineButton{ + { + Data: "model_" + modelName, + Text: text, + }, }) } return modelInfoMessage, modelOptions diff --git a/pkg/bot/router/router.go b/pkg/bot/router/router.go index 10cb83d..9161dc6 100644 --- a/pkg/bot/router/router.go +++ b/pkg/bot/router/router.go @@ -31,8 +31,9 @@ func (b Bot) Start() { callbackGroup.Handle("/models", models.GetModelsHandler(b.bot)) callbackGroup.Handle(tele.OnCallback, models.CallbackHandler(b.bot)) // Add all other handlers - b.bot.Handle("/resend", completion.Handler(b.bot, true)) - b.bot.Handle(tele.OnText, completion.Handler(b.bot, false)) + b.bot.Handle("/resend", completion.Handler(b.bot, completion.RegenerateCompletion)) + b.bot.Handle(tele.OnEdited, completion.Handler(b.bot, completion.EditCompletion)) + b.bot.Handle(tele.OnText, completion.Handler(b.bot, completion.NewCompletion)) b.bot.Start() }