Support vllm and tei rerank

Signed-off-by: junjie.jiang <[email protected]>

Author: junjiejiangjjj

configs/milvus.yaml

@@ -1316,3 +1316,10 @@ function:
       voyageai:
         credential:  # The name in the crendential configuration item
         url:  # Your voyageai embedding url, Default is the official embedding url
+  rerank:
+    model:
+      providers:
+        tei:
+          enable: true # Whether to enable TEI rerank service
+        vllm:
+          enable: true # Whether to enable vllm rerank service

internal/proxy/task_search.go

@@ -675,20 +675,26 @@ func (t *searchTask) hybridSearchRank(ctx context.Context, span trace.Span, mult
 		for i := 0; i < len(multipleMilvusResults); i++ {
 			multipleMilvusResults[i].Results.FieldsData = fields[i]
 		}
-		params := rerank.NewSearchParams(
-			t.Nq, t.rankParams.limit, t.rankParams.offset, t.rankParams.roundDecimal,
-			t.rankParams.groupByFieldId, t.rankParams.groupSize, t.rankParams.strictGroupSize, searchMetrics,
-		)
-
-		if t.result, err = t.functionScore.Process(ctx, params, multipleMilvusResults); err != nil {
-			return err
+		{
+			ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-call-rerank-function-udf")
+			defer sp.End()
+			params := rerank.NewSearchParams(
+				t.Nq, t.rankParams.limit, t.rankParams.offset, t.rankParams.roundDecimal,
+				t.rankParams.groupByFieldId, t.rankParams.groupSize, t.rankParams.strictGroupSize, searchMetrics,
+			)
+			if t.result, err = t.functionScore.Process(ctx, params, multipleMilvusResults); err != nil {
+				return err
+			}
 		}
 		if fields, err := t.reorganizeRequeryResults(ctx, queryResult.GetFieldsData(), []*schemapb.IDs{t.result.Results.Ids}); err != nil {
 			return err
 		} else {
 			t.result.Results.FieldsData = fields[0]
 		}
 	} else {
+		ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-call-rerank-function-udf")
+		defer sp.End()
+
 		params := rerank.NewSearchParams(
 			t.Nq, t.rankParams.limit, t.rankParams.offset, t.rankParams.roundDecimal,
 			t.rankParams.groupByFieldId, t.rankParams.groupSize, t.rankParams.strictGroupSize, searchMetrics,
@@ -816,11 +822,15 @@ func (t *searchTask) searchPostProcess(ctx context.Context, span trace.Span, toR
 	}
 
 	if t.functionScore != nil && len(result.Results.FieldsData) != 0 {
-		params := rerank.NewSearchParams(t.Nq, t.SearchRequest.GetTopk(), t.SearchRequest.GetOffset(),
-			t.queryInfos[0].RoundDecimal, t.queryInfos[0].GroupByFieldId, t.queryInfos[0].GroupSize, t.queryInfos[0].StrictGroupSize, []string{metricType})
-		// rank only returns id and score
-		if t.result, err = t.functionScore.Process(ctx, params, []*milvuspb.SearchResults{result}); err != nil {
-			return err
+		{
+			ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-call-rerank-function-udf")
+			defer sp.End()
+			params := rerank.NewSearchParams(t.Nq, t.SearchRequest.GetTopk(), t.SearchRequest.GetOffset(),
+				t.queryInfos[0].RoundDecimal, t.queryInfos[0].GroupByFieldId, t.queryInfos[0].GroupSize, t.queryInfos[0].StrictGroupSize, []string{metricType})
+			// rank only returns id and score
+			if t.result, err = t.functionScore.Process(ctx, params, []*milvuspb.SearchResults{result}); err != nil {
+				return err
+			}
 		}
 		if !t.needRequery {
 			fields, err := t.reorganizeRequeryResults(ctx, result.Results.FieldsData, []*schemapb.IDs{t.result.Results.Ids})

internal/util/function/common.go

@@ -109,16 +109,17 @@ const (
 	siliconflowAKEnvStr string = "MILVUSAI_SILICONFLOW_API_KEY"
 )
 
-// TEI
+// TEI and vllm
 
 const (
 	ingestionPromptParamKey     string = "ingestion_prompt"
 	searchPromptParamKey        string = "search_prompt"
 	maxClientBatchSizeParamKey  string = "max_client_batch_size"
 	truncationDirectionParamKey string = "truncation_direction"
-	endpointParamKey            string = "endpoint"
+	EndpointParamKey            string = "endpoint"
 
-	enableTeiEnvStr string = "MILVUSAI_ENABLE_TEI"
+	EnableTeiEnvStr  string = "MILVUSAI_ENABLE_TEI"
+	EnableVllmEnvStr string = "MILVUSAI_ENABLE_VLLM"
 )
 
 func parseAKAndURL(credentials *credentials.Credentials, params []*commonpb.KeyValuePair, confParams map[string]string, apiKeyEnv string) (string, string, error) {

internal/util/function/rerank/decay_function.go

@@ -55,24 +55,13 @@ type DecayFunction[T PKType, R int32 | int64 | float32 | float64] struct {
 }
 
 func newDecayFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemapb.FunctionSchema) (Reranker, error) {
-	pkType := schemapb.DataType_None
-	for _, field := range collSchema.Fields {
-		if field.IsPrimaryKey {
-			pkType = field.DataType
-		}
-	}
-
-	if pkType == schemapb.DataType_None {
-		return nil, fmt.Errorf("Collection %s can not found pk field", collSchema.Name)
-	}
-
-	base, err := newRerankBase(collSchema, funcSchema, decayFunctionName, false, pkType)
+	base, err := newRerankBase(collSchema, funcSchema, decayFunctionName, false)
 	if err != nil {
 		return nil, err
 	}
 
 	if len(base.GetInputFieldNames()) != 1 {
-		return nil, fmt.Errorf("Decay function only supoorts single input, but gets [%s] input", base.GetInputFieldNames())
+		return nil, fmt.Errorf("Decay function only supports single input, but gets [%s] input", base.GetInputFieldNames())
 	}
 
 	var inputType schemapb.DataType
@@ -82,7 +71,7 @@ func newDecayFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemap
 		}
 	}
 
-	if pkType == schemapb.DataType_Int64 {
+	if base.pkType == schemapb.DataType_Int64 {
 		switch inputType {
 		case schemapb.DataType_Int8, schemapb.DataType_Int16, schemapb.DataType_Int32:
 			return newFunction[int64, int32](base, funcSchema)
@@ -160,7 +149,7 @@ func newFunction[T PKType, R int32 | int64 | float32 | float64](base *RerankBase
 	}
 
 	if decayFunc.decay <= 0 || decayFunc.decay >= 1 {
-		return nil, fmt.Errorf("Decay function param: decay must 0 < decay < 1, but got %f", decayFunc.offset)
+		return nil, fmt.Errorf("Decay function param: decay must 0 < decay < 1, but got %f", decayFunc.decay)
 	}
 
 	switch decayFunc.functionName {

internal/util/function/rerank/decay_function_test.go

@@ -90,7 +90,7 @@ func (s *DecayFunctionSuite) TestNewDecayErrors() {
 	{
 		functionSchema.InputFieldNames = []string{"ts", "pk"}
 		_, err := newDecayFunction(schema, functionSchema)
-		s.ErrorContains(err, "Decay function only supoorts single input, but gets")
+		s.ErrorContains(err, "Decay function only supports single input, but gets")
 	}
 
 	{

internal/util/function/rerank/function_score.go

@@ -32,6 +32,7 @@ import (
 
 const (
 	decayFunctionName string = "decay"
+	modelFunctionName string = "model"
 )
 
 type SearchParams struct {
@@ -92,8 +93,10 @@ func createFunction(collSchema *schemapb.CollectionSchema, funcSchema *schemapb.
 	switch rerankerName {
 	case decayFunctionName:
 		rerankFunc, newRerankErr = newDecayFunction(collSchema, funcSchema)
+	case modelFunctionName:
+		rerankFunc, newRerankErr = newModelFunction(collSchema, funcSchema)
 	default:
-		return nil, fmt.Errorf("Unsupported rerank function: [%s] , list of supported [%s]", rerankerName, decayFunctionName)
+		return nil, fmt.Errorf("Unsupported rerank function: [%s] , list of supported [%s,%s]", rerankerName, decayFunctionName, modelFunctionName)
 	}
 
 	if newRerankErr != nil {