feat: rag初始版

dao/dataset_index.go

@@ -0,0 +1,59 @@
+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
+}