60 lines
1.6 KiB
Go
60 lines
1.6 KiB
Go
|
package dao
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"database/sql"
|
||
|
|
"fmt"
|
||
|
|
"rag/consts/public"
|
||
|
|
"rag/model/entity"
|
||
|
|
|
||
|
|
"gitea.com/red-future/common/db/gfdb"
|
||
|
|
)
|
||
|
|
|
||
|
|
var DatasetIndex = new(datasetIndexDao)
|
||
|
|
|
||
|
|
type datasetIndexDao struct{}
|
||
|
|
|
||
|
|
// Insert 插入数据集索引
|
||
|
|
func (d *datasetIndexDao) Insert(ctx context.Context, index *entity.DatasetIndex) (id int64, err error) {
|
||
|
|
_, err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).Data(index).Insert()
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
return 0, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// GetByDatasetId 根据数据集ID获取索引
|
||
|
|
func (d *datasetIndexDao) GetByDatasetId(ctx context.Context, datasetId int64) (result *entity.DatasetIndex, err error) {
|
||
|
|
err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).Where(entity.DatasetIndexCol.DatasetId, datasetId).Scan(&result)
|
||
|
|
if err != nil {
|
||
|
|
if err == sql.ErrNoRows {
|
||
|
|
return nil, nil
|
||
|
|
}
|
||
|
|
return
|
||
|
|
}
|
||
|
|
return result, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// IncVectorCount 增加或减少向量数量
|
||
|
|
func (d *datasetIndexDao) IncVectorCount(ctx context.Context, id int64, delta int64) (err error) {
|
||
|
|
_, err = gfdb.DB(ctx).Model(ctx, public.TableNameDatasetIndex).
|
||
|
|
Where(entity.DatasetIndexCol.Id, id).
|
||
|
|
Increment(entity.DatasetIndexCol.VectorCount, delta)
|
||
|
|
return
|
||
|
|
}
|
||
|
|
|
||
|
|
func (d *datasetIndexDao) InsertIndex(ctx context.Context, indexName string) (err error) {
|
||
|
|
prefix, err := gfdb.GetTablePrefix(ctx)
|
||
|
|
if err != nil {
|
||
|
|
return
|
||
|
|
}
|
||
|
|
sqlStr := fmt.Sprintf(`
|
||
|
|
CREATE INDEX IF NOT EXISTS %s
|
||
|
|
ON %s
|
||
|
|
USING ivfflat (vector vector_cosine_ops)
|
||
|
|
WHERE vector IS NOT NULL;
|
||
|
|
`, indexName, prefix+public.TableNameDocumentChunk)
|
||
|
|
_, err = gfdb.DB(ctx).Exec(ctx, sqlStr)
|
||
|
|
return
|
||
|
|
}
|