mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-08 08:07:21 +08:00
Go: implement Rerank in Gitee AI driver (#14656)
### What problem does this PR solve?
The Gitee AI Go driver shipped with a stub \`Rerank\` method that
returned \`Rerank not implemented\`, even though
\`conf/models/gitee.json\` already wires the rerank URL suffix at
\`\"rerank\": \"rerank\"\`. The same config did not list any
rerank model, so the picker had nothing to select.
So a Gitee tenant could not use BAAI/bge-reranker-v2-m3 as a reranker
through the Go layer today, even though the
infrastructure was one config entry and one method body away.
### What this PR includes
- \`conf/models/gitee.json\`: add \`BAAI/bge-reranker-v2-m3\` to the
\`models\` array.
- \`internal/entity/models/gitee.go\`: replace the \`Rerank\` stub with
a real implementation. Adds two small local types
that match the OpenAI-compatible \`/rerank\` shape already used by the
SiliconFlow and ZhipuAI drivers.
No factory change. No interface change.
### How the driver works
- Validate \`apiConfig\` and the API key, validate the model name,
resolve the region with a default fallback, build the
URL from \`BaseURL[region] + URLSuffix.Rerank\`.
- Use a per-call \`context.WithTimeout(30s)\` and
\`http.NewRequestWithContext\`, matching the pattern the
recently merged Aliyun Encode and the OpenAI driver already use.
- Send \`{model, query, documents, top_n, return_documents:false}\` in
the body.
- Parse \`results[*].relevance_score\` and copy each score into the
output slice indexed by \`results[*].index\`, so the
output order matches the input order even if the API returns items in a
different order.
- Empty input returns \`[]float64{}\` with no HTTP call.
- An out-of-range result index returns a clear error rather than
silently skipping the entry.
- Non-200 responses propagate the upstream status line and body.
### Type of change
- [x] New Feature (non-breaking change which adds functionality)
### How was this tested?
- \`go build ./internal/entity/models/...\` in a clean go 1.25 image
returns exit 0.
- The full method set on \`GiteeModel\` still matches the
\`ModelDriver\` interface.
- Pattern parity with the existing SiliconFlow Rerank and the recently
merged ZhipuAI Rerank (#14608).
Closes #14655
This commit is contained in:
@ -19,6 +19,7 @@ package models
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
@ -402,9 +403,105 @@ func (z *GiteeModel) Encode(modelName *string, texts []string, apiConfig *APICon
|
||||
return nil, fmt.Errorf("%s, no such method", z.Name())
|
||||
}
|
||||
|
||||
type giteeRerankRequest struct {
|
||||
Model string `json:"model"`
|
||||
Query string `json:"query"`
|
||||
Documents []string `json:"documents"`
|
||||
TopN int `json:"top_n"`
|
||||
ReturnDocuments bool `json:"return_documents"`
|
||||
}
|
||||
|
||||
type giteeRerankResponse struct {
|
||||
Results []struct {
|
||||
Index int `json:"index"`
|
||||
RelevanceScore float64 `json:"relevance_score"`
|
||||
} `json:"results"`
|
||||
}
|
||||
|
||||
// Rerank calculates similarity scores between query and texts
|
||||
func (z *GiteeModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
|
||||
return nil, fmt.Errorf("%s, Rerank not implemented", z.Name())
|
||||
if len(texts) == 0 {
|
||||
return []float64{}, nil
|
||||
}
|
||||
|
||||
if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" {
|
||||
return nil, fmt.Errorf("api key is required")
|
||||
}
|
||||
|
||||
if modelName == nil || *modelName == "" {
|
||||
return nil, fmt.Errorf("model name is required")
|
||||
}
|
||||
|
||||
region := "default"
|
||||
if apiConfig.Region != nil && *apiConfig.Region != "" {
|
||||
region = *apiConfig.Region
|
||||
}
|
||||
|
||||
baseURL := z.BaseURL["default"]
|
||||
if region != "default" {
|
||||
if regional, ok := z.BaseURL[region]; ok && regional != "" {
|
||||
baseURL = regional
|
||||
}
|
||||
}
|
||||
if baseURL == "" {
|
||||
return nil, fmt.Errorf("gitee: no base URL configured for default region")
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), z.URLSuffix.Rerank)
|
||||
|
||||
reqBody := giteeRerankRequest{
|
||||
Model: *modelName,
|
||||
Query: query,
|
||||
Documents: texts,
|
||||
TopN: len(texts),
|
||||
ReturnDocuments: false,
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to marshal request: %w", err)
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to create request: %w", err)
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
|
||||
|
||||
resp, err := z.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to send request: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to read response: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return nil, fmt.Errorf("Gitee rerank API error: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
var parsed giteeRerankResponse
|
||||
if err = json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
scores := make([]float64, len(texts))
|
||||
for _, r := range parsed.Results {
|
||||
if r.Index < 0 || r.Index >= len(texts) {
|
||||
return nil, fmt.Errorf("unexpected rerank index %d for %d inputs", r.Index, len(texts))
|
||||
}
|
||||
scores[r.Index] = r.RelevanceScore
|
||||
}
|
||||
|
||||
return scores, nil
|
||||
}
|
||||
|
||||
func (z *GiteeModel) ListModels(apiConfig *APIConfig) ([]string, error) {
|
||||
|
||||
Reference in New Issue
Block a user