Skip to content

Commit 3e1835a

Browse files
authored
Merge pull request #9 from DevilGenius/dev
feat: 添加图片计价功能和相关设置
2 parents d6a69af + 8d85668 commit 3e1835a

16 files changed

Lines changed: 369 additions & 27 deletions

File tree

backend/internal/billing/calculator.go

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,11 @@ type CalculateInput struct {
2929
// 用于计算 account_cost(账号实际消耗),写入 usage_log,仅供"账号计费"统计使用。
3030
// 与用户计费 (BillingRate) 完全独立,不影响 actual_cost / User.balance。
3131
AccountRate float64
32+
33+
// OutputBillingCostOverride 可覆盖 output_cost 在 actual_cost 管道里的计价结果。
34+
// 用于分组图片 1K/2K/4K 固定价:配置后 output 不再乘 BillingRate,
35+
// 未配置时保持 output_cost × BillingRate 的旧行为。
36+
OutputBillingCostOverride *float64
3237
}
3338

3439
// CalculateResult 计算结果
@@ -70,7 +75,11 @@ func (c *Calculator) Calculate(input CalculateInput) CalculateResult {
7075
accountRate = 1.0
7176
}
7277

73-
actualCost := totalCost * billingRate
78+
nonOutputCost := input.InputCost + input.CachedInputCost + input.CacheCreationCost
79+
actualCost := nonOutputCost*billingRate + input.OutputCost*billingRate
80+
if input.OutputBillingCostOverride != nil {
81+
actualCost = nonOutputCost*billingRate + *input.OutputBillingCostOverride
82+
}
7483

7584
billedCost := actualCost
7685
if input.SellRate > 0 {

backend/internal/billing/calculator_test.go

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,3 +158,31 @@ func TestCalculate_MarkupIndependentOfBillingRate(t *testing.T) {
158158
t.Fatalf("BilledCost should depend on SellRate but didn't change")
159159
}
160160
}
161+
162+
func TestCalculate_OutputBillingCostOverride(t *testing.T) {
163+
c := NewCalculator()
164+
outputOverride := 0.08
165+
res := c.Calculate(CalculateInput{
166+
InputCost: 0.10,
167+
OutputCost: 0.40,
168+
BillingRate: 0.50,
169+
OutputBillingCostOverride: &outputOverride,
170+
AccountRate: 1.25,
171+
})
172+
173+
if !almostEqual(res.TotalCost, 0.50) {
174+
t.Fatalf("TotalCost = %v, want 0.50", res.TotalCost)
175+
}
176+
if !almostEqual(res.ActualCost, 0.13) {
177+
t.Fatalf("ActualCost = %v, want 0.13", res.ActualCost)
178+
}
179+
if !almostEqual(res.BilledCost, res.ActualCost) {
180+
t.Fatalf("BilledCost = %v, want %v", res.BilledCost, res.ActualCost)
181+
}
182+
if !almostEqual(res.AccountCost, 0.625) {
183+
t.Fatalf("AccountCost = %v, want 0.625", res.AccountCost)
184+
}
185+
if !almostEqual(res.RateMultiplier, 0.50) {
186+
t.Fatalf("RateMultiplier = %v, want original billing rate 0.50", res.RateMultiplier)
187+
}
188+
}

backend/internal/plugin/host_service.go

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -853,14 +853,18 @@ func (h *HostService) recordHostForwardUsage(
853853
return
854854
}
855855

856-
calc := h.calculator.Calculate(billing.CalculateInput{
856+
calcInput := billing.CalculateInput{
857857
InputCost: usage.InputCost,
858858
OutputCost: usage.OutputCost,
859859
CachedInputCost: usage.CachedInputCost,
860860
CacheCreationCost: usage.CacheCreationCost,
861861
BillingRate: route.EffectiveRate,
862862
AccountRate: accFull.RateMultiplier,
863-
})
863+
}
864+
if override, ok := imageOutputBillingOverride(usage, route.GroupPluginSettings); ok {
865+
calcInput.OutputBillingCostOverride = &override
866+
}
867+
calc := h.calculator.Calculate(calcInput)
864868

865869
h.scheduler.AddWindowCost(ctx, accountID, calc.AccountCost)
866870

@@ -1136,7 +1140,7 @@ func hostForwardHeaders(req *pb.HostForwardRequest, route routing.Candidate) htt
11361140
}
11371141
for plugin, kv := range route.GroupPluginSettings {
11381142
for k, v := range kv {
1139-
if v == "" {
1143+
if v == "" || !shouldForwardPluginSetting(plugin, k) {
11401144
continue
11411145
}
11421146
headers.Set("X-Airgate-Plugin-"+canonicalHeaderToken(plugin)+"-"+canonicalHeaderToken(k), v)
Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package plugin
2+
3+
import (
4+
"math"
5+
"strconv"
6+
"strings"
7+
8+
sdk "github.com/DouDOU-start/airgate-sdk"
9+
)
10+
11+
const (
12+
openAIPluginSettingsKey = "openai"
13+
14+
imagePrice1KKey = "image_price_1k"
15+
imagePrice2KKey = "image_price_2k"
16+
imagePrice4KKey = "image_price_4k"
17+
)
18+
19+
func imageOutputBillingOverride(usage *sdk.Usage, settings map[string]map[string]string) (float64, bool) {
20+
if usage == nil || strings.TrimSpace(usage.ImageSize) == "" || usage.OutputCost <= 0 {
21+
return 0, false
22+
}
23+
tier, basePrice, ok := imageTierForSize(usage.ImageSize)
24+
if !ok || basePrice <= 0 {
25+
return 0, false
26+
}
27+
price, ok := imageTierPriceFromSettings(settings, tier)
28+
if !ok {
29+
return 0, false
30+
}
31+
imageCount := int(math.Round(usage.OutputCost / basePrice))
32+
if imageCount < 1 {
33+
imageCount = 1
34+
}
35+
return float64(imageCount) * price, true
36+
}
37+
38+
func shouldForwardPluginSetting(plugin, key string) bool {
39+
if !strings.EqualFold(plugin, openAIPluginSettingsKey) {
40+
return true
41+
}
42+
switch strings.ToLower(strings.TrimSpace(key)) {
43+
case imagePrice1KKey, imagePrice2KKey, imagePrice4KKey:
44+
return false
45+
default:
46+
return true
47+
}
48+
}
49+
50+
func imageTierPriceFromSettings(settings map[string]map[string]string, tier string) (float64, bool) {
51+
key := imageTierPriceKey(tier)
52+
if key == "" {
53+
return 0, false
54+
}
55+
for pluginName, kv := range settings {
56+
if !strings.EqualFold(pluginName, openAIPluginSettingsKey) {
57+
continue
58+
}
59+
for k, v := range kv {
60+
if !strings.EqualFold(k, key) {
61+
continue
62+
}
63+
raw := strings.TrimSpace(v)
64+
if raw == "" {
65+
return 0, false
66+
}
67+
price, err := strconv.ParseFloat(raw, 64)
68+
if err != nil || price < 0 {
69+
return 0, false
70+
}
71+
return price, true
72+
}
73+
}
74+
return 0, false
75+
}
76+
77+
func imageTierPriceKey(tier string) string {
78+
switch strings.ToLower(strings.TrimSpace(tier)) {
79+
case "1k":
80+
return imagePrice1KKey
81+
case "2k":
82+
return imagePrice2KKey
83+
case "4k":
84+
return imagePrice4KKey
85+
default:
86+
return ""
87+
}
88+
}
89+
90+
func imageTierForSize(size string) (tier string, basePrice float64, ok bool) {
91+
width, height, ok := parseImageSizeForBilling(size)
92+
if !ok {
93+
return "", 0, false
94+
}
95+
longest := width
96+
if height > longest {
97+
longest = height
98+
}
99+
switch {
100+
case longest <= 1536:
101+
return "1k", 0.10, true
102+
case longest <= 2048:
103+
return "2k", 0.20, true
104+
default:
105+
return "4k", 0.40, true
106+
}
107+
}
108+
109+
func parseImageSizeForBilling(size string) (int, int, bool) {
110+
parts := strings.Split(strings.ToLower(strings.TrimSpace(size)), "x")
111+
if len(parts) != 2 {
112+
return 0, 0, false
113+
}
114+
width, err := strconv.Atoi(strings.TrimSpace(parts[0]))
115+
if err != nil || width <= 0 {
116+
return 0, 0, false
117+
}
118+
height, err := strconv.Atoi(strings.TrimSpace(parts[1]))
119+
if err != nil || height <= 0 {
120+
return 0, 0, false
121+
}
122+
return width, height, true
123+
}
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
package plugin
2+
3+
import (
4+
"math"
5+
"testing"
6+
7+
sdk "github.com/DouDOU-start/airgate-sdk"
8+
)
9+
10+
func TestImageOutputBillingOverride_UsesConfiguredTier(t *testing.T) {
11+
usage := &sdk.Usage{
12+
OutputCost: 0.40,
13+
ImageSize: "1672x941",
14+
}
15+
settings := map[string]map[string]string{
16+
"openai": {
17+
"image_price_2k": "0.08",
18+
},
19+
}
20+
21+
got, ok := imageOutputBillingOverride(usage, settings)
22+
if !ok {
23+
t.Fatal("expected override")
24+
}
25+
if math.Abs(got-0.16) > 1e-9 {
26+
t.Fatalf("override = %v, want 0.16 for two 2K images", got)
27+
}
28+
}
29+
30+
func TestImageOutputBillingOverride_FallsBackWhenTierUnset(t *testing.T) {
31+
usage := &sdk.Usage{
32+
OutputCost: 0.40,
33+
ImageSize: "3840x2160",
34+
}
35+
settings := map[string]map[string]string{
36+
"openai": {
37+
"image_price_2k": "0.08",
38+
},
39+
}
40+
41+
if got, ok := imageOutputBillingOverride(usage, settings); ok {
42+
t.Fatalf("override = %v, want fallback", got)
43+
}
44+
}
45+
46+
func TestImageTierForSize(t *testing.T) {
47+
tests := []struct {
48+
size string
49+
wantTier string
50+
wantPrice float64
51+
}{
52+
{size: "1024x1024", wantTier: "1k", wantPrice: 0.10},
53+
{size: "1672x941", wantTier: "2k", wantPrice: 0.20},
54+
{size: "3840x2160", wantTier: "4k", wantPrice: 0.40},
55+
}
56+
57+
for _, tt := range tests {
58+
t.Run(tt.size, func(t *testing.T) {
59+
tier, price, ok := imageTierForSize(tt.size)
60+
if !ok {
61+
t.Fatal("expected tier")
62+
}
63+
if tier != tt.wantTier || price != tt.wantPrice {
64+
t.Fatalf("imageTierForSize() = (%q, %v), want (%q, %v)", tier, price, tt.wantTier, tt.wantPrice)
65+
}
66+
})
67+
}
68+
}
69+
70+
func TestShouldForwardPluginSetting_HidesImagePrices(t *testing.T) {
71+
if shouldForwardPluginSetting("openai", "image_price_1k") {
72+
t.Fatal("image price settings should stay inside core")
73+
}
74+
if !shouldForwardPluginSetting("openai", "image_enabled") {
75+
t.Fatal("image_enabled should still be forwarded to the plugin")
76+
}
77+
if !shouldForwardPluginSetting("claude", "claude_code_only") {
78+
t.Fatal("non-openai plugin settings should still be forwarded")
79+
}
80+
}

backend/internal/plugin/outcome.go

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -261,15 +261,19 @@ func (f *Forwarder) recordUsage(c *gin.Context, state *forwardState, execution f
261261
// billingRate: 平台对 reseller 的计费倍率(group/user 优先级链)
262262
// sellRate: reseller 对客户的销售倍率(独立 markup 管道)
263263
// accountRate: 账号自身的真实成本系数("账号计费"统计管道)
264-
calc := f.calculator.Calculate(billing.CalculateInput{
264+
calcInput := billing.CalculateInput{
265265
InputCost: usage.InputCost,
266266
OutputCost: usage.OutputCost,
267267
CachedInputCost: usage.CachedInputCost,
268268
CacheCreationCost: usage.CacheCreationCost,
269269
BillingRate: billing.ResolveBillingRate(state.keyInfo),
270270
SellRate: state.keyInfo.SellRate,
271271
AccountRate: state.account.RateMultiplier,
272-
})
272+
}
273+
if override, ok := imageOutputBillingOverride(usage, state.keyInfo.GroupPluginSettings); ok {
274+
calcInput.OutputBillingCostOverride = &override
275+
}
276+
calc := f.calculator.Calculate(calcInput)
273277

274278
// 窗口费用沿用 account_cost(= total × account_rate),与用户账单解耦。
275279
f.scheduler.AddWindowCost(c.Request.Context(), state.account.ID, calc.AccountCost)

backend/internal/plugin/request.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -283,7 +283,7 @@ func buildHeaders(source http.Header, keyInfo *auth.APIKeyInfo) http.Header {
283283
// 分组级插件开关:X-Airgate-Plugin-{plugin}-{key} 约定。
284284
for plugin, kv := range keyInfo.GroupPluginSettings {
285285
for k, v := range kv {
286-
if v == "" {
286+
if v == "" || !shouldForwardPluginSetting(plugin, k) {
287287
continue
288288
}
289289
headers.Set("X-Airgate-Plugin-"+canonicalHeaderToken(plugin)+"-"+canonicalHeaderToken(k), v)

backend/internal/server/dto/usage.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ type CustomerUsageLogResp struct {
6767
Stream bool `json:"stream"`
6868
DurationMs int64 `json:"duration_ms"`
6969
FirstTokenMs int64 `json:"first_token_ms"`
70+
Endpoint string `json:"endpoint,omitempty"`
7071
CreatedAt string `json:"created_at"`
7172
}
7273

backend/internal/server/handler/usage_handler_mapper.go

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ func toCustomerUsageLogResp(record appusage.LogRecord) dto.CustomerUsageLogResp
7474
Stream: record.Stream,
7575
DurationMs: record.DurationMs,
7676
FirstTokenMs: record.FirstTokenMs,
77+
Endpoint: record.Endpoint,
7778
CreatedAt: record.CreatedAt,
7879
}
7980
}

web/src/i18n/en.json

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -477,6 +477,11 @@
477477
"service_tier_all": "All Tiers",
478478
"select_platform": "Select platform",
479479
"rate_multiplier": "Rate Multiplier",
480+
"image_generation": "Image generation",
481+
"image_pricing": "Image fixed unit prices",
482+
"image_price_unit_short": "image",
483+
"image_price_fallback": "Empty uses rate",
484+
"image_price_fallback_detail": "No image unit price, uses rate multiplier",
480485
"service_tier": "Service Tier",
481486
"service_tier_default": "Default",
482487
"exclusive": "Exclusive",
@@ -657,6 +662,7 @@
657662
"billed_cost": "Customer Billed",
658663
"profit": "Profit",
659664
"image_size": "Image Size",
665+
"endpoint": "Endpoint",
660666
"first_token": "TTFT",
661667
"stream": "Stream",
662668
"type": "Type",

0 commit comments

Comments
 (0)