package psql import ( "database/sql" "errors" "ymgb/database" "github.com/vinovest/sqlx" ) type RPPreset struct { ID string Name string Description string PreHistory string `db:"pre_history"` PostHistory string `db:"post_history"` } type RPScenario struct { ID int Name string Description string Prompt string } type RPSetting struct { ID int Name string Description string Prompt string } type RPUser struct { UserID int64 `db:"user_id"` UserPrompt string `db:"user_prompt"` UsedTokens int64 `db:"used_tokens"` SelectedPreset string `db:"selected_preset"` Preset *RPPreset SelectedModel string `db:"selected_model"` Model *AIModel CompressMethod string `db:"compress_method"` CompressLimit int `db:"compress_limit"` } type RPWaifuPrompts struct { WaifuID int `db:"waifu_id"` Waifu *Waifu AppearancePrompt string `db:"appearance_prompt"` PersonalityPrompt string `db:"personality_prompt"` } type RPRepository struct { db *sqlx.DB } func newRpRepository(db *sqlx.DB) RPRepository { return RPRepository{db} } func NewRPRepository(db *database.Context) RPRepository { return newRpRepository(db.Postgres) } func (rep RPRepository) GetOrCreateUser(id int64) (RPUser, error) { user, err := rep.GetUser(id) if errors.Is(err, sql.ErrNoRows) { return rep.CreateUser(id) } return user, err } func (rep RPRepository) CreateUser(id int64) (RPUser, error) { user := RPUser{} err := rep.db.Get(&user, "INSERT INTO rp_users(user_id) VALUES ($1) RETURNING *;", id) if err != nil { return user, err } preset, err := rep.GetPreset(user.SelectedPreset) if err != nil { return user, err } aiRep := newAiRepository(rep.db) model, err := aiRep.GetModel(user.SelectedModel) if err != nil { return user, err } user.Preset = &preset user.Model = &model return user, err } func (rep RPRepository) GetUser(id int64) (RPUser, error) { user := RPUser{} err := rep.db.Get(&user, "SELECT * FROM rp_users WHERE user_id=$1", id) if err != nil { return user, err } preset, err := rep.GetPreset(user.SelectedPreset) if err != nil { return user, err } aiRep := newAiRepository(rep.db) model, err := aiRep.GetModel(user.SelectedModel) user.Model = &model user.Preset = &preset return user, err } func (rep RPRepository) UpdateUser(user RPUser) error { s, args, err := sqlx.In( "UPDATE rp_users SET selected_preset=?, used_tokens=?, user_prompt=?, selected_model=? WHERE user_id=?;", user.SelectedPreset, user.UsedTokens, user.UserPrompt, user.SelectedModel, user.UserID, ) if err != nil { return err } s = rep.db.Rebind(s) _, err = rep.db.Exec(s, args...) return err } func (rep RPRepository) UpdateUserPreset(user RPUser, presetId string) (RPPreset, error) { preset, err := rep.GetPreset(presetId) if err != nil { return preset, err } s, args, err := sqlx.In( "UPDATE rp_users SET selected_preset=? WHERE user_id=?;", presetId, user.UserID, ) if err != nil { return preset, err } s = rep.db.Rebind(s) _, err = rep.db.Exec(s, args...) return preset, err } func (rep RPRepository) GetUserPreset(user RPUser) (RPPreset, error) { preset, err := rep.GetPreset(user.SelectedPreset) if errors.Is(err, sql.ErrNoRows) { return rep.UpdateUserPreset(user, "soft") } return preset, err } func (rep RPRepository) GetAllPresets() ([]RPPreset, error) { presets := make([]RPPreset, 0) err := rep.db.Select(&presets, "SELECT * FROM rp_presets ORDER BY id;") return presets, err } func (rep RPRepository) GetPreset(id string) (RPPreset, error) { preset := RPPreset{} err := rep.db.Get(&preset, "SELECT * FROM rp_presets WHERE id=$1;", id) return preset, err } func (rep RPRepository) GetAllScenarios() ([]RPScenario, error) { scenarios := make([]RPScenario, 0) err := rep.db.Select(&scenarios, "SELECT * FROM rp_scenarios ORDER BY id;") return scenarios, err } func (rep RPRepository) GetScenario(id int) (RPScenario, error) { scenario := RPScenario{} err := rep.db.Get(&scenario, "SELECT * FROM rp_scenarios WHERE id=$1;", id) return scenario, err } func (rep RPRepository) GetAllSettings() ([]RPSetting, error) { settings := make([]RPSetting, 0) err := rep.db.Select(&settings, "SELECT * FROM rp_settings ORDER BY id;") return settings, err } func (rep RPRepository) GetSetting(id int) (RPSetting, error) { setting := RPSetting{} err := rep.db.Get(&setting, "SELECT * FROM rp_settings WHERE id=$1;", id) return setting, err } func (rep RPRepository) GetWaifuPrompts(waifuId int) (RPWaifuPrompts, error) { prompts := RPWaifuPrompts{} err := rep.db.Get(&prompts, "SELECT * FROM rp_waifu_prompts WHERE waifu_id=$1;", waifuId) return prompts, err } func (rep RPRepository) UpdateUserCompressSettings(user RPUser) (RPUser, error) { query, args, err := sqlx.In( "UPDATE rp_users SET compress_method=?, compress_limit=? WHERE user_id=?;", user.CompressMethod, user.CompressLimit, user.UserID, ) if err != nil { return user, err } query = rep.db.Rebind(query) _, err = rep.db.Exec(query, args...) return user, err }