diff --git a/bot.go b/bot.go index b46e007..443ae13 100644 --- a/bot.go +++ b/bot.go @@ -32,7 +32,7 @@ type Bot struct { botID uint // Reference to BotModel.ID } -// bot.go +// NewBot initializes and returns a new Bot instance. func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) (*Bot, error) { // Retrieve or create Bot entry in the database var botEntry BotModel @@ -94,6 +94,7 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) return b, nil } +// Start begins the bot's operation. func (b *Bot) Start(ctx context.Context) { b.tgBot.Start(ctx) } @@ -281,6 +282,14 @@ func initTelegramBot(token string, handleUpdate func(ctx context.Context, tgBot } func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, businessConnectionID string) error { + // Pass the outgoing message through the centralized screen for storage + _, err := b.screenOutgoingMessage(chatID, text, businessConnectionID) + if err != nil { + log.Printf("Error storing assistant message: %v", err) + return err + } + + // Prepare message parameters params := &bot.SendMessageParams{ ChatID: chatID, Text: text, @@ -290,7 +299,8 @@ func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, busin params.BusinessConnectionID = businessConnectionID } - _, err := b.tgBot.SendMessage(ctx, params) + // Send the message via Telegram client + _, err = b.tgBot.SendMessage(ctx, params) if err != nil { log.Printf("[%s] [ERROR] Error sending message to chat %d with BusinessConnectionID %s: %v", b.config.ID, chatID, businessConnectionID, err) @@ -319,20 +329,10 @@ func (b *Bot) sendStats(ctx context.Context, chatID int64, userID int64, usernam totalMessages, ) - // Store the user's /stats command - userMessage := b.createMessage(chatID, userID, username, "user", "/stats", true) - if err := b.storeMessage(userMessage); err != nil { - log.Printf("Error storing user message: %v", err) - } - - // Send and store the bot's response + // Send the response through the centralized screen if err := b.sendResponse(ctx, chatID, statsMessage, businessConnectionID); err != nil { log.Printf("Error sending stats message: %v", err) } - assistantMessage := b.createMessage(chatID, 0, "", "assistant", statsMessage, false) - if err := b.storeMessage(assistantMessage); err != nil { - log.Printf("Error storing assistant message: %v", err) - } } // getStats retrieves the total number of users and messages from the database. @@ -389,19 +389,49 @@ func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, userna caser.String(user.Role.Name), ) - // Store the user's /whoami command - userMessage := b.createMessage(chatID, userID, username, "user", "/whoami", true) - if err := b.storeMessage(userMessage); err != nil { - log.Printf("Error storing user message: %v", err) - } - - // Send and store the bot's response + // Send the response through the centralized screen if err := b.sendResponse(ctx, chatID, whoAmIMessage, businessConnectionID); err != nil { log.Printf("Error sending /whoami message: %v", err) } - assistantMessage := b.createMessage(chatID, 0, "", "assistant", whoAmIMessage, false) - if err := b.storeMessage(assistantMessage); err != nil { - log.Printf("Error storing assistant message: %v", err) - } - b.addMessageToChatMemory(b.getOrCreateChatMemory(chatID), assistantMessage) +} + +// screenIncomingMessage handles storing of incoming messages. +func (b *Bot) screenIncomingMessage(message *models.Message) (Message, error) { + userRole := string(anthropic.RoleUser) // Convert RoleUser to string + userMessage := b.createMessage(message.Chat.ID, message.From.ID, message.From.Username, userRole, message.Text, true) + + // If the message contains a sticker, include its details. + if message.Sticker != nil { + userMessage.StickerFileID = message.Sticker.FileID + if message.Sticker.Thumbnail != nil { + userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID + } + } + + // Store the message. + if err := b.storeMessage(userMessage); err != nil { + return Message{}, err + } + + // Update chat memory. + chatMemory := b.getOrCreateChatMemory(message.Chat.ID) + b.addMessageToChatMemory(chatMemory, userMessage) + + return userMessage, nil +} + +// screenOutgoingMessage handles storing of outgoing messages. +func (b *Bot) screenOutgoingMessage(chatID int64, response string, businessConnectionID string) (Message, error) { + assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false) + + // Store the message. + if err := b.storeMessage(assistantMessage); err != nil { + return Message{}, err + } + + // Update chat memory. + chatMemory := b.getOrCreateChatMemory(chatID) + b.addMessageToChatMemory(chatMemory, assistantMessage) + + return assistantMessage, nil } diff --git a/handlers.go b/handlers.go index 6121f74..96720fe 100644 --- a/handlers.go +++ b/handlers.go @@ -22,9 +22,6 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U return } - chatID := message.Chat.ID - userID := message.From.ID - // Extract businessConnectionID if available var businessConnectionID string if update.BusinessConnection != nil { @@ -33,6 +30,18 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U businessConnectionID = message.BusinessConnectionID } + chatID := message.Chat.ID + userID := message.From.ID + username := message.From.Username + text := message.Text + + // Pass the incoming message through the centralized screen for storage + _, err := b.screenIncomingMessage(message) + if err != nil { + log.Printf("Error storing user message: %v", err) + return + } + // Check if the message is a command if message.Entities != nil { for _, entity := range message.Entities { @@ -40,10 +49,10 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U command := strings.TrimSpace(message.Text[entity.Offset : entity.Offset+entity.Length]) switch command { case "/stats": - b.sendStats(ctx, chatID, userID, message.From.Username, businessConnectionID) + b.sendStats(ctx, chatID, userID, username, businessConnectionID) return case "/whoami": - b.sendWhoAmI(ctx, chatID, userID, message.From.Username, businessConnectionID) + b.sendWhoAmI(ctx, chatID, userID, username, businessConnectionID) return } } @@ -56,25 +65,21 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U return } - // Existing rate limit and message handling + // Rate limit check if !b.checkRateLimits(userID) { b.sendRateLimitExceededMessage(ctx, chatID, businessConnectionID) return } - username := message.From.Username - text := message.Text - // Proceed only if the message contains text if text == "" { - // Optionally, handle other message types or ignore log.Printf("Received a non-text message from user %d in chat %d", userID, chatID) return } // Determine if the user is the owner var isOwner bool - err := b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error + err = b.db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", userID, b.botID, true).First(&User{}).Error if err == nil { isOwner = true } @@ -93,35 +98,26 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U } } - userMessage := b.createMessage(chatID, userID, username, user.Role.Name, text, true) - userMessage.UserRole = string(anthropic.RoleUser) // Convert to string - if err := b.storeMessage(userMessage); err != nil { - log.Printf("Error storing user message: %v", err) - return - } + // Determine if the text contains only emojis + isEmojiOnly := isOnlyEmojis(text) + // Prepare context messages for Anthropic chatMemory := b.getOrCreateChatMemory(chatID) - b.addMessageToChatMemory(chatMemory, userMessage) - + b.addMessageToChatMemory(chatMemory, b.createMessage(chatID, userID, username, user.Role.Name, text, true)) contextMessages := b.prepareContextMessages(chatMemory) - isEmojiOnly := isOnlyEmojis(text) + // Get response from Anthropic response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), isOwner, isEmojiOnly) if err != nil { log.Printf("Error getting Anthropic response: %v", err) response = "I'm sorry, I'm having trouble processing your request right now." } + // Send the response through the centralized screen if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil { log.Printf("Error sending response: %v", err) return } - - assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false) - if err := b.storeMessage(assistantMessage); err != nil { - log.Printf("Error storing assistant message: %v", err) - } - b.addMessageToChatMemory(chatMemory, assistantMessage) } func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, businessConnectionID string) { @@ -131,7 +127,7 @@ func (b *Bot) sendRateLimitExceededMessage(ctx context.Context, chatID int64, bu func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, message *models.Message, businessConnectionID string) { username := message.From.Username - // Create and store the sticker message + // Create the user message (without storing it manually) userMessage := b.createMessage(chatID, userID, username, "user", "Sent a sticker.", true) userMessage.StickerFileID = message.Sticker.FileID @@ -140,11 +136,7 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me userMessage.StickerPNGFile = message.Sticker.Thumbnail.FileID } - if err := b.storeMessage(userMessage); err != nil { - log.Printf("Error storing user message: %v", err) - } - - // Update chat memory + // Update chat memory with the user message chatMemory := b.getOrCreateChatMemory(chatID) b.addMessageToChatMemory(chatMemory, userMessage) @@ -162,11 +154,8 @@ func (b *Bot) handleStickerMessage(ctx context.Context, chatID, userID int64, me } } + // Send the response through the centralized screen b.sendResponse(ctx, chatID, response, businessConnectionID) - - assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false) - b.storeMessage(assistantMessage) - b.addMessageToChatMemory(chatMemory, assistantMessage) } func (b *Bot) generateStickerResponse(ctx context.Context, message Message) (string, error) {