diff --git a/internal/entity/models/zhipu-ai.go b/internal/entity/models/zhipu-ai.go index 8785b26fb1..8da7b0f5e0 100644 --- a/internal/entity/models/zhipu-ai.go +++ b/internal/entity/models/zhipu-ai.go @@ -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 }