74 lines
1.6 KiB
Go
74 lines
1.6 KiB
Go
package kb
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
// Reranker 语义重排客户端 (Cross-Encoder)
|
|
type Reranker struct {
|
|
Endpoint string
|
|
APIKey string
|
|
Model string
|
|
HTTP *http.Client
|
|
}
|
|
|
|
// NewReranker 创建 Reranker 实例
|
|
func NewReranker(endpoint, key, model string) *Reranker {
|
|
return &Reranker{
|
|
Endpoint: normalizeAPIEndpoint(endpoint, "/rerank"),
|
|
APIKey: key,
|
|
Model: model,
|
|
HTTP: &http.Client{Timeout: 30 * time.Second},
|
|
}
|
|
}
|
|
|
|
// Rerank 对候选文档列表依据查询进行相关性重排打分
|
|
// 返回与 docs 等长的一组打分，下标一一对应
|
|
func (r *Reranker) Rerank(ctx context.Context, query string, docs []string) ([]float64, error) {
|
|
if r == nil || r.Endpoint == "" {
|
|
return nil, fmt.Errorf("reranker not configured")
|
|
}
|
|
if len(docs) == 0 {
|
|
return []float64{}, nil
|
|
}
|
|
|
|
body := map[string]any{
|
|
"model": r.Model,
|
|
"query": query,
|
|
"documents": docs,
|
|
}
|
|
|
|
raw, err := postJSON(ctx, r.HTTP, r.Endpoint, r.APIKey, body)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var resp struct {
|
|
Results []struct {
|
|
Index int `json:"index"`
|
|
RelevanceScore float64 `json:"relevance_score"`
|
|
Score float64 `json:"score"`
|
|
} `json:"results"`
|
|
}
|
|
|
|
if err := json.Unmarshal(raw, &resp); err != nil {
|
|
return nil, fmt.Errorf("decode rerank: %s", formatHTTPBodyError(0, raw))
|
|
}
|
|
|
|
out := make([]float64, len(docs))
|
|
for _, it := range resp.Results {
|
|
s := it.RelevanceScore
|
|
if s == 0 {
|
|
s = it.Score
|
|
}
|
|
if it.Index >= 0 && it.Index < len(out) {
|
|
out[it.Index] = s
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|