refactor: better reply handling

davidramiro · Dec 23, 2024 · 6781f62 · 6781f62
1 parent 68bb063
commit 6781f62
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 7 deletions.
diff --git a/internal/adapters/generator/claude.go b/internal/adapters/generator/claude.go
@@ -24,17 +24,17 @@ func NewClaudeGenerator(apiKey, systemPrompt string) *ClaudeGenerator {
 }
 
 func (c *ClaudeGenerator) GenerateFromPrompt(ctx context.Context, prompts []domain.Prompt) (string, error) {
-	var messages []anthropic.Message
+	messages := make([]anthropic.Message, len(prompts))
 
-	for _, prompt := range prompts {
+	for i, prompt := range prompts {
 		if prompt.Author == domain.System {
-			messages = append(messages, anthropic.NewAssistantTextMessage(prompt.Prompt))
+			messages[i] = anthropic.NewAssistantTextMessage(prompt.Prompt)
 		} else if prompt.Author == domain.User {
 			message, err := createUserMessage(ctx, prompt)
 			if err != nil {
 				return "", err
 			}
-			messages = append(messages, message)
+			messages[i] = message
 		}
 	}
 

diff --git a/internal/adapters/handler/command.go b/internal/adapters/handler/command.go
@@ -39,7 +39,16 @@ func (h *CommandHandler) Handle(ctx context.Context, b *bot.Bot, update *models.
 
 	replyToMessageID := new(int)
 	var quotedText string
+	var isReplyToBot bool
 	if update.Message.ReplyToMessage != nil {
+		botUser, err := b.GetMe(ctx)
+		if err != nil {
+			log.Debug().Str("command", cmd).Msg("failed to get bot user")
+			return
+		}
+		if update.Message.ReplyToMessage.From.ID == botUser.ID {
+			isReplyToBot = true
+		}
 		*replyToMessageID = update.Message.ReplyToMessage.ID
 		quotedText = update.Message.ReplyToMessage.Text
 	}
@@ -56,6 +65,7 @@ func (h *CommandHandler) Handle(ctx context.Context, b *bot.Bot, update *models.
 			ChatID:           update.Message.Chat.ID,
 			Text:             update.Message.Text,
 			ReplyToMessageID: replyToMessageID,
+			IsReplyToBot:     isReplyToBot,
 			QuotedText:       quotedText,
 			ImageURL:         <-imageURL,
 			AudioURL:         <-audioURL,
@@ -66,7 +76,7 @@ func (h *CommandHandler) Handle(ctx context.Context, b *bot.Bot, update *models.
 	}()
 }
 
-func getOptionalImage(ctx context.Context, b *bot.Bot, update *models.Update, url chan string) {
+func getOptionalImage(ctx context.Context, b *bot.Bot, update *models.Update, url chan<- string) {
 	var photos []models.PhotoSize
 
 	if update.Message.Photo != nil {
@@ -94,7 +104,7 @@ func getOptionalImage(ctx context.Context, b *bot.Bot, update *models.Update, ur
 	url <- b.FileDownloadLink(f)
 }
 
-func getOptionalAudio(ctx context.Context, b *bot.Bot, update *models.Update, url chan string) {
+func getOptionalAudio(ctx context.Context, b *bot.Bot, update *models.Update, url chan<- string) {
 	var fileID string
 	if update.Message.Audio != nil {
 		fileID = update.Message.Audio.FileID

diff --git a/internal/core/domain/commands/chat.go b/internal/core/domain/commands/chat.go
@@ -76,8 +76,14 @@ func (h *ChatHandler) Respond(ctx context.Context, timeout time.Duration, messag
 	}
 
 	if message.QuotedText != "" && message.ImageURL == "" {
+		// if there's a user message being replied to, add the previous message to the context
+		if !message.IsReplyToBot {
+			conversation.messages = append(conversation.messages, domain.Prompt{Author: domain.User,
+				Prompt: message.QuotedText})
+		}
+
 		conversation.messages = append(conversation.messages, domain.Prompt{Author: domain.User,
-			Prompt: fmt.Sprintf("%s:%s", promptText, message.QuotedText)})
+			Prompt: message.Text})
 	} else {
 		conversation.messages = append(conversation.messages, domain.Prompt{Author: domain.User,
 			Prompt: promptText, ImageURL: message.ImageURL})

diff --git a/internal/core/domain/model.go b/internal/core/domain/model.go
@@ -17,6 +17,7 @@ type Message struct {
 	ID               int
 	ChatID           int64
 	ReplyToMessageID *int
+	IsReplyToBot     bool
 	QuotedText       string
 	ImageURL         string
 	AudioURL         string