188 lines
6.7 KiB
Go
188 lines
6.7 KiB
Go
package dao
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
|
|
"rag-local/common"
|
|
"rag-local/kb/consts"
|
|
"rag-local/kb/model/entity"
|
|
|
|
"github.com/gogf/gf/v2/database/gdb"
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
"github.com/gogf/gf/v2/os/gtime"
|
|
"github.com/gogf/gf/v2/util/gconv"
|
|
)
|
|
|
|
var ModelConfig = &modelConfigDao{}
|
|
|
|
type modelConfigDao struct{}
|
|
|
|
func init() {
|
|
ctx := context.Background()
|
|
_, err := g.DB(consts.DbGroupSystem).Exec(ctx, `CREATE TABLE IF NOT EXISTS `+consts.TableNameModelConfig+` (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT NOT NULL DEFAULT '',
|
|
model_type TEXT NOT NULL DEFAULT 'chat',
|
|
model_name TEXT NOT NULL DEFAULT '',
|
|
endpoint_url TEXT NOT NULL DEFAULT '',
|
|
api_key TEXT NOT NULL DEFAULT '',
|
|
dimension INTEGER NOT NULL DEFAULT 1024,
|
|
extra TEXT NOT NULL DEFAULT '',
|
|
is_default INTEGER NOT NULL DEFAULT 0,
|
|
created_at DATETIME DEFAULT (datetime('now','localtime')),
|
|
updated_at DATETIME DEFAULT (datetime('now','localtime'))
|
|
)`)
|
|
if err != nil {
|
|
g.Log().Warningf(ctx, "create model_config table failed: %v", err)
|
|
}
|
|
// 迁移：旧库补 is_default 列，并把 system_config 中的旧默认对话模型键迁移为 is_default
|
|
// 注意：GetScan 只支持 struct/slice 指针，标量查询用 Value
|
|
cnt, err := g.DB(consts.DbGroupSystem).Ctx(ctx).GetValue(ctx,
|
|
"SELECT COUNT(*) FROM pragma_table_info('"+consts.TableNameModelConfig+"') WHERE name=?", "is_default")
|
|
if err == nil && cnt.Int64() == 0 {
|
|
if _, err := g.DB(consts.DbGroupSystem).Exec(ctx, "ALTER TABLE "+consts.TableNameModelConfig+" ADD COLUMN is_default INTEGER NOT NULL DEFAULT 0"); err != nil {
|
|
g.Log().Warningf(ctx, "alter model_config add is_default failed: %v", err)
|
|
} else {
|
|
old, err := g.DB(consts.DbGroupSystem).Ctx(ctx).GetValue(ctx,
|
|
"SELECT cfg_value FROM system_config WHERE cfg_key='default_chat_model'")
|
|
if err == nil && old.Int64() > 0 {
|
|
if _, err := g.DB(consts.DbGroupSystem).Exec(ctx,
|
|
"UPDATE "+consts.TableNameModelConfig+" SET is_default=CASE WHEN id=? THEN 1 ELSE 0 END WHERE model_type=?",
|
|
old.Int64(), consts.ModelTypeChat); err != nil {
|
|
g.Log().Warningf(ctx, "migrate default_chat_model to is_default failed: %v", err)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
// system_config 表已废弃（令牌改内存、默认模型迁入 is_default），清理残留
|
|
if _, err := g.DB(consts.DbGroupSystem).Exec(ctx, "DROP TABLE IF EXISTS system_config"); err != nil {
|
|
g.Log().Warningf(ctx, "drop system_config table failed: %v", err)
|
|
}
|
|
}
|
|
|
|
func (d *modelConfigDao) GetOne(ctx context.Context, id int64) (*entity.ModelConfig, error) {
|
|
var m entity.ModelConfig
|
|
err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).
|
|
Cache(gdb.CacheOption{Duration: common.CacheTTL(), Name: "model_config_GetOne_" + gconv.String(id)}).
|
|
Where("id", id).Scan(&m)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
// GetByName 按 model_name 查找已有配置（排除指定 id），供新配置继承端点与维度
|
|
func (d *modelConfigDao) GetByName(ctx context.Context, modelName string, excludeId int64) (*entity.ModelConfig, error) {
|
|
var m entity.ModelConfig
|
|
err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).
|
|
Where("model_name", modelName).
|
|
Where("id <> ?", excludeId).
|
|
OrderAsc("id").
|
|
Limit(1).
|
|
Scan(&m)
|
|
if err != nil {
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
if m.Id == 0 {
|
|
return nil, nil
|
|
}
|
|
return &m, nil
|
|
}
|
|
|
|
func (d *modelConfigDao) List(ctx context.Context, modelType string) ([]*entity.ModelConfig, error) {
|
|
var list []*entity.ModelConfig
|
|
m := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).OrderAsc("id")
|
|
if modelType != "" {
|
|
m = m.Where("model_type", modelType)
|
|
}
|
|
err := m.Scan(&list)
|
|
if list == nil {
|
|
list = make([]*entity.ModelConfig, 0)
|
|
}
|
|
return list, err
|
|
}
|
|
|
|
func (d *modelConfigDao) Insert(ctx context.Context, data *entity.ModelConfig) (int64, error) {
|
|
now := gtime.Now().Format("Y-m-d H:i:s")
|
|
if data.IsDefault == 1 {
|
|
// 新配置设为默认时，同类型旧默认先清零，保证唯一
|
|
if _, err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).
|
|
Where("model_type", data.ModelType).Update(g.Map{"is_default": 0}); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
r, err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).Data(g.Map{
|
|
"name": data.Name,
|
|
"model_type": data.ModelType,
|
|
"model_name": data.ModelName,
|
|
"endpoint_url": data.EndpointUrl,
|
|
"api_key": data.ApiKey,
|
|
"dimension": data.Dimension,
|
|
"extra": data.Extra,
|
|
"is_default": data.IsDefault,
|
|
"created_at": now,
|
|
"updated_at": now,
|
|
}).Insert()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
return r.LastInsertId()
|
|
}
|
|
|
|
// GetDefault 返回指定类型的默认模型配置 id；无默认返回 0
|
|
func (d *modelConfigDao) GetDefault(ctx context.Context, modelType string) (int64, error) {
|
|
v, err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).
|
|
Where("model_type", modelType).Where("is_default", 1).Value("id")
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if v.IsEmpty() {
|
|
return 0, nil
|
|
}
|
|
return v.Int64(), nil
|
|
}
|
|
|
|
// SetDefault 设置 id 为同类型默认，其余同类型取消默认（单 SQL 保证唯一）
|
|
func (d *modelConfigDao) SetDefault(ctx context.Context, id int64) error {
|
|
_, err := g.DB(consts.DbGroupSystem).Exec(ctx,
|
|
"UPDATE "+consts.TableNameModelConfig+" SET is_default=CASE WHEN id=? THEN 1 ELSE 0 END WHERE model_type=(SELECT model_type FROM "+consts.TableNameModelConfig+" WHERE id=?)",
|
|
id, id)
|
|
return err
|
|
}
|
|
|
|
func (d *modelConfigDao) Update(ctx context.Context, data *entity.ModelConfig) error {
|
|
_, err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).Data(g.Map{
|
|
"name": data.Name,
|
|
"model_type": data.ModelType,
|
|
"model_name": data.ModelName,
|
|
"endpoint_url": data.EndpointUrl,
|
|
"api_key": data.ApiKey,
|
|
"dimension": data.Dimension,
|
|
"extra": data.Extra,
|
|
"updated_at": gtime.Now().Format("Y-m-d H:i:s"),
|
|
}).Where("id", data.Id).Update()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// 查询缓存挂在 DB 内部缓存实例上且键带 SelectCache: 前缀，全局 gcache.Remove 无效
|
|
_, _ = g.DB(consts.DbGroupSystem).GetCache().Remove(ctx, "SelectCache:model_config_GetOne_"+gconv.String(data.Id))
|
|
return nil
|
|
}
|
|
|
|
func (d *modelConfigDao) Delete(ctx context.Context, id int64) error {
|
|
_, err := g.DB(consts.DbGroupSystem).Model(consts.TableNameModelConfig).Ctx(ctx).Where("id", id).Delete()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
_, _ = g.DB(consts.DbGroupSystem).GetCache().Remove(ctx, "SelectCache:model_config_GetOne_"+gconv.String(id))
|
|
return nil
|
|
}
|