117 lines
2.7 KiB
Go
117 lines
2.7 KiB
Go
|
package eino
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bytes"
|
|||
|
|
"context"
|
|||
|
|
"encoding/json"
|
|||
|
|
"fmt"
|
|||
|
|
"io"
|
|||
|
|
"net/http"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"github.com/cloudwego/eino/schema"
|
|||
|
|
"github.com/gogf/gf/v2/frame/g"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// DashScopeReranker 通义百炼 Rerank 精排（Cross-Encoder）
|
|||
|
|
type DashScopeReranker struct {
|
|||
|
|
httpClient *http.Client
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func NewDashScopeReranker() *DashScopeReranker {
|
|||
|
|
return &DashScopeReranker{
|
|||
|
|
httpClient: &http.Client{
|
|||
|
|
Timeout: 10 * time.Second,
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Rerank 对文档进行精排（Cross-Encoder 核心）
|
|||
|
|
func (d *DashScopeReranker) Rerank(ctx context.Context, query string, docs []*schema.Document) ([]*schema.Document, error) {
|
|||
|
|
if len(docs) == 0 {
|
|||
|
|
return docs, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 官方必过 URL
|
|||
|
|
url := "https://dashscope.aliyuncs.com/api/v1/services/rerank/text-rerank/text-rerank"
|
|||
|
|
apiKey := g.Cfg().MustGet(ctx, "eino.rerank.apiKey").String()
|
|||
|
|
model := g.Cfg().MustGet(ctx, "eino.rerank.model").String()
|
|||
|
|
|
|||
|
|
documents := make([]string, len(docs))
|
|||
|
|
for i, doc := range docs {
|
|||
|
|
documents[i] = doc.Content
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
reqBody := map[string]any{
|
|||
|
|
"model": model,
|
|||
|
|
"input": map[string]any{
|
|||
|
|
"query": query,
|
|||
|
|
"documents": documents,
|
|||
|
|
},
|
|||
|
|
"parameters": map[string]any{
|
|||
|
|
"top_n": len(docs),
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
bs, _ := json.Marshal(reqBody)
|
|||
|
|
|
|||
|
|
req, _ := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(bs))
|
|||
|
|
req.Header.Set("Authorization", "Bearer "+apiKey)
|
|||
|
|
req.Header.Set("Content-Type", "application/json")
|
|||
|
|
|
|||
|
|
resp, err := d.httpClient.Do(req)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
defer resp.Body.Close()
|
|||
|
|
|
|||
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|||
|
|
body, _ := io.ReadAll(resp.Body)
|
|||
|
|
return nil, fmt.Errorf("rerank api error: status=%d, body=%s", resp.StatusCode, string(body))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 解析结果
|
|||
|
|
var result struct {
|
|||
|
|
Output struct {
|
|||
|
|
Results []struct {
|
|||
|
|
Index int `json:"index"`
|
|||
|
|
RelevanceScore float64 `json:"relevance_score"`
|
|||
|
|
} `json:"results"`
|
|||
|
|
} `json:"output"`
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 按分数排序
|
|||
|
|
type scoredDoc struct {
|
|||
|
|
doc *schema.Document
|
|||
|
|
score float64
|
|||
|
|
}
|
|||
|
|
scored := make([]scoredDoc, len(docs))
|
|||
|
|
for i, doc := range docs {
|
|||
|
|
scored[i] = scoredDoc{doc: doc, score: 0}
|
|||
|
|
}
|
|||
|
|
for _, res := range result.Output.Results {
|
|||
|
|
scored[res.Index].score = res.RelevanceScore
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 分数从高到低排序
|
|||
|
|
for i := 0; i < len(scored); i++ {
|
|||
|
|
for j := i + 1; j < len(scored); j++ {
|
|||
|
|
if scored[j].score > scored[i].score {
|
|||
|
|
scored[i], scored[j] = scored[j], scored[i]
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 输出最终排好的文档
|
|||
|
|
ranked := make([]*schema.Document, 0, len(scored))
|
|||
|
|
for _, s := range scored {
|
|||
|
|
s.doc.MetaData["rerank_score"] = s.score
|
|||
|
|
ranked = append(ranked, s.doc)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
return ranked, nil
|
|||
|
|
}
|