Implemented screening

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
+	// Send the response through the centralized screen
-	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
 	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
+	// Send the response through the centralized screen
-	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
 	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)
+}
+
+// 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 {
-		log.Printf("Error storing assistant message: %v", err)
+		return Message{}, err
 	}
-	b.addMessageToChatMemory(b.getOrCreateChatMemory(chatID), assistantMessage)
+
+	// Update chat memory.
+	chatMemory := b.getOrCreateChatMemory(chatID)
+	b.addMessageToChatMemory(chatMemory, assistantMessage)
+
+	return assistantMessage, nil
 }

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)
+	// Determine if the text contains only emojis
-	userMessage.UserRole = string(anthropic.RoleUser) // Convert to string
+	isEmojiOnly := isOnlyEmojis(text)
-	if err := b.storeMessage(userMessage); err != nil {
-		log.Printf("Error storing user message: %v", err)
-		return
-	}
 
+	// 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 {
+	// Update chat memory with the user message
-		log.Printf("Error storing user message: %v", err)
-	}
-
-	// Update chat memory
 	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) {