mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-11 14:15:40 +08:00
Go: implement Rerank in ZhipuAI driver (#14608)
### What problem does this PR solve?
The ZhipuAI Go driver had a stub Rerank method that returned "not
implemented", even though conf/models/zhipu-ai.json already ships
glm-rerank as a rerank model and the rerank URL suffix is already wired
in url_suffix:
```json
"url_suffix": {
...
"rerank": "rerank"
},
"models": [
{"name": "glm-rerank", "model_types": ["rerank"]},
...
]
```
So the config was ready but the driver was not. A tenant who picked
glm-rerank in the Go layer could not actually run a rerank call. This PR
fills the gap so the listed model works end to end.
### What this PR includes
- `internal/entity/models/zhipu-ai.go`: real implementation of
`ZhipuAIModel.Rerank`, plus two small local types (`zhipuRerankRequest`,
`zhipuRerankResponse`) that mirror the standard OpenAI-compatible rerank
shape used by SiliconFlow.
No factory change. No JSON change. No interface change.
### How the driver works
- POST to `${BaseURL}/${URLSuffix.Rerank}` (resolves to
`https://open.bigmodel.cn/api/paas/v4/rerank` with the default config),
reusing the existing httpClient on the driver.
- Validate apiConfig and the API key, validate the model name, and
resolve the region. Return a clear local error before any HTTP call when
something is missing.
- Send `{model, query, documents, top_n, return_documents: false}` in
the body, the same shape the SiliconFlow driver already uses.
- Walk `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 results in a different order.
- Empty `texts` input returns an empty `[]float64` with no HTTP call.
- 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 (the
go.mod minimum) returns exit 0.
- The full method set on `ZhipuAIModel` still matches the `ModelDriver`
interface (NewInstance, Name, ChatWithMessages, ChatStreamlyWithSender,
Encode, ListModels, Balance, CheckConnection, Rerank).
- Pattern parity with the existing SiliconFlow Rerank implementation
(`internal/entity/models/siliconflow.go`).
Closes #14607
This commit is contained in:
@@ -489,7 +489,101 @@ func (z *ZhipuAIModel) CheckConnection(apiConfig *APIConfig) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Rerank calculates similarity scores between query and texts
|
||||
func (z *ZhipuAIModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
|
||||
return nil, fmt.Errorf("%s, Rerank not implemented", z.Name())
|
||||
// zhipuRerankRequest is the request body for the ZhipuAI rerank
|
||||
// endpoint. The shape matches the standard OpenAI-compatible rerank
|
||||
// API also used by SiliconFlow.
|
||||
type zhipuRerankRequest struct {
|
||||
Model string `json:"model"`
|
||||
Query string `json:"query"`
|
||||
Documents []string `json:"documents"`
|
||||
TopN int `json:"top_n"`
|
||||
ReturnDocuments bool `json:"return_documents"`
|
||||
}
|
||||
|
||||
// zhipuRerankResponse is the response shape for the ZhipuAI rerank
|
||||
// endpoint.
|
||||
type zhipuRerankResponse struct {
|
||||
Results []struct {
|
||||
Index int `json:"index"`
|
||||
RelevanceScore float64 `json:"relevance_score"`
|
||||
} `json:"results"`
|
||||
}
|
||||
|
||||
// Rerank calculates similarity scores between query and texts using
|
||||
// the ZhipuAI /rerank endpoint (e.g. glm-rerank). The result is one
|
||||
// score per input text, in the same order the texts were given.
|
||||
func (z *ZhipuAIModel) Rerank(modelName *string, query string, texts []string, apiConfig *APIConfig) ([]float64, error) {
|
||||
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 {
|
||||
region = *apiConfig.Region
|
||||
}
|
||||
|
||||
baseURL, ok := z.BaseURL[region]
|
||||
if !ok || baseURL == "" {
|
||||
return nil, fmt.Errorf("zhipu-ai: no base URL configured for region %q", region)
|
||||
}
|
||||
|
||||
url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), z.URLSuffix.Rerank)
|
||||
|
||||
reqBody := zhipuRerankRequest{
|
||||
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)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("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("ZhipuAI rerank API error: %s, body: %s", resp.Status, string(body))
|
||||
}
|
||||
|
||||
var rerankResp zhipuRerankResponse
|
||||
if err = json.Unmarshal(body, &rerankResp); err != nil {
|
||||
return nil, fmt.Errorf("failed to parse response: %w", err)
|
||||
}
|
||||
|
||||
scores := make([]float64, len(texts))
|
||||
for _, r := range rerankResp.Results {
|
||||
if r.Index >= 0 && r.Index < len(texts) {
|
||||
scores[r.Index] = r.RelevanceScore
|
||||
}
|
||||
}
|
||||
|
||||
return scores, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user