lone-services/services/analysis/internal/face/milvus.go

package face

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"

	"lone-services/pkg/utils"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"github.com/zeromicro/go-zero/core/logx"
)

const (
	collectionName   = "face_vectors"
	faceVectorField  = "face_vector"
	faceUserIDField  = "user_id"
	searchTopK       = 1
	similarityMin    = float32(0.55)
	hnswM            = 16
	hnswEfConstruct  = 200
	hnswEf           = 64
	collectionShards = 2
)

var (
	milvusMu     sync.RWMutex
	milvusClient client.Client
)

func InitMilvus() error {
	ctx := context.Background()
	host := utils.GetConfigString("milvus.milvus_host")
	port := utils.GetConfigString("milvus.milvus_grpc_port")
	if host == utils.StringEmpty || port == utils.StringEmpty {
		return fmt.Errorf("milvus config empty")
	}
	address := fmt.Sprintf("%s:%s", host, port)
	c, err := client.NewGrpcClient(ctx, address)
	if err != nil {
		return fmt.Errorf("connect milvus %s: %w", address, err)
	}
	if err = ensureCollection(ctx, c); err != nil {
		_ = c.Close()
		return err
	}
	if err = ensureIndex(ctx, c); err != nil {
		_ = c.Close()
		return err
	}
	if err = c.LoadCollection(ctx, collectionName, false); err != nil {
		_ = c.Close()
		return err
	}

	milvusMu.Lock()
	if milvusClient != nil {
		_ = milvusClient.Close()
	}
	milvusClient = c
	milvusMu.Unlock()
	logx.Infof("milvus connected: %s", address)
	return nil
}

func CloseMilvus() {
	milvusMu.Lock()
	defer milvusMu.Unlock()
	if milvusClient != nil {
		_ = milvusClient.Close()
		milvusClient = nil
	}
}

func getMilvus() client.Client {
	milvusMu.RLock()
	defer milvusMu.RUnlock()
	return milvusClient
}

func ensureCollection(ctx context.Context, c client.Client) error {
	has, err := c.HasCollection(ctx, collectionName)
	if err != nil {
		return err
	}
	if has {
		return nil
	}
	schema := &entity.Schema{
		CollectionName: collectionName,
		Description:    "user face vectors (ArcFace 512d)",
		AutoID:         false,
		Fields: []*entity.Field{
			{
				Name:       faceUserIDField,
				DataType:   entity.FieldTypeInt64,
				PrimaryKey: true,
			},
			{
				Name:     faceVectorField,
				DataType: entity.FieldTypeFloatVector,
				TypeParams: map[string]string{
					"dim": fmt.Sprintf("%d", faceEmbeddingDim),
				},
			},
		},
	}
	return c.CreateCollection(ctx, schema, int32(collectionShards))
}

func ensureIndex(ctx context.Context, c client.Client) error {
	indexes, err := c.DescribeIndex(ctx, collectionName, faceVectorField)
	if err != nil {
		if !strings.Contains(err.Error(), "index not found") {
			return err
		}
	} else if len(indexes) > utils.NumberZero {
		return nil
	}
	idx, err := entity.NewIndexHNSW(entity.IP, hnswM, hnswEfConstruct)
	if err != nil {
		return err
	}
	return c.CreateIndex(ctx, collectionName, faceVectorField, idx, false)
}

func UpsertFace(userID int64, vector []float32) error {
	c := getMilvus()
	if c == nil {
		return errors.New("milvus client is nil")
	}
	if userID < utils.NumberOne || len(vector) != faceEmbeddingDim {
		return errors.New("invalid face upsert args")
	}
	ctx := context.Background()
	_ = c.Delete(ctx, collectionName, "", fmt.Sprintf("%s == %d", faceUserIDField, userID))
	_, err := c.Insert(ctx, collectionName, "",
		entity.NewColumnInt64(faceUserIDField, []int64{userID}),
		entity.NewColumnFloatVector(faceVectorField, faceEmbeddingDim, [][]float32{vector}),
	)
	return err
}

func SearchUser(vector []float32) (int64, float32, error) {
	c := getMilvus()
	if c == nil {
		return 0, 0, errors.New("milvus client is nil")
	}
	if len(vector) != faceEmbeddingDim {
		return 0, 0, errors.New("invalid face vector")
	}
	ctx := context.Background()
	searchParam, err := entity.NewIndexHNSWSearchParam(hnswEf)
	if err != nil {
		return 0, 0, err
	}
	result, err := c.Search(
		ctx,
		collectionName,
		nil,
		"",
		[]string{faceUserIDField},
		[]entity.Vector{entity.FloatVector(vector)},
		faceVectorField,
		entity.IP,
		searchTopK,
		searchParam,
	)
	if err != nil {
		return 0, 0, err
	}
	if len(result) == utils.NumberZero || result[0].IDs.Len() == utils.NumberZero {
		return 0, 0, errors.New("face not found")
	}
	ids, ok := result[0].IDs.(*entity.ColumnInt64)
	if !ok {
		return 0, 0, errors.New("invalid milvus id column")
	}
	userID, err := ids.ValueByIdx(0)
	if err != nil {
		return 0, 0, err
	}
	score := result[0].Scores[0]
	if score < similarityMin {
		return 0, score, errors.New("face similarity too low")
	}
	return userID, score, nil
}