YaeMikoBot/database/psql/rp.go

package psql

import (
	"database/sql"
	"errors"

	"git.nix13.pw/scuroneko/laniakea"
	"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 RPRepository struct {
	db *sqlx.DB
}

func newRpRepository(db *sqlx.DB) RPRepository {
	return RPRepository{db}
}
func NewRPRepository(db *laniakea.DatabaseContext) RPRepository {
	return newRpRepository(db.PostgresSQL)
}

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) 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
}