aiforge/models/ai_model_manage.go

package models

import (
	"fmt"

	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/timeutil"
	"xorm.io/builder"
	"xorm.io/xorm"
)

type AiModelManage struct {
	ID            string             `xorm:"pk"`
	Name          string             `xorm:"NOT NULL"`
	Version       string             `xorm:"NOT NULL"`
	VersionCount  int                `xorm:"NOT NULL DEFAULT 0"`
	New           int                `xorm:"NOT NULL"`
	Type          int                `xorm:"NOT NULL"`
	Size          int64              `xorm:"NOT NULL"`
	Description   string             `xorm:"varchar(2000)"`
	Label         string             `xorm:"varchar(1000)"`
	Path          string             `xorm:"varchar(400) NOT NULL"`
	DownloadCount int                `xorm:"NOT NULL DEFAULT 0"`
	Engine        int64              `xorm:"NOT NULL DEFAULT 0"`
	Status        int                `xorm:"NOT NULL DEFAULT 0"`
	Accuracy      string             `xorm:"varchar(1000)"`
	AttachmentId  string             `xorm:"NULL"`
	RepoId        int64              `xorm:"NULL"`
	CodeBranch    string             `xorm:"varchar(400) NULL"`
	CodeCommitID  string             `xorm:"NULL"`
	UserId        int64              `xorm:"NOT NULL"`
	UserName      string             `xorm:"NULL"`
	TrainTaskInfo string             `xorm:"text NULL"`
	CreatedUnix   timeutil.TimeStamp `xorm:"created"`
	UpdatedUnix   timeutil.TimeStamp `xorm:"INDEX updated"`
}

type AiModelQueryOptions struct {
	ListOptions
	RepoID   int64 // include all repos if empty
	UserID   int64
	ModelID  string
	SortType string
	New      int
	// JobStatus     CloudbrainStatus
	Type int
}

func SaveModelToDb(model *AiModelManage) error {
	sess := x.NewSession()
	defer sess.Close()

	re, err := sess.Insert(model)
	if err != nil {
		log.Info("insert error." + err.Error())
		return err
	}
	log.Info("success to save db.re=" + fmt.Sprint((re)))
	return nil
}

func QueryModelById(id string) (*AiModelManage, error) {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table("ai_model_manage").
		Where("id='" + id + "'")
	aiModelManageList := make([]*AiModelManage, 0)
	err := sess.Find(&aiModelManageList)
	if err == nil {
		if len(aiModelManageList) == 1 {
			return aiModelManageList[0], nil
		}
	}
	return nil, err
}

func DeleteModelById(id string) error {
	sess := x.NewSession()
	defer sess.Close()

	re, err := sess.Delete(&AiModelManage{
		ID: id,
	})
	if err != nil {
		return err
	}
	log.Info("success to delete from  db.re=" + fmt.Sprint((re)))
	return nil

}

func ModifyModelDescription(id string, description string) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("description").Update(&AiModelManage{
		Description: description,
	})
	if err != nil {
		return err
	}
	log.Info("success to update description from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelNewProperty(id string, new int, versioncount int) error {
	var sess *xorm.Session
	sess = x.ID(id)
	defer sess.Close()
	re, err := sess.Cols("new", "version_count").Update(&AiModelManage{
		New:          new,
		VersionCount: versioncount,
	})
	if err != nil {
		return err
	}
	log.Info("success to update new property from  db.re=" + fmt.Sprint((re)))
	return nil
}

func ModifyModelDownloadCount(id string) error {
	sess := x.NewSession()
	defer sess.Close()
	if _, err := sess.Exec("UPDATE `ai_model_manage` SET download_count = download_count + 1 WHERE id = ?", id); err != nil {
		return err
	}

	return nil
}

func QueryModelByName(name string, repoId int64) []*AiModelManage {
	sess := x.NewSession()
	defer sess.Close()
	sess.Select("*").Table("ai_model_manage").
		Where("name='" + name + "' and repo_id=" + fmt.Sprint(repoId)).OrderBy("version desc")
	aiModelManageList := make([]*AiModelManage, 0)
	sess.Find(&aiModelManageList)
	return aiModelManageList
}

func QueryModel(opts *AiModelQueryOptions) ([]*AiModelManage, int64, error) {
	sess := x.NewSession()
	defer sess.Close()

	var cond = builder.NewCond()
	if opts.RepoID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.repo_id": opts.RepoID},
		)
	}

	if opts.UserID > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.user_id": opts.UserID},
		)
	}

	if opts.New >= 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.new": opts.New},
		)
	}

	if len(opts.ModelID) > 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.id": opts.ModelID},
		)
	}

	if (opts.Type) >= 0 {
		cond = cond.And(
			builder.Eq{"ai_model_manage.type": opts.Type},
		)
	}

	count, err := sess.Where(cond).Count(new(AiModelManage))
	if err != nil {
		return nil, 0, fmt.Errorf("Count: %v", err)
	}

	if opts.Page >= 0 && opts.PageSize > 0 {
		var start int
		if opts.Page == 0 {
			start = 0
		} else {
			start = (opts.Page - 1) * opts.PageSize
		}
		sess.Limit(opts.PageSize, start)
	}

	sess.OrderBy("ai_model_manage.created_unix DESC")
	aiModelManages := make([]*AiModelManage, 0, setting.UI.IssuePagingNum)
	if err := sess.Table("ai_model_manage").Where(cond).
		Find(&aiModelManages); err != nil {
		return nil, 0, fmt.Errorf("Find: %v", err)
	}
	sess.Close()

	return aiModelManages, count, nil
}