feat: support /v1/rerank router

relay/controller/helper.go

@@ -57,7 +57,37 @@ func getImageRequest(c *gin.Context, relayMode int) (*relaymodel.ImageRequest, e
 	}
 	return imageRequest, nil
 }
+func getRerankRequest(c *gin.Context, relayMode int) (*relaymodel.RerankRequest, error) {
+	rerankRequest := &relaymodel.RerankRequest{}
+	err := common.UnmarshalBodyReusable(c, rerankRequest)
+	if err != nil {
+		return nil, err
+	}
+	if rerankRequest.Model == "" {
+		return nil, errors.New("model parameter must be provided")
+	}
+	// Set default values if necessary
+	if rerankRequest.TopN == nil {
+		defaultTopN := 10 // Default to returning top 10 results
+		rerankRequest.TopN = &defaultTopN
+	}
+	if rerankRequest.Query == "" {
+		return nil, errors.New("query must not be empty")
+	}
+	if len(rerankRequest.Documents) == 0 {
+		return nil, errors.New("document list must not be empty")
+	}
+	// if rerankRequest.MaxChunksPerDoc == nil {
+	//     defaultMaxChunks := 5 // Default maximum chunks per document
+	//     rerankRequest.MaxChunksPerDoc = &defaultMaxChunks
+	// }
+	if rerankRequest.ReturnDocuments == nil {
+		defaultReturnDocs := true // Default to returning documents
+		rerankRequest.ReturnDocuments = &defaultReturnDocs
+	}
 
+	return rerankRequest, nil
+}
 func isValidImageSize(model string, size string) bool {
 	if model == "cogview-3" {
 		return true

relay/controller/rerank.go

@@ -0,0 +1,62 @@
+package controller
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/songquanpeng/one-api/common/logger"
+	"github.com/songquanpeng/one-api/relay"
+	"github.com/songquanpeng/one-api/relay/adaptor/openai"
+	"github.com/songquanpeng/one-api/relay/meta"
+	relaymodel "github.com/songquanpeng/one-api/relay/model"
+	"io"
+	"net/http"
+)
+
+func RerankHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatusCode {
+	ctx := c.Request.Context()
+	meta := meta.GetByContext(c)
+	rerankRequest, err := getRerankRequest(c, meta.Mode)
+	if err != nil {
+		logger.Errorf(ctx, "getRerankRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "invalid_rerank_request", http.StatusBadRequest)
+	}
+
+	// Map model name
+	var isModelMapped bool
+	meta.OriginModelName = rerankRequest.Model
+	rerankRequest.Model, isModelMapped = getMappedModelName(rerankRequest.Model, meta.ModelMapping)
+	meta.ActualModelName = rerankRequest.Model
+
+	var requestBody io.Reader
+	if isModelMapped {
+		jsonStr, err := json.Marshal(rerankRequest)
+		if err != nil {
+			return openai.ErrorWrapper(err, "marshal_rerank_request_failed", http.StatusInternalServerError)
+		}
+		requestBody = bytes.NewBuffer(jsonStr)
+	} else {
+		requestBody = c.Request.Body
+	}
+
+	adaptor := relay.GetAdaptor(meta.APIType)
+	if adaptor == nil {
+		return openai.ErrorWrapper(fmt.Errorf("invalid api type: %d", meta.APIType), "invalid_api_type", http.StatusBadRequest)
+	}
+
+	resp, err := adaptor.DoRequest(c, meta, requestBody)
+	if err != nil {
+		logger.Errorf(ctx, "DoRequest failed: %s", err.Error())
+		return openai.ErrorWrapper(err, "do_request_failed", http.StatusInternalServerError)
+	}
+
+	// do response
+	_, respErr := adaptor.DoResponse(c, resp, meta)
+	if respErr != nil {
+		logger.Errorf(ctx, "respErr is not nil: %+v", respErr)
+		return respErr
+	}
+
+	return nil
+}