From 6fef5aaf22114bb99ab0bcf61b74577c3cb263d8 Mon Sep 17 00:00:00 2001
From: "1808837298@qq.com" <1808837298@qq.com>
Date: Mon, 17 Mar 2025 16:44:53 +0800
Subject: [PATCH] feat: Update RerankerInfo structure and modify
 GenRelayInfoRerank function to accept RerankRequest

---
 dto/rerank.go                  |  9 ++++++++-
 relay/common/relay_info.go     |  8 +++++---
 relay/common_handler/rerank.go | 21 +++++++++++----------
 relay/relay_rerank.go          |  2 +-
 4 files changed, 25 insertions(+), 15 deletions(-)

diff --git a/dto/rerank.go b/dto/rerank.go
index 38aca907..21f6437c 100644
--- a/dto/rerank.go
+++ b/dto/rerank.go
@@ -5,11 +5,18 @@ type RerankRequest struct {
 	Query           string `json:"query"`
 	Model           string `json:"model"`
 	TopN            int    `json:"top_n"`
-	ReturnDocuments bool   `json:"return_documents,omitempty"`
+	ReturnDocuments *bool  `json:"return_documents,omitempty"`
 	MaxChunkPerDoc  int    `json:"max_chunk_per_doc,omitempty"`
 	OverLapTokens   int    `json:"overlap_tokens,omitempty"`
 }
 
+func (r *RerankRequest) GetReturnDocuments() bool {
+	if r.ReturnDocuments == nil {
+		return false
+	}
+	return *r.ReturnDocuments
+}
+
 type RerankResponseResult struct {
 	Document       any     `json:"document,omitempty"`
 	Index          int     `json:"index"`
diff --git a/relay/common/relay_info.go b/relay/common/relay_info.go
index 6b96419a..059c8284 100644
--- a/relay/common/relay_info.go
+++ b/relay/common/relay_info.go
@@ -34,7 +34,8 @@ const (
 )
 
 type RerankerInfo struct {
-	Documents []any
+	Documents       []any
+	ReturnDocuments bool
 }
 
 type RelayInfo struct {
@@ -116,11 +117,12 @@ func GenRelayInfoClaude(c *gin.Context) *RelayInfo {
 	return info
 }
 
-func GenRelayInfoRerank(c *gin.Context, documents []any) *RelayInfo {
+func GenRelayInfoRerank(c *gin.Context, req *dto.RerankRequest) *RelayInfo {
 	info := GenRelayInfo(c)
 	info.RelayMode = relayconstant.RelayModeRerank
 	info.RerankerInfo = &RerankerInfo{
-		Documents: documents,
+		Documents:       req.Documents,
+		ReturnDocuments: req.GetReturnDocuments(),
 	}
 	return info
 }
diff --git a/relay/common_handler/rerank.go b/relay/common_handler/rerank.go
index 05aaa8ae..496278b5 100644
--- a/relay/common_handler/rerank.go
+++ b/relay/common_handler/rerank.go
@@ -32,19 +32,20 @@ func RerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
 		}
 		jinaRespResults := make([]dto.RerankResponseResult, len(xinRerankResponse.Results))
 		for i, result := range xinRerankResponse.Results {
-			var document any
-			if result.Document == "" {
-				document = info.Documents[result.Index]
-			} else {
-				document = result.Document
-			}
-			jinaRespResults[i] = dto.RerankResponseResult{
+			respResult := dto.RerankResponseResult{
 				Index:          result.Index,
 				RelevanceScore: result.RelevanceScore,
-				Document: dto.RerankDocument{
-					Text: document,
-				},
 			}
+			if info.ReturnDocuments {
+				var document any
+				if result.Document == "" {
+					document = info.Documents[result.Index]
+				} else {
+					document = result.Document
+				}
+				respResult.Document = document
+			}
+			jinaRespResults[i] = respResult
 		}
 		jinaResp = dto.RerankResponse{
 			Results: jinaRespResults,
diff --git a/relay/relay_rerank.go b/relay/relay_rerank.go
index 69ab2247..6ca98de7 100644
--- a/relay/relay_rerank.go
+++ b/relay/relay_rerank.go
@@ -33,7 +33,7 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith
 		return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest)
 	}
 
-	relayInfo := relaycommon.GenRelayInfoRerank(c, rerankRequest.Documents)
+	relayInfo := relaycommon.GenRelayInfoRerank(c, rerankRequest)
 
 	if rerankRequest.Query == "" {
 		return service.OpenAIErrorWrapperLocal(fmt.Errorf("query is empty"), "invalid_query", http.StatusBadRequest)