From d8d49df35e8168f78e2af9050afa247c91aa45ce Mon Sep 17 00:00:00 2001 From: Panda Dev <56657208+pandadev66@users.noreply.github.com> Date: Fri, 8 May 2026 07:08:22 +0200 Subject: [PATCH] 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 --- conf/models/gitee.json | 7 +++ internal/entity/models/gitee.go | 99 ++++++++++++++++++++++++++++++++- 2 files changed, 105 insertions(+), 1 deletion(-) diff --git a/conf/models/gitee.json b/conf/models/gitee.json index 9ac683bc93..630106592f 100644 --- a/conf/models/gitee.json +++ b/conf/models/gitee.json @@ -32,6 +32,13 @@ "model_types": [ "chat" ] + }, + { + "name": "BAAI/bge-reranker-v2-m3", + "max_tokens": 8192, + "model_types": [ + "rerank" + ] } ] } \ No newline at end of file diff --git a/internal/entity/models/gitee.go b/internal/entity/models/gitee.go index 7f7050bfb9..85d4635611 100644 --- a/internal/entity/models/gitee.go +++ b/internal/entity/models/gitee.go @@ -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) {