diff --git a/conf/models/deepseek.json b/conf/models/deepseek.json index 5fdce2ac9f..146e11862a 100644 --- a/conf/models/deepseek.json +++ b/conf/models/deepseek.json @@ -5,7 +5,8 @@ }, "url_suffix": { "chat": "chat/completions", - "models": "models" + "models": "models", + "balance": "user/balance" }, "class": "deepseek", "models": [ diff --git a/internal/entity/models/deepseek.go b/internal/entity/models/deepseek.go index b74ef0db46..f1fd3116ac 100644 --- a/internal/entity/models/deepseek.go +++ b/internal/entity/models/deepseek.go @@ -24,6 +24,7 @@ import ( "io" "net/http" "ragflow/internal/common" + "strconv" "strings" "time" ) @@ -483,8 +484,92 @@ func (z *DeepSeekModel) ListModels(apiConfig *APIConfig) ([]string, error) { return models, nil } +// deepseekBalanceResponse is the shape returned by +// GET /user/balance. The balance fields are strings in the +// upstream API, so we parse them on our side. +type deepseekBalanceResponse struct { + IsAvailable bool `json:"is_available"` + BalanceInfos []struct { + Currency string `json:"currency"` + TotalBalance string `json:"total_balance"` + GrantedBalance string `json:"granted_balance"` + ToppedUpBalance string `json:"topped_up_balance"` + } `json:"balance_infos"` +} + +// Balance returns the user's available balance on DeepSeek by +// calling GET /user/balance with the configured Bearer token. +// The result map matches the shape used by the Moonshot driver, +// so the UI can render it without provider-specific code. func (z *DeepSeekModel) 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 + } + + // Look up the base URL for the requested region. If the region was + // supplied but is not configured (or is empty), fall back to the + // "default" region instead of erroring out, so a stray region value + // does not break an otherwise valid request. + baseURL := z.BaseURL["default"] + if region != "default" { + if regional, ok := z.BaseURL[region]; ok && regional != "" { + baseURL = regional + } + } + if baseURL == "" { + return nil, fmt.Errorf("deepseek: 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("DeepSeek balance API error: %s, body: %s", resp.Status, string(body)) + } + + var parsed deepseekBalanceResponse + if err = json.Unmarshal(body, &parsed); err != nil { + return nil, fmt.Errorf("failed to parse response: %w", err) + } + + if len(parsed.BalanceInfos) == 0 { + return nil, fmt.Errorf("no balance info in response") + } + + // Pick the first balance entry, the same way the Moonshot + // driver returns a single {balance, currency} pair to the UI. + first := parsed.BalanceInfos[0] + total, err := strconv.ParseFloat(first.TotalBalance, 64) + if err != nil { + return nil, fmt.Errorf("invalid total_balance %q: %w", first.TotalBalance, err) + } + + return map[string]interface{}{ + "balance": total, + "currency": first.Currency, + }, nil } func (z *DeepSeekModel) CheckConnection(apiConfig *APIConfig) error {