|
| 1 | +package agent |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "encoding/json" |
| 6 | + "fmt" |
| 7 | + "io" |
| 8 | + "net/http" |
| 9 | + "strconv" |
| 10 | + "strings" |
| 11 | + "time" |
| 12 | +) |
| 13 | + |
| 14 | +// BalanceResult 是一次账户余额查询的结果。Display 已是可直接渲染的金额串(含币种符号,如 "¥110.00")。 |
| 15 | +type BalanceResult struct { |
| 16 | + Display string // 可渲染金额串;Supported=false 时为空 |
| 17 | + Supported bool // 该供应商是否支持凭模型 API Key 查余额 |
| 18 | +} |
| 19 | + |
| 20 | +// ProbeBalance 用模型 API Key 查询账户剩余金额。 |
| 21 | +// |
| 22 | +// 现状:四家预置供应商里只有 DeepSeek 和 Kimi(Moonshot)提供可凭模型 Key 直接调用的余额接口。 |
| 23 | +// qwen 需阿里云 AccessKey/SecretKey 走 BSS OpenAPI、mimo 无任何余额接口 —— 这两家(及未知 custom 端点) |
| 24 | +// 返回 Supported=false,调用方据此显示 "-"。 |
| 25 | +// |
| 26 | +// 返回 (result, err): |
| 27 | +// - err != nil:网络 / 5xx 等瞬时错误 → 调用方不更新,下次再探; |
| 28 | +// - err == nil 且 Supported=false:该供应商不支持查询; |
| 29 | +// - err == nil 且 Supported=true:Display 为可渲染金额串。 |
| 30 | +func ProbeBalance(ctx context.Context, entry ModelEntry) (BalanceResult, error) { |
| 31 | + if entry.BaseURL == "" || entry.APIKey == "" { |
| 32 | + return BalanceResult{Supported: false}, nil |
| 33 | + } |
| 34 | + host := balanceHostOf(entry.BaseURL) |
| 35 | + switch { |
| 36 | + case strings.Contains(host, "deepseek"): |
| 37 | + return probeDeepSeekBalance(ctx, entry) |
| 38 | + case strings.Contains(host, "moonshot"): |
| 39 | + return probeKimiBalance(ctx, entry) |
| 40 | + default: |
| 41 | + // qwen / mimo / 未知 custom 端点:无可凭 Key 调用的余额接口。 |
| 42 | + return BalanceResult{Supported: false}, nil |
| 43 | + } |
| 44 | +} |
| 45 | + |
| 46 | +// probeDeepSeekBalance 调 GET {base_url}/user/balance(DeepSeek base_url 不含 /v1,拼出 /user/balance)。 |
| 47 | +func probeDeepSeekBalance(ctx context.Context, entry ModelEntry) (BalanceResult, error) { |
| 48 | + body, err := httpGetJSON(ctx, entry.BaseURL+"/user/balance", entry.APIKey) |
| 49 | + if err != nil { |
| 50 | + return BalanceResult{}, err |
| 51 | + } |
| 52 | + var r struct { |
| 53 | + IsAvailable bool `json:"is_available"` |
| 54 | + BalanceInfos []struct { |
| 55 | + Currency string `json:"currency"` |
| 56 | + TotalBalance string `json:"total_balance"` |
| 57 | + } `json:"balance_infos"` |
| 58 | + } |
| 59 | + if err := json.Unmarshal(body, &r); err != nil { |
| 60 | + return BalanceResult{}, err |
| 61 | + } |
| 62 | + if len(r.BalanceInfos) == 0 { |
| 63 | + return BalanceResult{Supported: true, Display: "—"}, nil |
| 64 | + } |
| 65 | + // DeepSeek 会按币种返回多条(常见 USD 0.00 在前、CNY 有钱在后)。挑金额最大的那条展示, |
| 66 | + // 全为 0 则退回第一条 —— 避免误显 "$0.00" 而把真实的 ¥24.86 漏掉。 |
| 67 | + best := r.BalanceInfos[0] |
| 68 | + bestAmt := -1.0 |
| 69 | + for _, b := range r.BalanceInfos { |
| 70 | + amt, err := strconv.ParseFloat(b.TotalBalance, 64) |
| 71 | + if err != nil { |
| 72 | + continue |
| 73 | + } |
| 74 | + if amt > bestAmt { |
| 75 | + bestAmt, best = amt, b |
| 76 | + } |
| 77 | + } |
| 78 | + return BalanceResult{Supported: true, Display: currencySymbol(best.Currency) + best.TotalBalance}, nil |
| 79 | +} |
| 80 | + |
| 81 | +// probeKimiBalance 调 GET {base_url}/users/me/balance(Kimi base_url 带 /v1,拼出 /v1/users/me/balance)。 |
| 82 | +func probeKimiBalance(ctx context.Context, entry ModelEntry) (BalanceResult, error) { |
| 83 | + body, err := httpGetJSON(ctx, entry.BaseURL+"/users/me/balance", entry.APIKey) |
| 84 | + if err != nil { |
| 85 | + return BalanceResult{}, err |
| 86 | + } |
| 87 | + var r struct { |
| 88 | + Status bool `json:"status"` |
| 89 | + Data struct { |
| 90 | + AvailableBalance float64 `json:"available_balance"` |
| 91 | + } `json:"data"` |
| 92 | + } |
| 93 | + if err := json.Unmarshal(body, &r); err != nil { |
| 94 | + return BalanceResult{}, err |
| 95 | + } |
| 96 | + // Kimi 余额均为人民币,且 available_balance = 现金 + 代金券,即"还能花多少"。 |
| 97 | + return BalanceResult{Supported: true, Display: fmt.Sprintf("¥%.2f", r.Data.AvailableBalance)}, nil |
| 98 | +} |
| 99 | + |
| 100 | +// httpGetJSON 发一个带 Bearer 鉴权的 GET,返回响应体(非 200 视为错误,调用方按瞬时错误处理)。 |
| 101 | +func httpGetJSON(ctx context.Context, url, apiKey string) ([]byte, error) { |
| 102 | + ctx, cancel := context.WithTimeout(ctx, 15*time.Second) |
| 103 | + defer cancel() |
| 104 | + req, err := http.NewRequestWithContext(ctx, "GET", url, nil) |
| 105 | + if err != nil { |
| 106 | + return nil, err |
| 107 | + } |
| 108 | + req.Header.Set("Accept", "application/json") |
| 109 | + req.Header.Set("Authorization", "Bearer "+apiKey) |
| 110 | + resp, err := http.DefaultClient.Do(req) |
| 111 | + if err != nil { |
| 112 | + return nil, err |
| 113 | + } |
| 114 | + defer resp.Body.Close() |
| 115 | + body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024)) |
| 116 | + if resp.StatusCode != 200 { |
| 117 | + return nil, fmt.Errorf("balance: http %d", resp.StatusCode) |
| 118 | + } |
| 119 | + return body, nil |
| 120 | +} |
| 121 | + |
| 122 | +// currencySymbol 把币种代码转成符号;未知币种回退成 "CODE "(带尾空格,数字前留隔)。 |
| 123 | +func currencySymbol(cur string) string { |
| 124 | + switch strings.ToUpper(cur) { |
| 125 | + case "CNY", "RMB": |
| 126 | + return "¥" |
| 127 | + case "USD": |
| 128 | + return "$" |
| 129 | + default: |
| 130 | + if cur == "" { |
| 131 | + return "" |
| 132 | + } |
| 133 | + return cur + " " |
| 134 | + } |
| 135 | +} |
| 136 | + |
| 137 | +// balanceHostOf 从 base_url 抽出 host(去 scheme / path),用于判定供应商。 |
| 138 | +func balanceHostOf(rawURL string) string { |
| 139 | + h := rawURL |
| 140 | + if i := strings.Index(h, "://"); i >= 0 { |
| 141 | + h = h[i+3:] |
| 142 | + } |
| 143 | + if i := strings.IndexAny(h, "/?"); i >= 0 { |
| 144 | + h = h[:i] |
| 145 | + } |
| 146 | + return h |
| 147 | +} |
0 commit comments