Started implementing owner feature

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,

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

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,

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

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."

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
 }

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