diff --git a/conf/models/siliconflow.json b/conf/models/siliconflow.json index d9340365d0..4da3e0dcab 100644 --- a/conf/models/siliconflow.json +++ b/conf/models/siliconflow.json @@ -7,7 +7,8 @@ "chat": "chat/completions", "models": "models", "embedding": "embeddings", - "rerank": "rerank" + "rerank": "rerank", + "balance": "user/info" }, "models": [ { diff --git a/internal/entity/models/siliconflow.go b/internal/entity/models/siliconflow.go index c1a1db07ef..6c9e9ce4af 100644 --- a/internal/entity/models/siliconflow.go +++ b/internal/entity/models/siliconflow.go @@ -24,6 +24,7 @@ import ( "io" "net/http" "ragflow/internal/common" + "strconv" "strings" "time" ) @@ -528,8 +529,90 @@ func (z *SiliconflowModel) ListModels(apiConfig *APIConfig) ([]string, error) { return models, nil } +type siliconflowBalanceResponse struct { + Code int `json:"code"` + Status bool `json:"status"` + Message string `json:"message"` + Data struct { + Balance string `json:"balance"` + TotalBalance string `json:"totalBalance"` + } `json:"data"` +} + func (z *SiliconflowModel) Balance(apiConfig *APIConfig) (map[string]interface{}, error) { - return nil, fmt.Errorf("%s, no such method", z.Name()) + if apiConfig == nil || apiConfig.ApiKey == nil || *apiConfig.ApiKey == "" { + return nil, fmt.Errorf("api key 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("siliconflow: no base URL configured for default region") + } + + url := fmt.Sprintf("%s/%s", strings.TrimSuffix(baseURL, "/"), z.URLSuffix.Balance) + + req, err := http.NewRequest("GET", url, nil) + if err != nil { + return nil, fmt.Errorf("failed to create request: %w", err) + } + + 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("SiliconFlow balance API error: %s, body: %s", resp.Status, string(body)) + } + + var parsed siliconflowBalanceResponse + if err = json.Unmarshal(body, &parsed); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + if !parsed.Status { + msg := parsed.Message + if msg == "" { + msg = "unknown API error" + } + return nil, fmt.Errorf("SiliconFlow API error (code %d): %s", parsed.Code, msg) + } + + raw := parsed.Data.TotalBalance + if raw == "" { + raw = parsed.Data.Balance + } + if raw == "" { + return nil, fmt.Errorf("no balance info in response") + } + + total, err := strconv.ParseFloat(raw, 64) + if err != nil { + return nil, fmt.Errorf("invalid balance %q: %w", raw, err) + } + + return map[string]interface{}{ + "balance": total, + "currency": "CNY", + }, nil } func (z *SiliconflowModel) CheckConnection(apiConfig *APIConfig) error {