feat: Update RerankerInfo structure and modify GenRelayInfoRerank function to accept RerankRequest
This commit is contained in:
@@ -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"`
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user