From aca2008404df33f3dd1020d1a7d093a2dee5443c Mon Sep 17 00:00:00 2001 From: HugeFrog24 <62775760+HugeFrog24@users.noreply.github.com> Date: Tue, 22 Oct 2024 14:32:16 +0200 Subject: [PATCH 1/4] Started implementing owner feature --- .gitignore | 0 anthropic.go | 0 bot.go | 83 ++++++++++++++++++++++++++++++++++++++++---- clock.go | 0 config.go | 1 + config/default.json | 1 + database.go | 4 +++ go.mod | 0 go.sum | 0 handlers.go | 13 +++++-- main.go | 0 models.go | 7 ++-- rate_limiter.go | 0 rate_limiter_test.go | 1 + 14 files changed, 98 insertions(+), 12 deletions(-) mode change 100755 => 100644 .gitignore mode change 100755 => 100644 anthropic.go mode change 100755 => 100644 bot.go mode change 100755 => 100644 clock.go mode change 100755 => 100644 config.go mode change 100755 => 100644 config/default.json mode change 100755 => 100644 database.go mode change 100755 => 100644 go.mod mode change 100755 => 100644 go.sum mode change 100755 => 100644 handlers.go mode change 100755 => 100644 main.go mode change 100755 => 100644 models.go mode change 100755 => 100644 rate_limiter.go mode change 100755 => 100644 rate_limiter_test.go diff --git a/.gitignore b/.gitignore old mode 100755 new mode 100644 diff --git a/anthropic.go b/anthropic.go old mode 100755 new mode 100644 diff --git a/bot.go b/bot.go old mode 100755 new mode 100644 index b90593c..2cc861f --- a/bot.go +++ b/bot.go @@ -43,6 +43,33 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { return nil, err } + // Ensure the owner exists in the Users table + var owner User + err = db.Where("telegram_id = ? AND bot_id = ?", config.OwnerTelegramID, botEntry.ID).First(&owner).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + // Assign the "owner" role + var ownerRole Role + err := db.Where("name = ?", "owner").First(&ownerRole).Error + if err != nil { + return nil, fmt.Errorf("owner role not found: %w", err) + } + + owner = User{ + BotID: botEntry.ID, + TelegramID: config.OwnerTelegramID, + Username: "Owner", // You might want to fetch the actual username + RoleID: ownerRole.ID, + IsOwner: true, + } + + if err := db.Create(&owner).Error; err != nil { + return nil, fmt.Errorf("failed to create owner user: %w", err) + } + } else if err != nil { + return nil, err + } + + // Initialize Anthropic client anthropicClient := anthropic.NewClient(os.Getenv("ANTHROPIC_API_KEY")) b := &Bot{ @@ -69,26 +96,70 @@ func (b *Bot) Start(ctx context.Context) { b.tgBot.Start(ctx) } -func (b *Bot) getOrCreateUser(userID int64, username string) (User, error) { +func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User, error) { var user User - err := b.db.Preload("Role").Where("telegram_id = ?", userID).First(&user).Error + err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - var defaultRole Role - if err := b.db.Where("name = ?", "user").First(&defaultRole).Error; err != nil { - return User{}, err + var role Role + if isOwner { + role, err = b.getRoleByName("owner") + if err != nil { + return User{}, err + } + } else { + role, err = b.getRoleByName("admin") + if err != nil { + return User{}, err + } } - user = User{TelegramID: userID, Username: username, RoleID: defaultRole.ID} + + user = User{ + BotID: b.botID, + TelegramID: userID, + Username: username, + RoleID: role.ID, + IsOwner: isOwner, + } + if err := b.db.Create(&user).Error; err != nil { return User{}, err } } else { return User{}, err } + } else { + if isOwner && !user.IsOwner { + // Check if another owner exists + var existingOwner User + err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error + if err == nil { + return User{}, fmt.Errorf("a bot can have only one owner") + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return User{}, err + } + // Promote to owner + user.Role, err = b.getRoleByName("owner") + if err != nil { + return User{}, err + } + user.RoleID = user.Role.ID + user.IsOwner = true + if err := b.db.Save(&user).Error; err != nil { + return User{}, err + } + } } + return user, nil } +func (b *Bot) getRoleByName(roleName string) (Role, error) { + var role Role + err := b.db.Where("name = ?", roleName).First(&role).Error + return role, err +} + func (b *Bot) createMessage(chatID, userID int64, username, userRole, text string, isUser bool) Message { message := Message{ ChatID: chatID, diff --git a/clock.go b/clock.go old mode 100755 new mode 100644 diff --git a/config.go b/config.go old mode 100755 new mode 100644 index 0ff9436..2d1e556 --- a/config.go +++ b/config.go @@ -19,6 +19,7 @@ type BotConfig struct { Model anthropic.Model `json:"model"` // Changed from string to anthropic.Model SystemPrompts map[string]string `json:"system_prompts"` Active bool `json:"active"` // New field to control bot activity + OwnerTelegramID int64 `json:"owner_telegram_id"` } // Custom unmarshalling to handle anthropic.Model diff --git a/config/default.json b/config/default.json old mode 100755 new mode 100644 index 9df35f7..fc0e75a --- a/config/default.json +++ b/config/default.json @@ -2,6 +2,7 @@ "id": "default_bot", "active": false, "telegram_token": "YOUR_TELEGRAM_BOT_TOKEN", + "owner_telegram_id": 000000000, "memory_size": 10, "messages_per_hour": 20, "messages_per_day": 100, diff --git a/database.go b/database.go old mode 100755 new mode 100644 index da7a29a..d35a25d --- a/database.go +++ b/database.go @@ -27,11 +27,15 @@ func initDB() (*gorm.DB, error) { return nil, fmt.Errorf("failed to connect to database: %w", err) } + // AutoMigrate with unique constraint for owners err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) if err != nil { return nil, fmt.Errorf("failed to migrate database schema: %w", err) } + // Add unique index for owners per bot + db.SetupJoinTable(&BotModel{}, "Users", &User{}) + err = createDefaultRoles(db) if err != nil { return nil, err diff --git a/go.mod b/go.mod old mode 100755 new mode 100644 diff --git a/go.sum b/go.sum old mode 100755 new mode 100644 diff --git a/handlers.go b/handlers.go old mode 100755 new mode 100644 index a7d9473..4d01477 --- a/handlers.go +++ b/handlers.go @@ -69,7 +69,14 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U return } - user, err := b.getOrCreateUser(userID, username) + // 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 + if err == nil { + isOwner = true + } + + user, err := b.getOrCreateUser(userID, username, isOwner) if err != nil { log.Printf("Error getting or creating user: %v", err) return @@ -84,8 +91,8 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U contextMessages := b.prepareContextMessages(chatMemory) - isEmojiOnly := isOnlyEmojis(text) // Ensure you have this variable defined - response, err := b.getAnthropicResponse(ctx, contextMessages, b.isNewChat(chatID), b.isAdminOrOwner(userID), isEmojiOnly) + isEmojiOnly := isOnlyEmojis(text) + 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." diff --git a/main.go b/main.go old mode 100755 new mode 100644 diff --git a/models.go b/models.go old mode 100755 new mode 100644 index f518754..3c2247f --- a/models.go +++ b/models.go @@ -11,7 +11,7 @@ type BotModel struct { Identifier string `gorm:"uniqueIndex"` // Renamed from ID to Identifier Name string Configs []ConfigModel `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` - Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Added foreign key + Users []User `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` // Associated users Messages []Message `gorm:"foreignKey:BotID;constraint:OnDelete:CASCADE"` } @@ -54,9 +54,10 @@ type Role struct { type User struct { gorm.Model - BotID uint `gorm:"index"` // Added foreign key to BotModel - TelegramID int64 `gorm:"uniqueIndex"` // Consider composite unique index if TelegramID is unique per Bot + BotID uint `gorm:"index"` // Foreign key to BotModel + TelegramID int64 `gorm:"uniqueIndex;not null"` // Unique per user Username string RoleID uint Role Role `gorm:"foreignKey:RoleID"` + IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner } diff --git a/rate_limiter.go b/rate_limiter.go old mode 100755 new mode 100644 diff --git a/rate_limiter_test.go b/rate_limiter_test.go old mode 100755 new mode 100644 index 6e4f54a..a2871f3 --- a/rate_limiter_test.go +++ b/rate_limiter_test.go @@ -22,6 +22,7 @@ func TestCheckRateLimits(t *testing.T) { TempBanDuration: "1m", // Temporary ban duration of 1 minute for testing SystemPrompts: make(map[string]string), TelegramToken: "YOUR_TELEGRAM_BOT_TOKEN", + OwnerTelegramID: 123456789, } // Initialize the Bot with mock data and MockClock From 972973633a8c97399b044674ab3eb9ea14cc7ad5 Mon Sep 17 00:00:00 2001 From: HugeFrog24 <62775760+HugeFrog24@users.noreply.github.com> Date: Tue, 22 Oct 2024 14:54:39 +0200 Subject: [PATCH 2/4] Updated owner management --- bot.go | 12 ++++- database.go | 13 ++++-- handlers.go | 17 +++++-- models.go | 5 ++ rate_limiter_test.go | 107 +++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 145 insertions(+), 9 deletions(-) diff --git a/bot.go b/bot.go index 2cc861f..9b3b80c 100644 --- a/bot.go +++ b/bot.go @@ -63,6 +63,10 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { } if err := db.Create(&owner).Error; err != nil { + // If unique constraint is violated, another owner already exists + if strings.Contains(err.Error(), "unique index") { + return nil, fmt.Errorf("an owner already exists for this bot") + } return nil, fmt.Errorf("failed to create owner user: %w", err) } } else if err != nil { @@ -123,6 +127,10 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User } if err := b.db.Create(&user).Error; err != nil { + // Handle unique constraint for owner + if isOwner && strings.Contains(err.Error(), "unique index") { + return User{}, fmt.Errorf("an owner already exists for this bot") + } return User{}, err } } else { @@ -139,11 +147,11 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User return User{}, err } // Promote to owner - user.Role, err = b.getRoleByName("owner") + role, err := b.getRoleByName("owner") if err != nil { return User{}, err } - user.RoleID = user.Role.ID + user.RoleID = role.ID user.IsOwner = true if err := b.db.Save(&user).Error; err != nil { return User{}, err diff --git a/database.go b/database.go index d35a25d..6c7a8f6 100644 --- a/database.go +++ b/database.go @@ -27,14 +27,21 @@ func initDB() (*gorm.DB, error) { return nil, fmt.Errorf("failed to connect to database: %w", err) } - // AutoMigrate with unique constraint for owners + // AutoMigrate the models err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) if err != nil { return nil, fmt.Errorf("failed to migrate database schema: %w", err) } - // Add unique index for owners per bot - db.SetupJoinTable(&BotModel{}, "Users", &User{}) + // Enforce unique owner per bot using raw SQL + // Note: SQLite doesn't support partial indexes, but we can simulate it by making a unique index on (BotID, IsOwner) + // and ensuring that IsOwner can only be true for one user per BotID. + // This approach allows multiple users with IsOwner=false for the same BotID, + // but only one user can have IsOwner=true per BotID. + err = db.Exec(`CREATE UNIQUE INDEX IF NOT EXISTS idx_bot_owner ON users (bot_id, is_owner) WHERE is_owner = 1;`).Error + if err != nil { + return nil, fmt.Errorf("failed to create unique index for bot owners: %w", err) + } err = createDefaultRoles(db) if err != nil { diff --git a/handlers.go b/handlers.go index 4d01477..149b772 100644 --- a/handlers.go +++ b/handlers.go @@ -84,7 +84,10 @@ 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 - b.storeMessage(userMessage) + if err := b.storeMessage(userMessage); err != nil { + log.Printf("Error storing user message: %v", err) + return + } chatMemory := b.getOrCreateChatMemory(chatID) b.addMessageToChatMemory(chatMemory, userMessage) @@ -98,10 +101,16 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U response = "I'm sorry, I'm having trouble processing your request right now." } - b.sendResponse(ctx, chatID, response, businessConnectionID) + if err := b.sendResponse(ctx, chatID, response, businessConnectionID); err != nil { + log.Printf("Error sending response: %v", err) + return + } - assistantMessage := b.createMessage(chatID, 0, "", string(anthropic.RoleAssistant), response, false) - b.storeMessage(assistantMessage) + assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false) + if err := b.storeMessage(assistantMessage); err != nil { + log.Printf("Error storing assistant message: %v", err) + return + } b.addMessageToChatMemory(chatMemory, assistantMessage) } diff --git a/models.go b/models.go index 3c2247f..6ce7b1e 100644 --- a/models.go +++ b/models.go @@ -61,3 +61,8 @@ type User struct { Role Role `gorm:"foreignKey:RoleID"` IsOwner bool `gorm:"default:false"` // Indicates if the user is the owner } + +// Compound unique index to ensure only one owner per bot +func (User) TableName() string { + return "users" +} diff --git a/rate_limiter_test.go b/rate_limiter_test.go index a2871f3..5965a3e 100644 --- a/rate_limiter_test.go +++ b/rate_limiter_test.go @@ -1,8 +1,12 @@ package main import ( + "strings" "testing" "time" + + "gorm.io/driver/sqlite" + "gorm.io/gorm" ) // TestCheckRateLimits tests the checkRateLimits method of the Bot. @@ -80,6 +84,109 @@ func TestCheckRateLimits(t *testing.T) { } } +func TestOwnerAssignment(t *testing.T) { + // Initialize in-memory database for testing + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("Failed to open in-memory database: %v", err) + } + + // Migrate the schema + err = db.AutoMigrate(&BotModel{}, &ConfigModel{}, &Message{}, &User{}, &Role{}) + if err != nil { + t.Fatalf("Failed to migrate database schema: %v", err) + } + + // Create default roles + err = createDefaultRoles(db) + if err != nil { + t.Fatalf("Failed to create default roles: %v", err) + } + + // Create a bot configuration + config := BotConfig{ + ID: "test_bot", + TelegramToken: "TEST_TELEGRAM_TOKEN", + MemorySize: 10, + MessagePerHour: 5, + MessagePerDay: 10, + TempBanDuration: "1m", + SystemPrompts: make(map[string]string), + Active: true, + OwnerTelegramID: 111111111, + } + + // Initialize MockClock + mockClock := &MockClock{ + currentTime: time.Now(), + } + + // Create the bot + bot, err := NewBot(db, config, mockClock) + if err != nil { + t.Fatalf("Failed to create bot: %v", err) + } + + // Verify that the owner exists + var owner User + err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, bot.botID, true).First(&owner).Error + if err != nil { + t.Fatalf("Owner was not created: %v", err) + } + + // Attempt to create another owner for the same bot + _, err = bot.getOrCreateUser(222222222, "AnotherOwner", true) + if err == nil { + t.Fatalf("Expected error when creating a second owner, but got none") + } + + // Verify that the error message is appropriate + expectedErrorMsg := "an owner already exists for this bot" + if err.Error() != expectedErrorMsg && !strings.Contains(err.Error(), "unique index") { + t.Fatalf("Unexpected error message: %v", err) + } + + // Assign admin role to a new user + adminUser, err := bot.getOrCreateUser(333333333, "AdminUser", false) + if err != nil { + t.Fatalf("Failed to create admin user: %v", err) + } + + if adminUser.Role.Name != "admin" { + t.Fatalf("Expected role 'admin', got '%s'", adminUser.Role.Name) + } + + // Assign owner role to a user from a different bot + otherBotConfig := BotConfig{ + ID: "other_bot", + TelegramToken: "OTHER_TELEGRAM_TOKEN", + MemorySize: 10, + MessagePerHour: 5, + MessagePerDay: 10, + TempBanDuration: "1m", + SystemPrompts: make(map[string]string), + Active: true, + OwnerTelegramID: 444444444, + } + + otherBot, err := NewBot(db, otherBotConfig, mockClock) + if err != nil { + t.Fatalf("Failed to create other bot: %v", err) + } + + _, err = otherBot.getOrCreateUser(config.OwnerTelegramID, "OwnerOfOtherBot", true) + if err != nil { + t.Fatalf("Failed to assign existing owner to another bot: %v", err) + } + + // Verify multiple bots can have the same owner telegram ID + var ownerOfOtherBot User + err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, otherBot.botID, true).First(&ownerOfOtherBot).Error + if err != nil { + t.Fatalf("Owner of other bot was not created: %v", err) + } +} + // To ensure thread safety and avoid race conditions during testing, // you can run the tests with the `-race` flag: // go test -race -v From acaf5d01ab925b6fc983c42eb878abaec92bec2a Mon Sep 17 00:00:00 2001 From: HugeFrog24 <62775760+HugeFrog24@users.noreply.github.com> Date: Tue, 22 Oct 2024 14:56:17 +0200 Subject: [PATCH 3/4] Updated json and gitignore --- .gitignore | 4 ++-- config/default.json | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.gitignore b/.gitignore index 21238dc..ae342fe 100644 --- a/.gitignore +++ b/.gitignore @@ -4,8 +4,8 @@ vendor/ # Environment variables .env -# Log file -bot.log +# Any log files +*.log # Database file bot.db diff --git a/config/default.json b/config/default.json index fc0e75a..036fd74 100644 --- a/config/default.json +++ b/config/default.json @@ -2,7 +2,7 @@ "id": "default_bot", "active": false, "telegram_token": "YOUR_TELEGRAM_BOT_TOKEN", - "owner_telegram_id": 000000000, + "owner_telegram_id": 111111111, "memory_size": 10, "messages_per_hour": 20, "messages_per_day": 100, From ce59b5f5f1117c5a4f4567907d829b56f3a5d417 Mon Sep 17 00:00:00 2001 From: HugeFrog24 <62775760+HugeFrog24@users.noreply.github.com> Date: Tue, 22 Oct 2024 15:50:51 +0200 Subject: [PATCH 4/4] Proceed with role management --- bot.go | 112 ++++++++++++++++++++++++++-------------- handlers.go | 12 ++++- main.go | 14 ++++- rate_limiter_test.go | 56 ++++++++++---------- telegram_client.go | 16 ++++++ telegram_client_mock.go | 35 +++++++++++++ 6 files changed, 172 insertions(+), 73 deletions(-) create mode 100644 telegram_client.go create mode 100644 telegram_client_mock.go diff --git a/bot.go b/bot.go index 9b3b80c..7b851cd 100644 --- a/bot.go +++ b/bot.go @@ -13,11 +13,13 @@ import ( "github.com/go-telegram/bot" "github.com/go-telegram/bot/models" "github.com/liushuangls/go-anthropic/v2" + "golang.org/x/text/cases" + "golang.org/x/text/language" "gorm.io/gorm" ) type Bot struct { - tgBot *bot.Bot + tgBot TelegramClient db *gorm.DB anthropicClient *anthropic.Client chatMemories map[int64]*ChatMemory @@ -30,7 +32,8 @@ type Bot struct { botID uint // Reference to BotModel.ID } -func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { +// bot.go +func NewBot(db *gorm.DB, config BotConfig, clock Clock, tgClient TelegramClient) (*Bot, error) { // Retrieve or create Bot entry in the database var botEntry BotModel err := db.Where("identifier = ?", config.ID).First(&botEntry).Error @@ -57,7 +60,7 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { owner = User{ BotID: botEntry.ID, TelegramID: config.OwnerTelegramID, - Username: "Owner", // You might want to fetch the actual username + Username: "", // Initialize as empty; will be updated upon interaction RoleID: ownerRole.ID, IsOwner: true, } @@ -85,14 +88,9 @@ func NewBot(db *gorm.DB, config BotConfig, clock Clock) (*Bot, error) { userLimiters: make(map[int64]*userLimiter), clock: clock, botID: botEntry.ID, // Ensure BotModel has ID field + tgBot: tgClient, } - tgBot, err := initTelegramBot(config.TelegramToken, b.handleUpdate) - if err != nil { - return nil, err - } - b.tgBot = tgBot - return b, nil } @@ -105,17 +103,28 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User err := b.db.Preload("Role").Where("telegram_id = ? AND bot_id = ?", userID, b.botID).First(&user).Error if err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { - var role Role + // Check if an owner already exists for this bot if isOwner { - role, err = b.getRoleByName("owner") - if err != nil { - return User{}, err + var existingOwner User + err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error + if err == nil { + return User{}, fmt.Errorf("an owner already exists for this bot") + } else if !errors.Is(err, gorm.ErrRecordNotFound) { + return User{}, fmt.Errorf("failed to check existing owner: %w", err) } + } + + var role Role + var roleName string + if isOwner { + roleName = "owner" } else { - role, err = b.getRoleByName("admin") - if err != nil { - return User{}, err - } + roleName = "user" // Assign "user" role to non-owner users + } + + err := b.db.Where("name = ?", roleName).First(&role).Error + if err != nil { + return User{}, fmt.Errorf("failed to get role: %w", err) } user = User{ @@ -123,39 +132,23 @@ func (b *Bot) getOrCreateUser(userID int64, username string, isOwner bool) (User TelegramID: userID, Username: username, RoleID: role.ID, + Role: role, IsOwner: isOwner, } if err := b.db.Create(&user).Error; err != nil { - // Handle unique constraint for owner - if isOwner && strings.Contains(err.Error(), "unique index") { + // If unique constraint is violated, another owner already exists + if strings.Contains(err.Error(), "unique index") { return User{}, fmt.Errorf("an owner already exists for this bot") } - return User{}, err + return User{}, fmt.Errorf("failed to create user: %w", err) } } else { return User{}, err } } else { if isOwner && !user.IsOwner { - // Check if another owner exists - var existingOwner User - err := b.db.Where("bot_id = ? AND is_owner = ?", b.botID, true).First(&existingOwner).Error - if err == nil { - return User{}, fmt.Errorf("a bot can have only one owner") - } else if !errors.Is(err, gorm.ErrRecordNotFound) { - return User{}, err - } - // Promote to owner - role, err := b.getRoleByName("owner") - if err != nil { - return User{}, err - } - user.RoleID = role.ID - user.IsOwner = true - if err := b.db.Save(&user).Error; err != nil { - return User{}, err - } + return User{}, fmt.Errorf("cannot change existing user to owner") } } @@ -274,12 +267,17 @@ func (b *Bot) isAdminOrOwner(userID int64) bool { return user.Role.Name == "admin" || user.Role.Name == "owner" } -func initTelegramBot(token string, handleUpdate func(ctx context.Context, tgBot *bot.Bot, update *models.Update)) (*bot.Bot, error) { +func initTelegramBot(token string, handleUpdate func(ctx context.Context, tgBot *bot.Bot, update *models.Update)) (TelegramClient, error) { opts := []bot.Option{ bot.WithDefaultHandler(handleUpdate), } - return bot.New(token, opts...) + tgBot, err := bot.New(token, opts...) + if err != nil { + return nil, err + } + + return tgBot, nil } func (b *Bot) sendResponse(ctx context.Context, chatID int64, text string, businessConnectionID string) error { @@ -369,3 +367,37 @@ func isEmoji(r rune) bool { (r >= 0x2600 && r <= 0x26FF) || // Misc symbols (r >= 0x2700 && r <= 0x27BF) // Dingbats } + +func (b *Bot) sendWhoAmI(ctx context.Context, chatID int64, userID int64, username string, businessConnectionID string) { + user, err := b.getOrCreateUser(userID, username, false) + if err != nil { + log.Printf("Error getting or creating user: %v", err) + b.sendResponse(ctx, chatID, "Sorry, I couldn't retrieve your information.", businessConnectionID) + return + } + + caser := cases.Title(language.English) + whoAmIMessage := fmt.Sprintf( + "👤 Your Information:\n\n"+ + "- Username: %s\n"+ + "- Role: %s", + user.Username, + 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 + 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) +} diff --git a/handlers.go b/handlers.go index 149b772..39b668e 100644 --- a/handlers.go +++ b/handlers.go @@ -42,6 +42,9 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U case "/stats": b.sendStats(ctx, chatID, userID, message.From.Username, businessConnectionID) return + case "/whoami": + b.sendWhoAmI(ctx, chatID, userID, message.From.Username, businessConnectionID) + return } } } @@ -82,6 +85,14 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U return } + // Update the username if it's empty or has changed + if user.Username != username { + user.Username = username + if err := b.db.Save(&user).Error; err != nil { + log.Printf("Error updating user username: %v", err) + } + } + 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 { @@ -109,7 +120,6 @@ func (b *Bot) handleUpdate(ctx context.Context, tgBot *bot.Bot, update *models.U assistantMessage := b.createMessage(chatID, 0, "", "assistant", response, false) if err := b.storeMessage(assistantMessage); err != nil { log.Printf("Error storing assistant message: %v", err) - return } b.addMessageToChatMemory(chatMemory, assistantMessage) } diff --git a/main.go b/main.go index c93a754..1e6c2aa 100644 --- a/main.go +++ b/main.go @@ -52,14 +52,24 @@ func main() { go func(cfg BotConfig) { defer wg.Done() - // Create Bot instance with RealClock + // Create Bot instance without TelegramClient initially realClock := RealClock{} - bot, err := NewBot(db, cfg, realClock) + bot, err := NewBot(db, cfg, realClock, nil) if err != nil { log.Printf("Error creating bot %s: %v", cfg.ID, err) return } + // Initialize TelegramClient with the bot's handleUpdate method + tgClient, err := initTelegramBot(cfg.TelegramToken, bot.handleUpdate) + if err != nil { + log.Printf("Error initializing Telegram client for bot %s: %v", cfg.ID, err) + return + } + + // Assign the TelegramClient to the bot + bot.tgBot = tgClient + // Start the bot log.Printf("Starting bot %s...", cfg.ID) bot.Start(ctx) diff --git a/rate_limiter_test.go b/rate_limiter_test.go index 5965a3e..7b85f38 100644 --- a/rate_limiter_test.go +++ b/rate_limiter_test.go @@ -1,10 +1,13 @@ package main import ( - "strings" + "context" + "fmt" "testing" "time" + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" "gorm.io/driver/sqlite" "gorm.io/gorm" ) @@ -121,8 +124,20 @@ func TestOwnerAssignment(t *testing.T) { currentTime: time.Now(), } - // Create the bot - bot, err := NewBot(db, config, mockClock) + // Initialize MockTelegramClient + mockTGClient := &MockTelegramClient{ + SendMessageFunc: func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) { + chatID, ok := params.ChatID.(int64) + if !ok { + return nil, fmt.Errorf("ChatID is not of type int64") + } + // Simulate successful message sending + return &models.Message{ID: 1, Chat: models.Chat{ID: chatID}}, nil + }, + } + + // Create the bot with the mock Telegram client + bot, err := NewBot(db, config, mockClock, mockTGClient) if err != nil { t.Fatalf("Failed to create bot: %v", err) } @@ -142,7 +157,7 @@ func TestOwnerAssignment(t *testing.T) { // Verify that the error message is appropriate expectedErrorMsg := "an owner already exists for this bot" - if err.Error() != expectedErrorMsg && !strings.Contains(err.Error(), "unique index") { + if err.Error() != expectedErrorMsg { t.Fatalf("Unexpected error message: %v", err) } @@ -156,34 +171,15 @@ func TestOwnerAssignment(t *testing.T) { t.Fatalf("Expected role 'admin', got '%s'", adminUser.Role.Name) } - // Assign owner role to a user from a different bot - otherBotConfig := BotConfig{ - ID: "other_bot", - TelegramToken: "OTHER_TELEGRAM_TOKEN", - MemorySize: 10, - MessagePerHour: 5, - MessagePerDay: 10, - TempBanDuration: "1m", - SystemPrompts: make(map[string]string), - Active: true, - OwnerTelegramID: 444444444, + // Attempt to change an existing user to owner + _, err = bot.getOrCreateUser(333333333, "AdminUser", true) + if err == nil { + t.Fatalf("Expected error when changing existing user to owner, but got none") } - otherBot, err := NewBot(db, otherBotConfig, mockClock) - if err != nil { - t.Fatalf("Failed to create other bot: %v", err) - } - - _, err = otherBot.getOrCreateUser(config.OwnerTelegramID, "OwnerOfOtherBot", true) - if err != nil { - t.Fatalf("Failed to assign existing owner to another bot: %v", err) - } - - // Verify multiple bots can have the same owner telegram ID - var ownerOfOtherBot User - err = db.Where("telegram_id = ? AND bot_id = ? AND is_owner = ?", config.OwnerTelegramID, otherBot.botID, true).First(&ownerOfOtherBot).Error - if err != nil { - t.Fatalf("Owner of other bot was not created: %v", err) + expectedErrorMsg = "cannot change existing user to owner" + if err.Error() != expectedErrorMsg { + t.Fatalf("Unexpected error message: %v", err) } } diff --git a/telegram_client.go b/telegram_client.go new file mode 100644 index 0000000..10dc9bf --- /dev/null +++ b/telegram_client.go @@ -0,0 +1,16 @@ +// telegram_client.go +package main + +import ( + "context" + + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" +) + +// TelegramClient defines the methods required from the Telegram bot. +type TelegramClient interface { + SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) + Start(ctx context.Context) + // Add other methods if needed. +} diff --git a/telegram_client_mock.go b/telegram_client_mock.go new file mode 100644 index 0000000..61520f7 --- /dev/null +++ b/telegram_client_mock.go @@ -0,0 +1,35 @@ +// telegram_client_mock.go +package main + +import ( + "context" + + "github.com/go-telegram/bot" + "github.com/go-telegram/bot/models" +) + +// MockTelegramClient is a mock implementation of TelegramClient for testing. +type MockTelegramClient struct { + // You can add fields to keep track of calls if needed. + SendMessageFunc func(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) + StartFunc func(ctx context.Context) // Optional: track Start calls +} + +// SendMessage mocks sending a message. +func (m *MockTelegramClient) SendMessage(ctx context.Context, params *bot.SendMessageParams) (*models.Message, error) { + if m.SendMessageFunc != nil { + return m.SendMessageFunc(ctx, params) + } + // Default behavior: return an empty message without error. + return &models.Message{}, nil +} + +// Start mocks starting the Telegram client. +func (m *MockTelegramClient) Start(ctx context.Context) { + if m.StartFunc != nil { + m.StartFunc(ctx) + } + // Default behavior: do nothing. +} + +// Add other mocked methods if your Bot uses more TelegramClient methods.