Add edit support and enhance models command

Signed-off-by: Sid Sun <sid@sidsun.com>

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,

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
+)

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()))

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{
+		modelOptions = append(modelOptions, []tele.InlineButton{
+			{
 				Data: "model_" + modelName,
 				Text: text,
+			},
 		})
 	}
 	return modelInfoMessage, modelOptions

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()
 }