many new

database/mdb/rp_chats.go

@@ -0,0 +1,40 @@
+package mdb
+
+import (
+	"context"
+	"kurumibot/database"
+	"kurumibot/laniakea"
+	"time"
+
+	"go.mongodb.org/mongo-driver/v2/bson"
+)
+
+type RPChatMessage struct {
+	ChatID  string `bson:"chat_id"`
+	Role    string `bson:"role"`
+	Message string `bson:"message"`
+}
+
+func GetChatHistory(db *laniakea.DatabaseContext, chatId string) ([]*RPChatMessage, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	col := database.GetMongoCollection(db, "rp_chat_history")
+	cursor, err := col.Find(ctx, bson.M{"chat_id": chatId})
+	if err != nil {
+		return nil, err
+	}
+	result := make([]*RPChatMessage, 0)
+	err = cursor.All(ctx, &result)
+	return result, err
+}
+func UpdateChatHistory(db *laniakea.DatabaseContext, chatId, role, message string) error {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	col := database.GetMongoCollection(db, "rp_chat_history")
+	_, err := col.InsertOne(ctx, RPChatMessage{
+		chatId,
+		role,
+		message,
+	})
+	return err
+}

database/psql/users.go

@@ -13,17 +13,16 @@ import (
 )
 
 type User struct {
-	ID         int
-	TelegramID int `gorm:"column:tg_id"`
-	Balance    decimal.Decimal
-	Name       string
-	GroupID    int `gorm:"default:1"`
-	Group      *Group
-	Level      int `gorm:"default:1"`
-	Exp        int `gorm:"default:0"`
-	WorkID     int `gorm:"default:1"`
-	Work       *Work
-	WorkTime   time.Time
+	ID       int
+	Balance  decimal.Decimal
+	Name     string
+	GroupID  int `gorm:"default:1"`
+	Group    *Group
+	Level    int `gorm:"default:1"`
+	Exp      int `gorm:"default:0"`
+	WorkID   int `gorm:"default:1"`
+	Work     *Work
+	WorkTime time.Time
 
 	AutoID     sql.NullInt64
 	Auto       *ShopAuto
@@ -64,10 +63,10 @@ func GetOrCreateUser(tgId int, name string) (*User, error) {
 	return user, err
 }
 
-func CreateUser(tgId int, name string) (*User, error) {
+func CreateUser(id int, name string) (*User, error) {
 	user := &User{
-		TelegramID: tgId,
-		Name:       name,
+		ID:   id,
+		Name: name,
 	}
 	tx := database.PostgresDatabase.Create(user)
 	return user, tx.Error
@@ -75,7 +74,7 @@ func CreateUser(tgId int, name string) (*User, error) {
 
 func GetUser(telegramId int) (*User, error) {
 	user := new(User)
-	tx := database.PostgresDatabase.Joins("Group").Joins("Work").Joins("Auto").Joins("Business").Joins("Maid").Joins("Miner").Joins("Fraction").Preload("Pair").Take(user, "tg_id=?", telegramId)
+	tx := database.PostgresDatabase.Joins("Group").Joins("Work").Joins("Auto").Joins("Business").Joins("Maid").Joins("Miner").Joins("Fraction").Preload("Pair").Take(user, "users.id=?", telegramId)
 	return user, tx.Error
 }
 

database/red/rp_chats.go

@@ -0,0 +1,44 @@
+package red
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"kurumibot/laniakea"
+)
+
+var ctx = context.Background()
+
+func RPSetSelectedWaifu(db *laniakea.DatabaseContext, userId, waifuId uint) error {
+	key := fmt.Sprintf("ai.chats.rp.%d", userId)
+	return db.Redis.Set(ctx, key, waifuId, 0).Err()
+}
+func RPGetSelectedWaifu(db *laniakea.DatabaseContext, userId uint) uint {
+	key := fmt.Sprintf("ai.chats.rp.%d", userId)
+	res := db.Redis.Get(ctx, key)
+	if res.Err() != nil {
+		return 0
+	}
+	i, _ := res.Int()
+	return uint(i)
+}
+func GetChatId(db *laniakea.DatabaseContext, userId uint) string {
+	waifuId := RPGetSelectedWaifu(db, userId)
+	if waifuId == 0 {
+		return "0"
+	}
+	key := fmt.Sprintf("ai.chats.rp.%d.%d", userId, waifuId)
+	res := db.Redis.Get(ctx, key)
+	if res.Err() != nil {
+		return ""
+	}
+	return res.Val()
+}
+func SaveChatId(db *laniakea.DatabaseContext, userId uint, chatId string) error {
+	waifuId := RPGetSelectedWaifu(db, userId)
+	if waifuId == 0 {
+		return errors.New("no chat found")
+	}
+	key := fmt.Sprintf("ai.chats.rp.%d.%d", userId, waifuId)
+	return db.Redis.Set(context.Background(), key, chatId, 0).Err()
+}

database/redis.go

@@ -1,12 +1,17 @@
 package database
 
-import "github.com/redis/go-redis/v9"
+import (
+	"fmt"
+	"os"
+
+	"github.com/redis/go-redis/v9"
+)
 
 var RedisClient *redis.Client
 
 func ConnectRedis() {
 	RedisClient = redis.NewClient(&redis.Options{
-		Addr:     "redis:6379",
+		Addr:     fmt.Sprintf("%s:6379", os.Getenv("REDIS_HOST")),
 		Password: "",
 		DB:       0,
 	})