Qdrant向量数据库的增删改查
Qdrant是一个开源的向量相似度搜索引擎,它提供了一个生产就绪的服务,通过便捷的API来存储、搜索和管理带有额外有效载荷的向量。
存储高维向量数据
快速进行相似度搜索
管理带有元数据的向量
支持多种距离度量方式
go版本操作
安装第三方库
go get github.com/qdrant/go-client/qdrant
func QdrantInit() {if err := viper.UnmarshalKey("qdrant", &globals.AppConfig.Qdrant); err != nil {globals.Log.Panicf("无法解码为结构: %s", err)}var err errorglobals.Qdrant, err = qdrant.NewClient(&qdrant.Config{// 192.168.10.4Host: globals.AppConfig.Qdrant.Host,// 6334 grpc端口Port: globals.AppConfig.Qdrant.Port,APIKey: globals.AppConfig.Qdrant.ApiKey,SkipCompatibilityCheck: true,})if err != nil {globals.Log.Panicf("Qdrant连接失败: %v", err)} else {globals.Log.Infof("Qdrant连接成功")}
}
// Collection 初始化向量存储
func Collection(client *qdrant.Client) {// 判断集合是否存在,存在的集合不能重复创建exists, err := client.CollectionExists(ctx, collectionName)if err != nil {globals.Log.Errorf("Collection->判断集合是否存在失败, err:%v", err)}if exists {// 清空指定集合err := client.DeleteCollection(ctx, collectionName)if err != nil {globals.Log.Errorf("Collection->清空集合失败, err:%v", err)}}// 创建集合err = client.CreateCollection(ctx,&qdrant.CreateCollection{CollectionName: "question_vector",VectorsConfig: qdrant.NewVectorsConfig(&qdrant.VectorParams{Size: 2560,Distance: qdrant.Distance_Cosine,}),})if err != nil {globals.Log.Errorf("Collection->创建集合失败, err:%v", err)return}
}
// StoreVector 存储向量
func StoreVector(list []string, msg []map[string]interface{}, client *qdrant.Client) {points := make([]*qdrant.PointStruct, len(list))// 生成向量vector, err := utils.GenerateVector(list)if err != nil || vector == nil {globals.Log.Errorf("StoreVector->生成向量失败, err:%v", err)return}for i, data := range vector.Vector {id, err := getID(msg[i])if err != nil {globals.Log.Errorf("StoreVector->获取id失败, err:%v", err)return}points[i] = &qdrant.PointStruct{Id: qdrant.NewIDNum(id),Vectors: qdrant.NewVectors(data.Values...),Payload: qdrant.NewValueMap(msg[i]),}}// 加入重试操作err = retry(3, 2*time.Second, func() error {// 批量插入向量_, err = client.Upsert(context.Background(), &qdrant.UpsertPoints{CollectionName: collectionName,Points: points,})return err})if err != nil {globals.Log.Errorf("StoreVector->批量插入向量失败, err:%v", err)return}
}
// DeleteVector 删除向量
func DeleteVector(id []int, client *qdrant.Client) {// 获取要删除向量的idpoints := make([]*qdrant.PointId, len(id))for i, v := range id {points[i] = qdrant.NewIDNum(uint64(v))}err := retry(3, 2*time.Second, func() error {// 根据id删除指定向量_, err := client.Delete(context.Background(), &qdrant.DeletePoints{CollectionName: collectionName,Points: qdrant.NewPointsSelectorIDs(points),})return err})if err != nil {globals.Log.Errorf("DeleteVector->删除向量失败: %v", err)}return
}
// UpdateVector 更新向量元数据
func UpdateVector(id []int, msg []map[string]interface{}, client *qdrant.Client) {// 批量修改向量元数据for i, m := range msg {m, err := getPayload(id[i], m, client)if err != nil {globals.Log.Errorf("UpdateVector->获取向量元数据失败: %v", err)continue}pointID := qdrant.NewIDNum(uint64(id[i]))payload := qdrant.NewValueMap(m)err = retry(3, 2*time.Second, func() error {_, err = client.SetPayload(context.Background(), &qdrant.SetPayloadPoints{CollectionName: collectionName,Payload: payload,PointsSelector: qdrant.NewPointsSelector(pointID),})return err})if err != nil {globals.Log.Errorf("UpdateVector->更新向量元数据失败: %v", err)continue}}
}
// SearchSimilar 搜索相似标题
func SearchSimilar(list []string, limit, offset *uint64, client *qdrant.Client) ([]map[string]interface{}, int, error) {// 获得缓存向量vec, exist := getCachedVector(list[0])if !exist { // 缓存中不存在该向量// 获取向量数据库中的向量vec, exist = getDatabaseVector(list[0], client)if !exist { // 向量数据库中不存在该向量vector, err := utils.GenerateVector(list)if err != nil || vector == nil {globals.Log.Errorf("SearchSimilar->生成向量失败, err:%v", err)return nil, 0, err}// 存储缓存向量cacheVector(list[0], vector.Vector[0].Values)vec = vector.Vector[0].Values}}// 当偏移量超过集合点数时,修正为0offset, err := correctOffset(offset, client)if err != nil {globals.Log.Errorf("SearchSimilar->修正offset失败, err:%v", err)return nil, 0, err}threshold := float32(questionThreshold)// 获得相似向量的总数countTotal, err := client.Query(context.Background(), &qdrant.QueryPoints{CollectionName: collectionName,Query: qdrant.NewQuery(vec...),ScoreThreshold: &threshold,})if err != nil {globals.Log.Errorf("SearchSimilar->查询总数失败, err:%v", err)return nil, 0, err}total := len(countTotal)// 搜索相似向量query, err := client.Query(context.Background(), &qdrant.QueryPoints{CollectionName: collectionName,Query: qdrant.NewQuery(vec...),ScoreThreshold: &threshold,WithPayload: qdrant.NewWithPayload(true),Limit: limit,Offset: offset,})if err != nil {globals.Log.Errorf("SearchSimilar->分页查询失败, err:%v", err)return nil, 0, err}// 解析查询结果中的 payloadvar res []map[string]interface{}for _, point := range query {payload := convertPayload(point.Payload)res = append(res, payload)}return res, total, nil
}
// getDatabaseVector 获取向量数据库中的向量
func getDatabaseVector(title string, client *qdrant.Client) ([]float32, bool) {limit, offset := uint64(1), uint64(0)query, err := client.Query(context.Background(),&qdrant.QueryPoints{CollectionName: collectionName,Filter: &qdrant.Filter{// 必须满足的过滤条件Must: []*qdrant.Condition{qdrant.NewMatchKeyword("title", title),},},Limit: &limit,Offset: &offset,WithVectors: qdrant.NewWithVectors(true),},)if err != nil || len(query) == 0 {globals.Log.Error("getDatabaseVector->查询存在的向量失败 error:", err)return nil, false}data := query[0].Vectors.GetVector().Datareturn data, true
}