mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-06-29 23:41:12 +08:00
Go: implement Balance in DeepSeek driver (#14632)
Closes #14631 ### What problem does this PR solve? The DeepSeek Go driver shipped with a stub \`Balance\` method that returned \`no such method\`, even though DeepSeek exposes a public \`GET /user/balance\` endpoint that works with the same Bearer token used for chat. So the "Balance" panel in the model provider UI always shows an error for DeepSeek tenants, while it already works for Moonshot and Gitee. This PR fills the gap. ### What this PR includes - \`conf/models/deepseek.json\`: add \`\"balance\": \"user/balance\"\` under \`url_suffix\` so the driver can build the URL from config the same way the other endpoints do. - \`internal/entity/models/deepseek.go\`: replace the \`Balance\` stub with a real implementation. Adds a small local response type \`deepseekBalanceResponse\` that matches the upstream shape. No factory change. No interface change. ### How the driver works - Validate \`apiConfig\` and the API key, resolve the region (with a \`default\` fallback), and build the URL from \`BaseURL[region] + URLSuffix.Balance\`. - GET the URL with \`Authorization: Bearer <api_key>\`. - Parse the upstream response: \`\`\`json { \"is_available\": true, \"balance_infos\": [ {\"currency\": \"USD\", \"total_balance\": \"10.00\", ...}, {\"currency\": \"CNY\", \"total_balance\": \"70.00\", ...} ] } \`\`\` \`total_balance\` is a string in the upstream API, so the driver parses it with \`strconv.ParseFloat\`. - Return the first balance entry as \`{\"balance\": <float>, \"currency\": <string>}\`, the same shape the Moonshot driver returns. The UI can render it with no provider-specific code. ### Edge cases - Missing or empty API key returns a clear local error before any HTTP call. - Empty \`balance_infos\` returns a clear \"no balance info in response\" error rather than a zero-value silent success. - Non-numeric \`total_balance\` returns a clear parse error. - Non-200 responses propagate the upstream status line and body so the user can see why the call failed. ### 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 \`DeepSeekModel\` still matches the \`ModelDriver\` interface. - Pattern parity with the existing Moonshot and Gitee Balance implementations.
This commit is contained in:
@@ -5,7 +5,8 @@
|
||||
},
|
||||
"url_suffix": {
|
||||
"chat": "chat/completions",
|
||||
"models": "models"
|
||||
"models": "models",
|
||||
"balance": "user/balance"
|
||||
},
|
||||
"class": "deepseek",
|
||||
"models": [
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in New Issue
Block a user