ai work

database/psql/rp.go

@@ -0,0 +1,88 @@
+package psql
+
+import (
+	"database/sql"
+	"errors"
+
+	"github.com/vinovest/sqlx"
+)
+
+type RPGeneralPreset struct {
+	ID          string
+	Name        string
+	Description string
+	PreHistory  string `db:"pre_history"`
+	PostHistory string `db:"post_history"`
+}
+type RPScenarios struct {
+	ID     int
+	Name   string
+	Prompt string
+}
+type RPUser struct {
+	UserID         int64  `db:"user_id"`
+	SelectedPreset string `db:"selected_preset"`
+	UsedTokens     int64  `db:"used_tokens"`
+}
+
+type RPRepository struct {
+	db *sqlx.DB
+}
+
+func NewRPRepository(db *sqlx.DB) *RPRepository {
+	return &RPRepository{db: db}
+}
+func (rep *RPRepository) GetOrCreateUser(id int64) (*RPUser, error) {
+	user, err := rep.GetUser(id)
+	if errors.Is(err, sql.ErrNoRows) {
+		user, err = rep.CreateUser(id)
+	}
+	return user, err
+}
+func (rep *RPRepository) CreateUser(id int64) (*RPUser, error) {
+	user := new(RPUser)
+	err := rep.db.Get(user, "INSERT INTO rp_users(user_id) VALUES ($1) RETURNING *;", id)
+	return user, err
+}
+func (rep *RPRepository) GetUser(id int64) (*RPUser, error) {
+	user := new(RPUser)
+	err := rep.db.Get(user, "SELECT * FROM rp_users WHERE user_id=$1", id)
+	return user, err
+}
+func (rep *RPRepository) UpdateUser(user *RPUser) error {
+	_, err := rep.db.NamedExec(
+		"UPDATE rp_users SET selected_preset=:selected_preset, used_tokens=:used_tokens WHERE user_id=:user_id;",
+		user,
+	)
+	return err
+}
+func (rep *RPRepository) UpdateUserPreset(user *RPUser, presetId string) (*RPGeneralPreset, error) {
+	preset, err := rep.GetPreset(presetId)
+	if err != nil {
+		return preset, err
+	}
+	_, err = rep.db.Exec("UPDATE rp_users SET selected_preset=$1 WHERE user_id=$2;", presetId, user.UserID)
+	return preset, err
+}
+func (rep *RPRepository) GetUserPreset(user *RPUser) (*RPGeneralPreset, 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() ([]*RPGeneralPreset, error) {
+	presets := make([]*RPGeneralPreset, 0)
+	err := rep.db.Select(&presets, "SELECT * FROM rp_general_presets;")
+	return presets, err
+}
+func (rep *RPRepository) GetPreset(id string) (*RPGeneralPreset, error) {
+	preset := new(RPGeneralPreset)
+	err := rep.db.Get(preset, "SELECT * FROM rp_general_presets WHERE id=$1;", id)
+	return preset, err
+}
+
+func (rep *RPRepository) GetAllScenarios() ([]*RPScenarios, error) {
+	return nil, nil
+}