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:
Panda Dev
2026-05-07 11:56:30 +02:00
committed by GitHub
parent 59bb184e63
commit bb10b83e61

View File

@@ -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
}