oadmin/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Update RerankerInfo structure and modify GenRelayInfoRerank function to accept RerankRequest

Browse Source
This commit is contained in:
1808837298@qq.com
2025-03-17 16:44:53 +08:00
parent b5aa3c129b
commit 6fef5aaf22
4 changed files with 25 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
dto/rerank.go
View File

@@ -5,11 +5,18 @@ type RerankRequest struct {
Query string `json:"query"` Query string `json:"query"`
Model string `json:"model"` Model string `json:"model"`
TopN int `json:"top_n"` 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"` MaxChunkPerDoc int `json:"max_chunk_per_doc,omitempty"`
OverLapTokens int `json:"overlap_tokens,omitempty"` OverLapTokens int `json:"overlap_tokens,omitempty"`
} }
func (r *RerankRequest) GetReturnDocuments() bool {
if r.ReturnDocuments == nil {
return false
}
return *r.ReturnDocuments
}
type RerankResponseResult struct { type RerankResponseResult struct {
Document any `json:"document,omitempty"` Document any `json:"document,omitempty"`
Index int `json:"index"` Index int `json:"index"`

8
relay/common/relay_info.go
View File

@@ -34,7 +34,8 @@ const (
) )
type RerankerInfo struct { type RerankerInfo struct {
Documents []any Documents []any
ReturnDocuments bool
} }
type RelayInfo struct { type RelayInfo struct {
@@ -116,11 +117,12 @@ func GenRelayInfoClaude(c *gin.Context) *RelayInfo {
return info return info
} }
func GenRelayInfoRerank(c *gin.Context, documents []any) *RelayInfo { func GenRelayInfoRerank(c *gin.Context, req *dto.RerankRequest) *RelayInfo {
info := GenRelayInfo(c) info := GenRelayInfo(c)
info.RelayMode = relayconstant.RelayModeRerank info.RelayMode = relayconstant.RelayModeRerank
info.RerankerInfo = &RerankerInfo{ info.RerankerInfo = &RerankerInfo{
Documents: documents, Documents: req.Documents,
ReturnDocuments: req.GetReturnDocuments(),
} }
return info return info
} }

21
relay/common_handler/rerank.go
View File

@@ -32,19 +32,20 @@ func RerankHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Respo
} }
jinaRespResults := make([]dto.RerankResponseResult, len(xinRerankResponse.Results)) jinaRespResults := make([]dto.RerankResponseResult, len(xinRerankResponse.Results))
for i, result := range xinRerankResponse.Results { for i, result := range xinRerankResponse.Results {
var document any respResult := dto.RerankResponseResult{
if result.Document == "" {
document = info.Documents[result.Index]
} else {
document = result.Document
}
jinaRespResults[i] = dto.RerankResponseResult{
Index: result.Index, Index: result.Index,
RelevanceScore: result.RelevanceScore, 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{ jinaResp = dto.RerankResponse{
Results: jinaRespResults, Results: jinaRespResults,

2
relay/relay_rerank.go
View File

@@ -33,7 +33,7 @@ func RerankHelper(c *gin.Context, relayMode int) (openaiErr *dto.OpenAIErrorWith
return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest) return service.OpenAIErrorWrapperLocal(err, "invalid_text_request", http.StatusBadRequest)
} }
relayInfo := relaycommon.GenRelayInfoRerank(c, rerankRequest.Documents) relayInfo := relaycommon.GenRelayInfoRerank(c, rerankRequest)
if rerankRequest.Query == "" { if rerankRequest.Query == "" {
return service.OpenAIErrorWrapperLocal(fmt.Errorf("query is empty"), "invalid_query", http.StatusBadRequest) return service.OpenAIErrorWrapperLocal(fmt.Errorf("query is empty"), "invalid_query", http.StatusBadRequest)