diff --git a/backend/app/db/mongo/mongo.go b/backend/app/db/mongo/mongo.go index 2baab24..d5c0b09 100644 --- a/backend/app/db/mongo/mongo.go +++ b/backend/app/db/mongo/mongo.go @@ -22,6 +22,7 @@ type Client struct { type MongoClient interface { Disconnect(ctx context.Context) error GetUser(ctx context.Context) (*models.MongoUser, error) + GetUserIds(ctx context.Context, page int, pageSize int) ([]string, error) GetUsersCount(ctx context.Context) (int64, error) GetUsersCountForSubscription(ctx context.Context, subscription string) (int64, error) MigrateUsersToSubscription(ctx context.Context, from, to string) error @@ -189,3 +190,26 @@ func (c *Client) UpdateUserStripeCustomerId(ctx context.Context, stripeCustomerI _, err := collection.UpdateOne(ctx, filter, update, options) return err } + +func (c *Client) GetUserIds(ctx context.Context, page int, pageSize int) ([]string, error) { + collection := c.Database(config.CONFIG.MongoDBName).Collection("users") + findOptions := options.Find() + findOptions.SetSkip(int64(page * pageSize)) + findOptions.SetLimit(int64(pageSize)) + cursor, err := collection.Find(ctx, bson.M{}, findOptions) + if err != nil { + return nil, fmt.Errorf("GetUserIds: failed to find users: %w", err) + } + defer cursor.Close(ctx) + + var userIds []string + for cursor.Next(ctx) { + var user models.MongoUser + err := cursor.Decode(&user) + if err != nil { + return nil, fmt.Errorf("GetUserIds: failed to decode user: %w", err) + } + userIds = append(userIds, user.ID) + } + return userIds, nil +} diff --git a/backend/app/lib/lib.go b/backend/app/lib/lib.go index 1824af0..040f1be 100644 --- a/backend/app/lib/lib.go +++ b/backend/app/lib/lib.go @@ -121,6 +121,7 @@ func AddBotSuffixToGroupCommands(ctx context.Context, message string) string { client := ctx.Value(models.ClientContext{}).(string) if client == string(TelegramClientName) && strings.HasPrefix(chatString, "-") { message = strings.ReplaceAll(message, "/chatgpt", "/chatgpt@"+config.CONFIG.BotName) + message = strings.ReplaceAll(message, "/voicegpt", "/voicegpt@"+config.CONFIG.BotName) message = strings.ReplaceAll(message, "/clear", "/clear@"+config.CONFIG.BotName) message = strings.ReplaceAll(message, "/downgrade", "/downgrade@"+config.CONFIG.BotName) message = strings.ReplaceAll(message, "/grammar", "/grammar@"+config.CONFIG.BotName) @@ -132,7 +133,8 @@ func AddBotSuffixToGroupCommands(ctx context.Context, message string) string { message = strings.ReplaceAll(message, "/terms", "/terms@"+config.CONFIG.BotName) message = strings.ReplaceAll(message, "/transcribe", "/transcribe@"+config.CONFIG.BotName) message = strings.ReplaceAll(message, "/upgrade", "/upgrade@"+config.CONFIG.BotName) - message = strings.ReplaceAll(message, "/voicegpt", "/voicegpt@"+config.CONFIG.BotName) + message = strings.ReplaceAll(message, "/translate", "/translate@"+config.CONFIG.BotName) + message = strings.ReplaceAll(message, "/billing", "/billing@"+config.CONFIG.BotName) } return message } diff --git a/backend/app/lib/modes.go b/backend/app/lib/modes.go index 9785da3..0b4fbf2 100644 --- a/backend/app/lib/modes.go +++ b/backend/app/lib/modes.go @@ -177,6 +177,13 @@ var summarizeSeed = []models.Message{ }, } +var translateSeed = []models.Message{ + { + Role: "system", + Content: "You will translate a text to {{target}} only. You will keep the meaning and style of the original text. If the original text is not in a {{target}}, you will translate it to {{target}}. If the original text is in {{target}}, you will respond - [correct].", + }, +} + func GetSeedDataAndPrimer(mode ModeName) ([]models.Message, string) { var seedData []models.Message userMessagePrimer := "Text to correct:\n" //default @@ -197,6 +204,9 @@ func GetSeedDataAndPrimer(mode ModeName) ([]models.Message, string) { case Summarize: seedData = summarizeSeed userMessagePrimer = "" + case Translate: + seedData = translateSeed + userMessagePrimer = "Text to translate to {{target}}:\n" default: seedData = chatGPTSeed userMessagePrimer = "" diff --git a/backend/app/telegram/commands.go b/backend/app/telegram/commands.go index 9d49615..8a469e1 100644 --- a/backend/app/telegram/commands.go +++ b/backend/app/telegram/commands.go @@ -40,23 +40,24 @@ Here are some of the things I can do: Enjoy and let me know if any /support is needed!` const ( - CancelSubscriptionCommand Command = "/downgrade" - EmiliCommand Command = "/emily" - EmptyCommand Command = "" + StartCommand Command = "/start" ChatGPTCommand Command = "/chatgpt" VoiceGPTCommand Command = "/voicegpt" GrammarCommand Command = "/grammar" - StartCommand Command = "/start" + ClearThreadCommand Command = "/clear" + TeacherCommand Command = "/teacher" + TranscribeCommand Command = "/transcribe" + SummarizeCommand Command = "/summarize" + TranslateCommand Command = "/translate" StatusCommand Command = "/status" SupportCommand Command = "/support" TermsCommand Command = "/terms" - TeacherCommand Command = "/teacher" UpgradeCommand Command = "/upgrade" - VasilisaCommand Command = "/vasilisa" - TranscribeCommand Command = "/transcribe" - SummarizeCommand Command = "/summarize" - ClearThreadCommand Command = "/clear" + CancelSubscriptionCommand Command = "/downgrade" BillingCommand Command = "/billing" + VasilisaCommand Command = "/vasilisa" + EmiliCommand Command = "/emily" + EmptyCommand Command = "" // commands setting for BotFather Commands string = ` @@ -634,8 +635,8 @@ func IsCreateImageCommand(prompt string) bool { log.Debugf("Cleaned prompt: %s", cleanPrompt) - triggerWords := []string{"create", "draw", "picture", "imagine", "image"} - stopWords := []string{"text", "write", "article"} + triggerWords := []string{"create", "draw", "drawing", "picture", "imagine", "image"} + stopWords := []string{"write", "article"} // Split the prompt into words words := strings.Fields(cleanPrompt) diff --git a/backend/app/telegram/commands_test.go b/backend/app/telegram/commands_test.go index 7d226fb..c156704 100644 --- a/backend/app/telegram/commands_test.go +++ b/backend/app/telegram/commands_test.go @@ -15,6 +15,7 @@ import ( func TestIsCreateImageCommandTrue(t *testing.T) { prompts := []string{ + "Can you create a drawin of a sunset?", "Can you create an image of a sunset?", "Draw, please, an image of a cat", "I'd like a picture! Of a mountain", @@ -31,7 +32,6 @@ func TestIsCreateImageCommandTrue(t *testing.T) { func TestIsCreateImageCommandFalse(t *testing.T) { prompts := []string{ - "Can you create a text?", "Imagene, please, an article about a cat", "I'd like a video! Of a mountain", } diff --git a/backend/app/telegram/telegram_system.go b/backend/app/telegram/telegram_system.go index b85e373..38f4c41 100644 --- a/backend/app/telegram/telegram_system.go +++ b/backend/app/telegram/telegram_system.go @@ -35,6 +35,7 @@ const ( SYSTEMUserCommand Command = "/user" SYSTEMUsersCountCommand Command = "/userscount" SYSTEMUsersForSubscriptionCommand Command = "/usersforsubscription" + SYSTEMSendMessageToUsers Command = "/sendmessagetousers" ) var SystemCommandHandlers CommandHandlers = CommandHandlers{} @@ -139,130 +140,221 @@ func setupSystemCommandHandlers() { newCommandHandler(EmptyCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { bot.SendMessage(tu.Message(SystemBOT.ChatID, "There is no message provided to correct or comment on. If you have a message you would like me to review, please provide it.")) }), - newCommandHandler(SYSTEMStatusCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - systemStatus := redis.RedisClient.Get(context.Background(), "system-status") - bot.SendMessage(tu.Message(SystemBOT.ChatID, systemStatus.Val())) - }), - newCommandHandler(SYSTEMUserCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) - return - } - userId := commandArray[1] - user, err := mongo.MongoDBClient.GetUser(context.WithValue(context.Background(), models.UserContext{}, userId)) - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get user: %s", err))) - return - } - userJson, err := json.Marshal(user) - if err != nil { - log.Errorf("Failed to marshal user: %s", err) - } - userString := "DB:\n" + string(userJson) + "\n\n" + newCommandHandler(SYSTEMStatusCommand, handleStatus), + newCommandHandler(SYSTEMUserCommand, handleUser), + newCommandHandler(SYSTEMStripeResetCommand, handleStripeReset), + newCommandHandler(SYSTEMUsersCountCommand, handleUsersCount), + newCommandHandler(SYSTEMUsageResetCommand, handleUsageReset), + newCommandHandler(SYSTEMUsersForSubscriptionCommand, handleUsersForSubscription), + newCommandHandler(SYSTEMBanUserCommand, handleBanUser), + newCommandHandler(SYSTEMUnbanUserCommand, handleUnbanUser), + newCommandHandler(SYSTEMSendMessageToUsers, handleSendMessageToUsers), + } +} - userString += "Redis:\ntotal cost - " + redis.RedisClient.Get(ctx, lib.UserTotalCostKey(userId)).Val() + "\n" - userString += "total audio minutes - " + redis.RedisClient.Get(ctx, lib.UserTotalAudioMinutesKey(userId)).Val() + "\n" - userString += "total tokens - " + redis.RedisClient.Get(ctx, lib.UserTotalTokensKey(userId)).Val() + "\n\n" - userString += "current thread - " + redis.RedisClient.Get(ctx, lib.UserCurrentThreadKey(userId, "")).Val() + "\n" - userString += "current prompt tokens - " + redis.RedisClient.Get(ctx, lib.UserCurrentThreadPromptKey(userId, "")).Val() +func handleStatus(ctx context.Context, bot *Bot, message *telego.Message) { + systemStatus := redis.RedisClient.Get(context.Background(), "system-status") + bot.SendMessage(tu.Message(SystemBOT.ChatID, systemStatus.Val())) +} - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("User: %+v", userString))) - }), - newCommandHandler(SYSTEMStripeResetCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) - return - } - userId := commandArray[1] - ctx = context.WithValue(ctx, models.UserContext{}, userId) - user, err := mongo.MongoDBClient.GetUser(ctx) - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get user: %s", err))) - return - } - stripeCustomerId := user.StripeCustomerId - if stripeCustomerId != "" { - payments.StripeCancelSubscription(ctx, stripeCustomerId) - customer, err := customer.Del(stripeCustomerId, nil) - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to delete customer: %s", err))) - } else { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Customer deleted from Stripe: %+v", customer))) - } - err = mongo.MongoDBClient.UpdateUserStripeCustomerId(ctx, "") - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to update user: %s", err))) - return - } - // downgrade engine to GPT-3.5 Turbo - redis.SaveEngine(userId, models.ChatGpt35Turbo) - } +func handleUser(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) + return + } + userId := commandArray[1] + user, err := mongo.MongoDBClient.GetUser(context.WithValue(context.Background(), models.UserContext{}, userId)) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get user: %s", err))) + return + } + userJson, err := json.Marshal(user) + if err != nil { + log.Errorf("Failed to marshal user: %s", err) + } + userString := "DB:\n" + string(userJson) + "\n\n" - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Stripe Reset for: %+v", user))) - }), - newCommandHandler(SYSTEMUsersCountCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - users, err := mongo.MongoDBClient.GetUsersCount(context.Background()) + userString += "Redis:\ntotal cost - " + redis.RedisClient.Get(ctx, lib.UserTotalCostKey(userId)).Val() + "\n" + userString += "total audio minutes - " + redis.RedisClient.Get(ctx, lib.UserTotalAudioMinutesKey(userId)).Val() + "\n" + userString += "total tokens - " + redis.RedisClient.Get(ctx, lib.UserTotalTokensKey(userId)).Val() + "\n\n" + userString += "current thread - " + redis.RedisClient.Get(ctx, lib.UserCurrentThreadKey(userId, "")).Val() + "\n" + userString += "current prompt tokens - " + redis.RedisClient.Get(ctx, lib.UserCurrentThreadPromptKey(userId, "")).Val() + + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("User: %+v", userString))) +} + +func handleStripeReset(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) + return + } + userId := commandArray[1] + ctx = context.WithValue(ctx, models.UserContext{}, userId) + user, err := mongo.MongoDBClient.GetUser(ctx) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get user: %s", err))) + return + } + stripeCustomerId := user.StripeCustomerId + if stripeCustomerId != "" { + payments.StripeCancelSubscription(ctx, stripeCustomerId) + customer, err := customer.Del(stripeCustomerId, nil) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to delete customer: %s", err))) + } else { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Customer deleted from Stripe: %+v", customer))) + } + err = mongo.MongoDBClient.UpdateUserStripeCustomerId(ctx, "") + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to update user: %s", err))) + return + } + // downgrade engine to GPT-3.5 Turbo + redis.SaveEngine(userId, models.ChatGpt35Turbo) + } + + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Stripe Reset for: %+v", user))) +} + +func handleUsersCount(ctx context.Context, bot *Bot, message *telego.Message) { + users, err := mongo.MongoDBClient.GetUsersCount(context.Background()) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get users: %s", err))) + return + } + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Users: %+v", users))) +} + +func handleUsageReset(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) + return + } + userId := commandArray[1] + redis.RedisClient.Del(ctx, lib.UserTotalCostKey(userId)) + redis.RedisClient.Del(ctx, lib.UserTotalAudioMinutesKey(userId)) + redis.RedisClient.Del(ctx, lib.UserTotalTokensKey(userId)) + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Usage reset for user: "+userId)) +} + +func handleUsersForSubscription(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide subscription name")) + return + } + subscriptionName := commandArray[1] + usersCount, err := mongo.MongoDBClient.GetUsersCountForSubscription(context.Background(), subscriptionName) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get users: %s", err))) + return + } + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Users count for %s subscription: %d", subscriptionName, usersCount))) +} + +func handleBanUser(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) + return + } + userId := commandArray[1] + err := redis.RedisClient.Set(ctx, userId+":banned", "true", 0).Err() + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to ban user: %v", err))) + return + } + bot.SendMessage(tu.Message(SystemBOT.ChatID, "User "+userId+" banned")) +} + +func handleUnbanUser(ctx context.Context, bot *Bot, message *telego.Message) { + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 2 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) + return + } + userId := commandArray[1] + err := redis.RedisClient.Del(ctx, userId+":banned").Err() + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to unban user: %v", err))) + return + } + bot.SendMessage(tu.Message(SystemBOT.ChatID, "User "+userId+" unbanned")) +} + +func handleSendMessageToUsers(ctx context.Context, bot *Bot, message *telego.Message) { + commandUsage := fmt.Sprintf("Usage: %s ", SYSTEMSendMessageToUsers) + commandArray := strings.Split(message.Text, " ") + if len(commandArray) < 3 { + bot.SendMessage(tu.Message(SystemBOT.ChatID, commandUsage)) + return + } + maximumUsers, err := strconv.Atoi(commandArray[1]) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, commandUsage)) + return + } + messageText := strings.Join(commandArray[2:], " ") + + page := 0 + pageSize := 10 + sleepBetweenPages := 1 * time.Second + notifiedUsers := 0.0 + skippedUsers := 0.0 + + defer func() { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Message sent to %f users, %f users skipped", notifiedUsers, skippedUsers))) + config.CONFIG.DataDogClient.Incr("custom_message_sent_total", []string{}, notifiedUsers) + config.CONFIG.DataDogClient.Incr("custom_message_skipped_total", []string{}, skippedUsers) + }() + + for { + users, err := mongo.MongoDBClient.GetUserIds(context.Background(), page, pageSize) + if err != nil { + bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get users page %d (page size %d): %s", page, pageSize, err))) + return + } + if len(users) == 0 { + break + } + for _, user := range users { + userId, err := strconv.ParseInt(user, 10, 64) if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get users: %s", err))) - return - } - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Users: %+v", users))) - }), - newCommandHandler(SYSTEMUsageResetCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) - return - } - userId := commandArray[1] - redis.RedisClient.Del(ctx, lib.UserTotalCostKey(userId)) - redis.RedisClient.Del(ctx, lib.UserTotalAudioMinutesKey(userId)) - redis.RedisClient.Del(ctx, lib.UserTotalTokensKey(userId)) - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Usage reset for user: "+userId)) - }), - newCommandHandler(SYSTEMUsersForSubscriptionCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide subscription name")) - return + log.Errorf("Failed to convert user id %s to int: %s", user, err) + skippedUsers++ + continue } - subscriptionName := commandArray[1] - usersCount, err := mongo.MongoDBClient.GetUsersCountForSubscription(context.Background(), subscriptionName) - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to get users: %s", err))) - return + + if userId < 0 { + // skip groups + skippedUsers++ + continue } - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Users count for %s subscription: %d", subscriptionName, usersCount))) - }), - newCommandHandler(SYSTEMBanUserCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) - return + + if userId != SystemBOT.ChatID.ID { + // skip anything else but system user for now + skippedUsers++ + continue } - userId := commandArray[1] - err := redis.RedisClient.Set(ctx, userId+":banned", "true", 0).Err() + + _, err = BOT.SendMessage(tu.Message(tu.ID(userId), messageText)) if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to ban user: %v", err))) - return + skippedUsers++ + log.Errorf("Failed to send message to user %d: %s", userId, err) + } else { + notifiedUsers++ + log.Infof("Custom message sent to user %d", userId) } - bot.SendMessage(tu.Message(SystemBOT.ChatID, "User "+userId+" banned")) - }), - newCommandHandler(SYSTEMUnbanUserCommand, func(ctx context.Context, bot *Bot, message *telego.Message) { - commandArray := strings.Split(message.Text, " ") - if len(commandArray) < 2 { - bot.SendMessage(tu.Message(SystemBOT.ChatID, "Please provide user id")) - return - } - userId := commandArray[1] - err := redis.RedisClient.Del(ctx, userId+":banned").Err() - if err != nil { - bot.SendMessage(tu.Message(SystemBOT.ChatID, fmt.Sprintf("Failed to unban user: %v", err))) + + if notifiedUsers >= float64(maximumUsers) { + log.Infof("Maximum users reached: %f", notifiedUsers) return } - bot.SendMessage(tu.Message(SystemBOT.ChatID, "User "+userId+" unbanned")) - }), + } + time.Sleep(sleepBetweenPages) + page++ } }