oadmin/sub2api
1
0
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases Wiki Activity

Merge PR #166: feat: 图片生成计费功能

Browse Source
This commit is contained in:
shaw
2026-01-06 11:29:35 +08:00
parent f5603b0780 195e227c04
commit cb72262ad8
44 changed files with 2799 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
backend/internal/handler/admin/account_handler.go
View File

@@ -34,15 +34,16 @@ func NewOAuthHandler(oauthService *service.OAuthService) *OAuthHandler {
// AccountHandler handles admin account management
type AccountHandler struct {
adminService service.AdminService
oauthService *service.OAuthService
openaiOAuthService *service.OpenAIOAuthService
geminiOAuthService *service.GeminiOAuthService
rateLimitService *service.RateLimitService
accountUsageService *service.AccountUsageService
accountTestService *service.AccountTestService
concurrencyService *service.ConcurrencyService
crsSyncService *service.CRSSyncService
adminService service.AdminService
oauthService *service.OAuthService
openaiOAuthService *service.OpenAIOAuthService
geminiOAuthService *service.GeminiOAuthService
antigravityOAuthService *service.AntigravityOAuthService
rateLimitService *service.RateLimitService
accountUsageService *service.AccountUsageService
accountTestService *service.AccountTestService
concurrencyService *service.ConcurrencyService
crsSyncService *service.CRSSyncService
}
// NewAccountHandler creates a new admin account handler
@@ -51,6 +52,7 @@ func NewAccountHandler(
oauthService *service.OAuthService,
openaiOAuthService *service.OpenAIOAuthService,
geminiOAuthService *service.GeminiOAuthService,
antigravityOAuthService *service.AntigravityOAuthService,
rateLimitService *service.RateLimitService,
accountUsageService *service.AccountUsageService,
accountTestService *service.AccountTestService,
@@ -58,15 +60,16 @@ func NewAccountHandler(
crsSyncService *service.CRSSyncService,
) *AccountHandler {
return &AccountHandler{
adminService: adminService,
oauthService: oauthService,
openaiOAuthService: openaiOAuthService,
geminiOAuthService: geminiOAuthService,
rateLimitService: rateLimitService,
accountUsageService: accountUsageService,
accountTestService: accountTestService,
concurrencyService: concurrencyService,
crsSyncService: crsSyncService,
adminService: adminService,
oauthService: oauthService,
openaiOAuthService: openaiOAuthService,
geminiOAuthService: geminiOAuthService,
antigravityOAuthService: antigravityOAuthService,
rateLimitService: rateLimitService,
accountUsageService: accountUsageService,
accountTestService: accountTestService,
concurrencyService: concurrencyService,
crsSyncService: crsSyncService,
}
}
@@ -420,6 +423,19 @@ func (h *AccountHandler) Refresh(c *gin.Context) {
newCredentials[k] = v
}
}
} else if account.Platform == service.PlatformAntigravity {
tokenInfo, err := h.antigravityOAuthService.RefreshAccountToken(c.Request.Context(), account)
if err != nil {
response.ErrorFrom(c, err)
return
}
newCredentials = h.antigravityOAuthService.BuildAccountCredentials(tokenInfo)
for k, v := range account.Credentials {
if _, exists := newCredentials[k]; !exists {
newCredentials[k] = v
}
}
} else {
// Use Anthropic/Claude OAuth service to refresh token
tokenInfo, err := h.oauthService.RefreshAccountToken(c.Request.Context(), account)

14
backend/internal/handler/admin/group_handler.go
View File

@@ -33,6 +33,10 @@ type CreateGroupRequest struct {
DailyLimitUSD *float64 `json:"daily_limit_usd"`
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
// 图片生成计费配置（antigravity 和 gemini 平台使用，负数表示清除配置）
ImagePrice1K *float64 `json:"image_price_1k"`
ImagePrice2K *float64 `json:"image_price_2k"`
ImagePrice4K *float64 `json:"image_price_4k"`
}
// UpdateGroupRequest represents update group request
@@ -47,6 +51,10 @@ type UpdateGroupRequest struct {
DailyLimitUSD *float64 `json:"daily_limit_usd"`
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
// 图片生成计费配置（antigravity 和 gemini 平台使用，负数表示清除配置）
ImagePrice1K *float64 `json:"image_price_1k"`
ImagePrice2K *float64 `json:"image_price_2k"`
ImagePrice4K *float64 `json:"image_price_4k"`
}
// List handles listing all groups with pagination
@@ -139,6 +147,9 @@ func (h *GroupHandler) Create(c *gin.Context) {
DailyLimitUSD: req.DailyLimitUSD,
WeeklyLimitUSD: req.WeeklyLimitUSD,
MonthlyLimitUSD: req.MonthlyLimitUSD,
ImagePrice1K: req.ImagePrice1K,
ImagePrice2K: req.ImagePrice2K,
ImagePrice4K: req.ImagePrice4K,
})
if err != nil {
response.ErrorFrom(c, err)
@@ -174,6 +185,9 @@ func (h *GroupHandler) Update(c *gin.Context) {
DailyLimitUSD: req.DailyLimitUSD,
WeeklyLimitUSD: req.WeeklyLimitUSD,
MonthlyLimitUSD: req.MonthlyLimitUSD,
ImagePrice1K: req.ImagePrice1K,
ImagePrice2K: req.ImagePrice2K,
ImagePrice4K: req.ImagePrice4K,
})
if err != nil {
response.ErrorFrom(c, err)

5
backend/internal/handler/dto/mappers.go
View File

@@ -78,6 +78,9 @@ func GroupFromServiceShallow(g *service.Group) *Group {
DailyLimitUSD: g.DailyLimitUSD,
WeeklyLimitUSD: g.WeeklyLimitUSD,
MonthlyLimitUSD: g.MonthlyLimitUSD,
ImagePrice1K: g.ImagePrice1K,
ImagePrice2K: g.ImagePrice2K,
ImagePrice4K: g.ImagePrice4K,
CreatedAt: g.CreatedAt,
UpdatedAt: g.UpdatedAt,
AccountCount: g.AccountCount,
@@ -247,6 +250,8 @@ func UsageLogFromService(l *service.UsageLog) *UsageLog {
Stream: l.Stream,
DurationMs: l.DurationMs,
FirstTokenMs: l.FirstTokenMs,
ImageCount: l.ImageCount,
ImageSize: l.ImageSize,
CreatedAt: l.CreatedAt,
User: UserFromServiceShallow(l.User),
APIKey: APIKeyFromService(l.APIKey),

9
backend/internal/handler/dto/types.go
View File

@@ -47,6 +47,11 @@ type Group struct {
WeeklyLimitUSD *float64 `json:"weekly_limit_usd"`
MonthlyLimitUSD *float64 `json:"monthly_limit_usd"`
// 图片生成计费配置（仅 antigravity 平台使用）
ImagePrice1K *float64 `json:"image_price_1k"`
ImagePrice2K *float64 `json:"image_price_2k"`
ImagePrice4K *float64 `json:"image_price_4k"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
@@ -169,6 +174,10 @@ type UsageLog struct {
DurationMs *int `json:"duration_ms"`
FirstTokenMs *int `json:"first_token_ms"`
// 图片生成字段
ImageCount int `json:"image_count"`
ImageSize *string `json:"image_size"`
CreatedAt time.Time `json:"created_at"`
User *User `json:"user,omitempty"`

7
backend/internal/pkg/antigravity/gemini_types.go
View File

@@ -67,6 +67,13 @@ type GeminiGenerationConfig struct {
TopK *int `json:"topK,omitempty"`
ThinkingConfig *GeminiThinkingConfig `json:"thinkingConfig,omitempty"`
StopSequences []string `json:"stopSequences,omitempty"`
ImageConfig *GeminiImageConfig `json:"imageConfig,omitempty"`
}
// GeminiImageConfig Gemini 图片生成配置（仅 gemini-3-pro-image 支持）
type GeminiImageConfig struct {
AspectRatio string `json:"aspectRatio,omitempty"` // "1:1", "16:9", "9:16", "4:3", "3:4"
ImageSize string `json:"imageSize,omitempty"` // "1K", "2K", "4K"
}
// GeminiThinkingConfig Gemini thinking 配置

3
backend/internal/repository/api_key_repo.go
View File

@@ -321,6 +321,9 @@ func groupEntityToService(g *dbent.Group) *service.Group {
DailyLimitUSD: g.DailyLimitUsd,
WeeklyLimitUSD: g.WeeklyLimitUsd,
MonthlyLimitUSD: g.MonthlyLimitUsd,
ImagePrice1K: g.ImagePrice1k,
ImagePrice2K: g.ImagePrice2k,
ImagePrice4K: g.ImagePrice4k,
DefaultValidityDays: g.DefaultValidityDays,
CreatedAt: g.CreatedAt,
UpdatedAt: g.UpdatedAt,

6
backend/internal/repository/group_repo.go
View File

@@ -43,6 +43,9 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er
SetNillableDailyLimitUsd(groupIn.DailyLimitUSD).
SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD).
SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD).
SetNillableImagePrice1k(groupIn.ImagePrice1K).
SetNillableImagePrice2k(groupIn.ImagePrice2K).
SetNillableImagePrice4k(groupIn.ImagePrice4K).
SetDefaultValidityDays(groupIn.DefaultValidityDays)
created, err := builder.Save(ctx)
@@ -80,6 +83,9 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
SetNillableDailyLimitUsd(groupIn.DailyLimitUSD).
SetNillableWeeklyLimitUsd(groupIn.WeeklyLimitUSD).
SetNillableMonthlyLimitUsd(groupIn.MonthlyLimitUSD).
SetNillableImagePrice1k(groupIn.ImagePrice1K).
SetNillableImagePrice2k(groupIn.ImagePrice2K).
SetNillableImagePrice4k(groupIn.ImagePrice4K).
SetDefaultValidityDays(groupIn.DefaultValidityDays).
Save(ctx)
if err != nil {

25
backend/internal/repository/usage_log_repo.go
View File

@@ -22,7 +22,7 @@ import (
"github.com/lib/pq"
)
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, billing_type, stream, duration_ms, first_token_ms, created_at"
const usageLogSelectColumns = "id, user_id, api_key_id, account_id, request_id, model, group_id, subscription_id, input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens, cache_creation_5m_tokens, cache_creation_1h_tokens, input_cost, output_cost, cache_creation_cost, cache_read_cost, total_cost, actual_cost, rate_multiplier, billing_type, stream, duration_ms, first_token_ms, image_count, image_size, created_at"
type usageLogRepository struct {
client *dbent.Client
@@ -109,6 +109,8 @@ func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog)
stream,
duration_ms,
first_token_ms,
image_count,
image_size,
created_at
) VALUES (
$1, $2, $3, $4, $5,
@@ -116,7 +118,8 @@ func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog)
$8, $9, $10, $11,
$12, $13,
$14, $15, $16, $17, $18, $19,
$20, $21, $22, $23, $24, $25
$20, $21, $22, $23, $24,
$25, $26, $27
)
ON CONFLICT (request_id, api_key_id) DO NOTHING
RETURNING id, created_at
@@ -126,6 +129,7 @@ func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog)
subscriptionID := nullInt64(log.SubscriptionID)
duration := nullInt(log.DurationMs)
firstToken := nullInt(log.FirstTokenMs)
imageSize := nullString(log.ImageSize)
var requestIDArg any
if requestID != "" {
@@ -157,6 +161,8 @@ func (r *usageLogRepository) Create(ctx context.Context, log *service.UsageLog)
log.Stream,
duration,
firstToken,
log.ImageCount,
imageSize,
createdAt,
}
if err := scanSingleRow(ctx, sqlq, query, args, &log.ID, &log.CreatedAt); err != nil {
@@ -1789,6 +1795,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
stream bool
durationMs sql.NullInt64
firstTokenMs sql.NullInt64
imageCount int
imageSize sql.NullString
createdAt time.Time
)
@@ -1818,6 +1826,8 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
&stream,
&durationMs,
&firstTokenMs,
&imageCount,
&imageSize,
&createdAt,
); err != nil {
return nil, err
@@ -1844,6 +1854,7 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
RateMultiplier: rateMultiplier,
BillingType: int8(billingType),
Stream: stream,
ImageCount: imageCount,
CreatedAt: createdAt,
}
@@ -1866,6 +1877,9 @@ func scanUsageLog(scanner interface{ Scan(...any) error }) (*service.UsageLog, e
value := int(firstTokenMs.Int64)
log.FirstTokenMs = &value
}
if imageSize.Valid {
log.ImageSize = &imageSize.String
}
return log, nil
}
@@ -1938,6 +1952,13 @@ func nullInt(v *int) sql.NullInt64 {
return sql.NullInt64{Int64: int64(*v), Valid: true}
}
func nullString(v *string) sql.NullString {
if v == nil || *v == "" {
return sql.NullString{}
}
return sql.NullString{String: *v, Valid: true}
}
func setToSlice(set map[int64]struct{}) []int64 {
out := make([]int64, 0, len(set))
for id := range set {

2
backend/internal/server/api_contract_test.go
View File

@@ -241,6 +241,8 @@ func TestAPIContracts(t *testing.T) {
"stream": true,
"duration_ms": 100,
"first_token_ms": 50,
"image_count": 0,
"image_size": null,
"created_at": "2025-01-02T03:04:05Z"
}
],

34
backend/internal/service/admin_service.go
View File

@@ -98,6 +98,10 @@ type CreateGroupInput struct {
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
// 图片生成计费配置（仅 antigravity 平台使用）
ImagePrice1K *float64
ImagePrice2K *float64
ImagePrice4K *float64
}
type UpdateGroupInput struct {
@@ -111,6 +115,10 @@ type UpdateGroupInput struct {
DailyLimitUSD *float64 // 日限额 (USD)
WeeklyLimitUSD *float64 // 周限额 (USD)
MonthlyLimitUSD *float64 // 月限额 (USD)
// 图片生成计费配置（仅 antigravity 平台使用）
ImagePrice1K *float64
ImagePrice2K *float64
ImagePrice4K *float64
}
type CreateAccountInput struct {
@@ -498,6 +506,11 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
weeklyLimit := normalizeLimit(input.WeeklyLimitUSD)
monthlyLimit := normalizeLimit(input.MonthlyLimitUSD)
// 图片价格：负数表示清除（使用默认价格），0 保留（表示免费）
imagePrice1K := normalizePrice(input.ImagePrice1K)
imagePrice2K := normalizePrice(input.ImagePrice2K)
imagePrice4K := normalizePrice(input.ImagePrice4K)
group := &Group{
Name: input.Name,
Description: input.Description,
@@ -509,6 +522,9 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
DailyLimitUSD: dailyLimit,
WeeklyLimitUSD: weeklyLimit,
MonthlyLimitUSD: monthlyLimit,
ImagePrice1K: imagePrice1K,
ImagePrice2K: imagePrice2K,
ImagePrice4K: imagePrice4K,
}
if err := s.groupRepo.Create(ctx, group); err != nil {
return nil, err
@@ -524,6 +540,14 @@ func normalizeLimit(limit *float64) *float64 {
return limit
}
// normalizePrice 将负数转换为 nil（表示使用默认价格），0 保留（表示免费）
func normalizePrice(price *float64) *float64 {
if price == nil || *price < 0 {
return nil
}
return price
}
func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error) {
group, err := s.groupRepo.GetByID(ctx, id)
if err != nil {
@@ -563,6 +587,16 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
if input.MonthlyLimitUSD != nil {
group.MonthlyLimitUSD = normalizeLimit(input.MonthlyLimitUSD)
}
// 图片生成计费配置：负数表示清除（使用默认价格）
if input.ImagePrice1K != nil {
group.ImagePrice1K = normalizePrice(input.ImagePrice1K)
}
if input.ImagePrice2K != nil {
group.ImagePrice2K = normalizePrice(input.ImagePrice2K)
}
if input.ImagePrice4K != nil {
group.ImagePrice4K = normalizePrice(input.ImagePrice4K)
}
if err := s.groupRepo.Update(ctx, group); err != nil {
return nil, err

197
backend/internal/service/admin_service_group_test.go Normal file
View File

@@ -0,0 +1,197 @@
//go:build unit
package service
import (
"context"
"testing"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/stretchr/testify/require"
)
// groupRepoStubForAdmin 用于测试 AdminService 的 GroupRepository Stub
type groupRepoStubForAdmin struct {
created *Group // 记录 Create 调用的参数
updated *Group // 记录 Update 调用的参数
getByID *Group // GetByID 返回值
getErr error // GetByID 返回的错误
}
func (s *groupRepoStubForAdmin) Create(_ context.Context, g *Group) error {
s.created = g
return nil
}
func (s *groupRepoStubForAdmin) Update(_ context.Context, g *Group) error {
s.updated = g
return nil
}
func (s *groupRepoStubForAdmin) GetByID(_ context.Context, _ int64) (*Group, error) {
if s.getErr != nil {
return nil, s.getErr
}
return s.getByID, nil
}
func (s *groupRepoStubForAdmin) Delete(_ context.Context, _ int64) error {
panic("unexpected Delete call")
}
func (s *groupRepoStubForAdmin) DeleteCascade(_ context.Context, _ int64) ([]int64, error) {
panic("unexpected DeleteCascade call")
}
func (s *groupRepoStubForAdmin) List(_ context.Context, _ pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected List call")
}
func (s *groupRepoStubForAdmin) ListWithFilters(_ context.Context, _ pagination.PaginationParams, _, _ string, _ *bool) ([]Group, *pagination.PaginationResult, error) {
panic("unexpected ListWithFilters call")
}
func (s *groupRepoStubForAdmin) ListActive(_ context.Context) ([]Group, error) {
panic("unexpected ListActive call")
}
func (s *groupRepoStubForAdmin) ListActiveByPlatform(_ context.Context, _ string) ([]Group, error) {
panic("unexpected ListActiveByPlatform call")
}
func (s *groupRepoStubForAdmin) ExistsByName(_ context.Context, _ string) (bool, error) {
panic("unexpected ExistsByName call")
}
func (s *groupRepoStubForAdmin) GetAccountCount(_ context.Context, _ int64) (int64, error) {
panic("unexpected GetAccountCount call")
}
func (s *groupRepoStubForAdmin) DeleteAccountGroupsByGroupID(_ context.Context, _ int64) (int64, error) {
panic("unexpected DeleteAccountGroupsByGroupID call")
}
// TestAdminService_CreateGroup_WithImagePricing 测试创建分组时 ImagePrice 字段正确传递
func TestAdminService_CreateGroup_WithImagePricing(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
price1K := 0.10
price2K := 0.15
price4K := 0.30
input := &CreateGroupInput{
Name: "test-group",
Description: "Test group",
Platform: PlatformAntigravity,
RateMultiplier: 1.0,
ImagePrice1K: &price1K,
ImagePrice2K: &price2K,
ImagePrice4K: &price4K,
}
group, err := svc.CreateGroup(context.Background(), input)
require.NoError(t, err)
require.NotNil(t, group)
// 验证 repo 收到了正确的字段
require.NotNil(t, repo.created)
require.NotNil(t, repo.created.ImagePrice1K)
require.NotNil(t, repo.created.ImagePrice2K)
require.NotNil(t, repo.created.ImagePrice4K)
require.InDelta(t, 0.10, *repo.created.ImagePrice1K, 0.0001)
require.InDelta(t, 0.15, *repo.created.ImagePrice2K, 0.0001)
require.InDelta(t, 0.30, *repo.created.ImagePrice4K, 0.0001)
}
// TestAdminService_CreateGroup_NilImagePricing 测试 ImagePrice 为 nil 时正常创建
func TestAdminService_CreateGroup_NilImagePricing(t *testing.T) {
repo := &groupRepoStubForAdmin{}
svc := &adminServiceImpl{groupRepo: repo}
input := &CreateGroupInput{
Name: "test-group",
Description: "Test group",
Platform: PlatformAntigravity,
RateMultiplier: 1.0,
// ImagePrice 字段全部为 nil
}
group, err := svc.CreateGroup(context.Background(), input)
require.NoError(t, err)
require.NotNil(t, group)
// 验证 ImagePrice 字段为 nil
require.NotNil(t, repo.created)
require.Nil(t, repo.created.ImagePrice1K)
require.Nil(t, repo.created.ImagePrice2K)
require.Nil(t, repo.created.ImagePrice4K)
}
// TestAdminService_UpdateGroup_WithImagePricing 测试更新分组时 ImagePrice 字段正确更新
func TestAdminService_UpdateGroup_WithImagePricing(t *testing.T) {
existingGroup := &Group{
ID: 1,
Name: "existing-group",
Platform: PlatformAntigravity,
Status: StatusActive,
}
repo := &groupRepoStubForAdmin{getByID: existingGroup}
svc := &adminServiceImpl{groupRepo: repo}
price1K := 0.12
price2K := 0.18
price4K := 0.36
input := &UpdateGroupInput{
ImagePrice1K: &price1K,
ImagePrice2K: &price2K,
ImagePrice4K: &price4K,
}
group, err := svc.UpdateGroup(context.Background(), 1, input)
require.NoError(t, err)
require.NotNil(t, group)
// 验证 repo 收到了更新后的字段
require.NotNil(t, repo.updated)
require.NotNil(t, repo.updated.ImagePrice1K)
require.NotNil(t, repo.updated.ImagePrice2K)
require.NotNil(t, repo.updated.ImagePrice4K)
require.InDelta(t, 0.12, *repo.updated.ImagePrice1K, 0.0001)
require.InDelta(t, 0.18, *repo.updated.ImagePrice2K, 0.0001)
require.InDelta(t, 0.36, *repo.updated.ImagePrice4K, 0.0001)
}
// TestAdminService_UpdateGroup_PartialImagePricing 测试仅更新部分 ImagePrice 字段
func TestAdminService_UpdateGroup_PartialImagePricing(t *testing.T) {
oldPrice2K := 0.15
existingGroup := &Group{
ID: 1,
Name: "existing-group",
Platform: PlatformAntigravity,
Status: StatusActive,
ImagePrice2K: &oldPrice2K, // 已有 2K 价格
}
repo := &groupRepoStubForAdmin{getByID: existingGroup}
svc := &adminServiceImpl{groupRepo: repo}
// 只更新 1K 价格
price1K := 0.10
input := &UpdateGroupInput{
ImagePrice1K: &price1K,
// ImagePrice2K 和 ImagePrice4K 为 nil，不更新
}
group, err := svc.UpdateGroup(context.Background(), 1, input)
require.NoError(t, err)
require.NotNil(t, group)
// 验证：1K 被更新，2K 保持原值，4K 仍为 nil
require.NotNil(t, repo.updated)
require.NotNil(t, repo.updated.ImagePrice1K)
require.InDelta(t, 0.10, *repo.updated.ImagePrice1K, 0.0001)
require.NotNil(t, repo.updated.ImagePrice2K)
require.InDelta(t, 0.15, *repo.updated.ImagePrice2K, 0.0001) // 原值保持
require.Nil(t, repo.updated.ImagePrice4K)
}

45
backend/internal/service/antigravity_gateway_service.go
View File

@@ -860,6 +860,9 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
return nil, s.writeGoogleError(c, http.StatusBadRequest, "Request body is empty")
}
// 解析请求以获取 image_size（用于图片计费）
imageSize := s.extractImageSize(body)
switch action {
case "generateContent", "streamGenerateContent":
// ok
@@ -1059,6 +1062,13 @@ handleSuccess:
usage = &ClaudeUsage{}
}
// 判断是否为图片生成模型
imageCount := 0
if isImageGenerationModel(mappedModel) {
// Gemini 图片生成 API 每次请求只生成一张图片（API 限制）
imageCount = 1
}
return &ForwardResult{
RequestID: requestID,
Usage: *usage,
@@ -1066,6 +1076,8 @@ handleSuccess:
Stream: stream,
Duration: time.Since(startTime),
FirstTokenMs: firstTokenMs,
ImageCount: imageCount,
ImageSize: imageSize,
}, nil
}
@@ -1572,3 +1584,36 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context
}
}
// extractImageSize 从 Gemini 请求中提取 image_size 参数
func (s *AntigravityGatewayService) extractImageSize(body []byte) string {
var req antigravity.GeminiRequest
if err := json.Unmarshal(body, &req); err != nil {
return "2K" // 默认 2K
}
if req.GenerationConfig != nil && req.GenerationConfig.ImageConfig != nil {
size := strings.ToUpper(strings.TrimSpace(req.GenerationConfig.ImageConfig.ImageSize))
if size == "1K" || size == "2K" || size == "4K" {
return size
}
}
return "2K" // 默认 2K
}
// isImageGenerationModel 判断模型是否为图片生成模型
// 支持的模型：gemini-3-pro-image, gemini-3-pro-image-preview, gemini-2.5-flash-image 等
func isImageGenerationModel(model string) bool {
modelLower := strings.ToLower(model)
// 移除 models/ 前缀
modelLower = strings.TrimPrefix(modelLower, "models/")
// 精确匹配或前缀匹配
return modelLower == "gemini-3-pro-image" ||
modelLower == "gemini-3-pro-image-preview" ||
strings.HasPrefix(modelLower, "gemini-3-pro-image-") ||
modelLower == "gemini-2.5-flash-image" ||
modelLower == "gemini-2.5-flash-image-preview" ||
strings.HasPrefix(modelLower, "gemini-2.5-flash-image-")
}

123
backend/internal/service/antigravity_image_test.go Normal file
View File

@@ -0,0 +1,123 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
// TestIsImageGenerationModel_GeminiProImage 测试 gemini-3-pro-image 识别
func TestIsImageGenerationModel_GeminiProImage(t *testing.T) {
require.True(t, isImageGenerationModel("gemini-3-pro-image"))
require.True(t, isImageGenerationModel("gemini-3-pro-image-preview"))
require.True(t, isImageGenerationModel("models/gemini-3-pro-image"))
}
// TestIsImageGenerationModel_GeminiFlashImage 测试 gemini-2.5-flash-image 识别
func TestIsImageGenerationModel_GeminiFlashImage(t *testing.T) {
require.True(t, isImageGenerationModel("gemini-2.5-flash-image"))
require.True(t, isImageGenerationModel("gemini-2.5-flash-image-preview"))
}
// TestIsImageGenerationModel_RegularModel 测试普通模型不被识别为图片模型
func TestIsImageGenerationModel_RegularModel(t *testing.T) {
require.False(t, isImageGenerationModel("claude-3-opus"))
require.False(t, isImageGenerationModel("claude-sonnet-4-20250514"))
require.False(t, isImageGenerationModel("gpt-4o"))
require.False(t, isImageGenerationModel("gemini-2.5-pro")) // 非图片模型
require.False(t, isImageGenerationModel("gemini-2.5-flash"))
// 验证不会误匹配包含关键词的自定义模型名
require.False(t, isImageGenerationModel("my-gemini-3-pro-image-test"))
require.False(t, isImageGenerationModel("custom-gemini-2.5-flash-image-wrapper"))
}
// TestIsImageGenerationModel_CaseInsensitive 测试大小写不敏感
func TestIsImageGenerationModel_CaseInsensitive(t *testing.T) {
require.True(t, isImageGenerationModel("GEMINI-3-PRO-IMAGE"))
require.True(t, isImageGenerationModel("Gemini-3-Pro-Image"))
require.True(t, isImageGenerationModel("GEMINI-2.5-FLASH-IMAGE"))
}
// TestExtractImageSize_ValidSizes 测试有效尺寸解析
func TestExtractImageSize_ValidSizes(t *testing.T) {
svc := &AntigravityGatewayService{}
// 1K
body := []byte(`{"generationConfig":{"imageConfig":{"imageSize":"1K"}}}`)
require.Equal(t, "1K", svc.extractImageSize(body))
// 2K
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":"2K"}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
// 4K
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":"4K"}}}`)
require.Equal(t, "4K", svc.extractImageSize(body))
}
// TestExtractImageSize_CaseInsensitive 测试大小写不敏感
func TestExtractImageSize_CaseInsensitive(t *testing.T) {
svc := &AntigravityGatewayService{}
body := []byte(`{"generationConfig":{"imageConfig":{"imageSize":"1k"}}}`)
require.Equal(t, "1K", svc.extractImageSize(body))
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":"4k"}}}`)
require.Equal(t, "4K", svc.extractImageSize(body))
}
// TestExtractImageSize_Default 测试无 imageConfig 返回默认 2K
func TestExtractImageSize_Default(t *testing.T) {
svc := &AntigravityGatewayService{}
// 无 generationConfig
body := []byte(`{"contents":[]}`)
require.Equal(t, "2K", svc.extractImageSize(body))
// 有 generationConfig 但无 imageConfig
body = []byte(`{"generationConfig":{"temperature":0.7}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
// 有 imageConfig 但无 imageSize
body = []byte(`{"generationConfig":{"imageConfig":{}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
}
// TestExtractImageSize_InvalidJSON 测试非法 JSON 返回默认 2K
func TestExtractImageSize_InvalidJSON(t *testing.T) {
svc := &AntigravityGatewayService{}
body := []byte(`not valid json`)
require.Equal(t, "2K", svc.extractImageSize(body))
body = []byte(`{"broken":`)
require.Equal(t, "2K", svc.extractImageSize(body))
}
// TestExtractImageSize_EmptySize 测试空 imageSize 返回默认 2K
func TestExtractImageSize_EmptySize(t *testing.T) {
svc := &AntigravityGatewayService{}
body := []byte(`{"generationConfig":{"imageConfig":{"imageSize":""}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
// 空格
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":" "}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
}
// TestExtractImageSize_InvalidSize 测试无效尺寸返回默认 2K
func TestExtractImageSize_InvalidSize(t *testing.T) {
svc := &AntigravityGatewayService{}
body := []byte(`{"generationConfig":{"imageConfig":{"imageSize":"3K"}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":"8K"}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
body = []byte(`{"generationConfig":{"imageConfig":{"imageSize":"invalid"}}}`)
require.Equal(t, "2K", svc.extractImageSize(body))
}

85
backend/internal/service/billing_service.go
View File

@@ -295,3 +295,88 @@ func (s *BillingService) ForceUpdatePricing() error {
}
return fmt.Errorf("pricing service not initialized")
}
// ImagePriceConfig 图片计费配置
type ImagePriceConfig struct {
Price1K *float64 // 1K 尺寸价格（nil 表示使用默认值）
Price2K *float64 // 2K 尺寸价格（nil 表示使用默认值）
Price4K *float64 // 4K 尺寸价格（nil 表示使用默认值）
}
// CalculateImageCost 计算图片生成费用
// model: 请求的模型名称（用于获取 LiteLLM 默认价格）
// imageSize: 图片尺寸 "1K", "2K", "4K"
// imageCount: 生成的图片数量
// groupConfig: 分组配置的价格（可能为 nil，表示使用默认值）
// rateMultiplier: 费率倍数
func (s *BillingService) CalculateImageCost(model string, imageSize string, imageCount int, groupConfig *ImagePriceConfig, rateMultiplier float64) *CostBreakdown {
if imageCount <= 0 {
return &CostBreakdown{}
}
// 获取单价
unitPrice := s.getImageUnitPrice(model, imageSize, groupConfig)
// 计算总费用
totalCost := unitPrice * float64(imageCount)
// 应用倍率
if rateMultiplier <= 0 {
rateMultiplier = 1.0
}
actualCost := totalCost * rateMultiplier
return &CostBreakdown{
TotalCost: totalCost,
ActualCost: actualCost,
}
}
// getImageUnitPrice 获取图片单价
func (s *BillingService) getImageUnitPrice(model string, imageSize string, groupConfig *ImagePriceConfig) float64 {
// 优先使用分组配置的价格
if groupConfig != nil {
switch imageSize {
case "1K":
if groupConfig.Price1K != nil {
return *groupConfig.Price1K
}
case "2K":
if groupConfig.Price2K != nil {
return *groupConfig.Price2K
}
case "4K":
if groupConfig.Price4K != nil {
return *groupConfig.Price4K
}
}
}
// 回退到 LiteLLM 默认价格
return s.getDefaultImagePrice(model, imageSize)
}
// getDefaultImagePrice 获取 LiteLLM 默认图片价格
func (s *BillingService) getDefaultImagePrice(model string, imageSize string) float64 {
basePrice := 0.0
// 从 PricingService 获取 output_cost_per_image
if s.pricingService != nil {
pricing := s.pricingService.GetModelPricing(model)
if pricing != nil && pricing.OutputCostPerImage > 0 {
basePrice = pricing.OutputCostPerImage
}
}
// 如果没有找到价格，使用硬编码默认值（$0.134，来自 gemini-3-pro-image-preview）
if basePrice <= 0 {
basePrice = 0.134
}
// 4K 尺寸翻倍
if imageSize == "4K" {
return basePrice * 2
}
return basePrice
}

149
backend/internal/service/billing_service_image_test.go Normal file
View File

@@ -0,0 +1,149 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
// TestCalculateImageCost_DefaultPricing 测试无分组配置时使用默认价格
func TestCalculateImageCost_DefaultPricing(t *testing.T) {
svc := &BillingService{} // pricingService 为 nil，使用硬编码默认值
// 2K 尺寸，默认价格 $0.134
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
require.InDelta(t, 0.134, cost.ActualCost, 0.0001)
// 多张图片
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 3, nil, 1.0)
require.InDelta(t, 0.402, cost.TotalCost, 0.0001)
}
// TestCalculateImageCost_GroupCustomPricing 测试分组自定义价格
func TestCalculateImageCost_GroupCustomPricing(t *testing.T) {
svc := &BillingService{}
price1K := 0.10
price2K := 0.15
price4K := 0.30
groupConfig := &ImagePriceConfig{
Price1K: &price1K,
Price2K: &price2K,
Price4K: &price4K,
}
// 1K 使用分组价格
cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 2, groupConfig, 1.0)
require.InDelta(t, 0.20, cost.TotalCost, 0.0001)
// 2K 使用分组价格
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, groupConfig, 1.0)
require.InDelta(t, 0.15, cost.TotalCost, 0.0001)
// 4K 使用分组价格
cost = svc.CalculateImageCost("gemini-3-pro-image", "4K", 1, groupConfig, 1.0)
require.InDelta(t, 0.30, cost.TotalCost, 0.0001)
}
// TestCalculateImageCost_4KDoublePrice 测试 4K 默认价格翻倍
func TestCalculateImageCost_4KDoublePrice(t *testing.T) {
svc := &BillingService{}
// 4K 尺寸，默认价格翻倍 $0.134 * 2 = $0.268
cost := svc.CalculateImageCost("gemini-3-pro-image", "4K", 1, nil, 1.0)
require.InDelta(t, 0.268, cost.TotalCost, 0.0001)
}
// TestCalculateImageCost_RateMultiplier 测试费率倍数
func TestCalculateImageCost_RateMultiplier(t *testing.T) {
svc := &BillingService{}
// 费率倍数 1.5x
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.5)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001) // TotalCost 不变
require.InDelta(t, 0.201, cost.ActualCost, 0.0001) // ActualCost = 0.134 * 1.5
// 费率倍数 2.0x
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 2, nil, 2.0)
require.InDelta(t, 0.268, cost.TotalCost, 0.0001)
require.InDelta(t, 0.536, cost.ActualCost, 0.0001)
}
// TestCalculateImageCost_ZeroCount 测试 imageCount=0
func TestCalculateImageCost_ZeroCount(t *testing.T) {
svc := &BillingService{}
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 0, nil, 1.0)
require.Equal(t, 0.0, cost.TotalCost)
require.Equal(t, 0.0, cost.ActualCost)
}
// TestCalculateImageCost_NegativeCount 测试 imageCount=-1
func TestCalculateImageCost_NegativeCount(t *testing.T) {
svc := &BillingService{}
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", -1, nil, 1.0)
require.Equal(t, 0.0, cost.TotalCost)
require.Equal(t, 0.0, cost.ActualCost)
}
// TestCalculateImageCost_ZeroRateMultiplier 测试费率倍数为 0 时默认使用 1.0
func TestCalculateImageCost_ZeroRateMultiplier(t *testing.T) {
svc := &BillingService{}
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
require.InDelta(t, 0.134, cost.ActualCost, 0.0001) // 0 倍率当作 1.0 处理
}
// TestGetImageUnitPrice_GroupPriorityOverDefault 测试分组价格优先于默认价格
func TestGetImageUnitPrice_GroupPriorityOverDefault(t *testing.T) {
svc := &BillingService{}
price2K := 0.20
groupConfig := &ImagePriceConfig{
Price2K: &price2K,
}
// 分组配置了 2K 价格，应该使用分组价格而不是默认的 $0.134
cost := svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, groupConfig, 1.0)
require.InDelta(t, 0.20, cost.TotalCost, 0.0001)
}
// TestGetImageUnitPrice_PartialGroupConfig 测试分组部分配置时回退默认
func TestGetImageUnitPrice_PartialGroupConfig(t *testing.T) {
svc := &BillingService{}
// 只配置 1K 价格
price1K := 0.10
groupConfig := &ImagePriceConfig{
Price1K: &price1K,
}
// 1K 使用分组价格
cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, groupConfig, 1.0)
require.InDelta(t, 0.10, cost.TotalCost, 0.0001)
// 2K 回退默认价格 $0.134
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, groupConfig, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
// 4K 回退默认价格 $0.268 (翻倍)
cost = svc.CalculateImageCost("gemini-3-pro-image", "4K", 1, groupConfig, 1.0)
require.InDelta(t, 0.268, cost.TotalCost, 0.0001)
}
// TestGetDefaultImagePrice_FallbackHardcoded 测试 PricingService 无数据时使用硬编码默认值
func TestGetDefaultImagePrice_FallbackHardcoded(t *testing.T) {
svc := &BillingService{} // pricingService 为 nil
// 1K 和 2K 使用相同的默认价格 $0.134
cost := svc.CalculateImageCost("gemini-3-pro-image", "1K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
cost = svc.CalculateImageCost("gemini-3-pro-image", "2K", 1, nil, 1.0)
require.InDelta(t, 0.134, cost.TotalCost, 0.0001)
}

51
backend/internal/service/gateway_service.go
View File

@@ -104,6 +104,10 @@ type ForwardResult struct {
Stream bool
Duration time.Duration
FirstTokenMs *int // 首字时间（流式请求）
// 图片生成计费字段（仅 gemini-3-pro-image 使用）
ImageCount int // 生成的图片数量
ImageSize string // 图片尺寸 "1K", "2K", "4K"
}
// UpstreamFailoverError indicates an upstream error that should trigger account failover.
@@ -2009,25 +2013,40 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
account := input.Account
subscription := input.Subscription
// 计算费用
tokens := UsageTokens{
InputTokens: result.Usage.InputTokens,
OutputTokens: result.Usage.OutputTokens,
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
CacheReadTokens: result.Usage.CacheReadInputTokens,
}
// 获取费率倍数
multiplier := s.cfg.Default.RateMultiplier
if apiKey.GroupID != nil && apiKey.Group != nil {
multiplier = apiKey.Group.RateMultiplier
}
cost, err := s.billingService.CalculateCost(result.Model, tokens, multiplier)
if err != nil {
log.Printf("Calculate cost failed: %v", err)
// 使用默认费用继续
cost = &CostBreakdown{ActualCost: 0}
var cost *CostBreakdown
// 根据请求类型选择计费方式
if result.ImageCount > 0 {
// 图片生成计费
var groupConfig *ImagePriceConfig
if apiKey.Group != nil {
groupConfig = &ImagePriceConfig{
Price1K: apiKey.Group.ImagePrice1K,
Price2K: apiKey.Group.ImagePrice2K,
Price4K: apiKey.Group.ImagePrice4K,
}
}
cost = s.billingService.CalculateImageCost(result.Model, result.ImageSize, result.ImageCount, groupConfig, multiplier)
} else {
// Token 计费
tokens := UsageTokens{
InputTokens: result.Usage.InputTokens,
OutputTokens: result.Usage.OutputTokens,
CacheCreationTokens: result.Usage.CacheCreationInputTokens,
CacheReadTokens: result.Usage.CacheReadInputTokens,
}
var err error
cost, err = s.billingService.CalculateCost(result.Model, tokens, multiplier)
if err != nil {
log.Printf("Calculate cost failed: %v", err)
cost = &CostBreakdown{ActualCost: 0}
}
}
// 判断计费方式：订阅模式 vs 余额模式
@@ -2039,6 +2058,10 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
// 创建使用日志
durationMs := int(result.Duration.Milliseconds())
var imageSize *string
if result.ImageSize != "" {
imageSize = &result.ImageSize
}
usageLog := &UsageLog{
UserID: user.ID,
APIKeyID: apiKey.ID,
@@ -2060,6 +2083,8 @@ func (s *GatewayService) RecordUsage(ctx context.Context, input *RecordUsageInpu
Stream: result.Stream,
DurationMs: &durationMs,
FirstTokenMs: result.FirstTokenMs,
ImageCount: result.ImageCount,
ImageSize: imageSize,
CreatedAt: time.Now(),
}

21
backend/internal/service/group.go
View File

@@ -17,6 +17,11 @@ type Group struct {
MonthlyLimitUSD *float64
DefaultValidityDays int
// 图片生成计费配置（antigravity 和 gemini 平台使用）
ImagePrice1K *float64
ImagePrice2K *float64
ImagePrice4K *float64
CreatedAt time.Time
UpdatedAt time.Time
@@ -47,3 +52,19 @@ func (g *Group) HasWeeklyLimit() bool {
func (g *Group) HasMonthlyLimit() bool {
return g.MonthlyLimitUSD != nil && *g.MonthlyLimitUSD > 0
}
// GetImagePrice 根据 image_size 返回对应的图片生成价格
// 如果分组未配置价格，返回 nil（调用方应使用默认值）
func (g *Group) GetImagePrice(imageSize string) *float64 {
switch imageSize {
case "1K":
return g.ImagePrice1K
case "2K":
return g.ImagePrice2K
case "4K":
return g.ImagePrice4K
default:
// 未知尺寸默认按 2K 计费
return g.ImagePrice2K
}
}

92
backend/internal/service/group_test.go Normal file
View File

@@ -0,0 +1,92 @@
//go:build unit
package service
import (
"testing"
"github.com/stretchr/testify/require"
)
// TestGroup_GetImagePrice_1K 测试 1K 尺寸返回正确价格
func TestGroup_GetImagePrice_1K(t *testing.T) {
price := 0.10
group := &Group{
ImagePrice1K: &price,
}
result := group.GetImagePrice("1K")
require.NotNil(t, result)
require.InDelta(t, 0.10, *result, 0.0001)
}
// TestGroup_GetImagePrice_2K 测试 2K 尺寸返回正确价格
func TestGroup_GetImagePrice_2K(t *testing.T) {
price := 0.15
group := &Group{
ImagePrice2K: &price,
}
result := group.GetImagePrice("2K")
require.NotNil(t, result)
require.InDelta(t, 0.15, *result, 0.0001)
}
// TestGroup_GetImagePrice_4K 测试 4K 尺寸返回正确价格
func TestGroup_GetImagePrice_4K(t *testing.T) {
price := 0.30
group := &Group{
ImagePrice4K: &price,
}
result := group.GetImagePrice("4K")
require.NotNil(t, result)
require.InDelta(t, 0.30, *result, 0.0001)
}
// TestGroup_GetImagePrice_UnknownSize 测试未知尺寸回退 2K
func TestGroup_GetImagePrice_UnknownSize(t *testing.T) {
price2K := 0.15
group := &Group{
ImagePrice2K: &price2K,
}
// 未知尺寸 "3K" 应该回退到 2K
result := group.GetImagePrice("3K")
require.NotNil(t, result)
require.InDelta(t, 0.15, *result, 0.0001)
// 空字符串也回退到 2K
result = group.GetImagePrice("")
require.NotNil(t, result)
require.InDelta(t, 0.15, *result, 0.0001)
}
// TestGroup_GetImagePrice_NilValues 测试未配置时返回 nil
func TestGroup_GetImagePrice_NilValues(t *testing.T) {
group := &Group{
// 所有 ImagePrice 字段都是 nil
}
require.Nil(t, group.GetImagePrice("1K"))
require.Nil(t, group.GetImagePrice("2K"))
require.Nil(t, group.GetImagePrice("4K"))
require.Nil(t, group.GetImagePrice("unknown"))
}
// TestGroup_GetImagePrice_PartialConfig 测试部分配置
func TestGroup_GetImagePrice_PartialConfig(t *testing.T) {
price1K := 0.10
group := &Group{
ImagePrice1K: &price1K,
// ImagePrice2K 和 ImagePrice4K 未配置
}
result := group.GetImagePrice("1K")
require.NotNil(t, result)
require.InDelta(t, 0.10, *result, 0.0001)
// 2K 和 4K 返回 nil
require.Nil(t, group.GetImagePrice("2K"))
require.Nil(t, group.GetImagePrice("4K"))
}

5
backend/internal/service/pricing_service.go
View File

@@ -34,6 +34,7 @@ type LiteLLMModelPricing struct {
LiteLLMProvider string `json:"litellm_provider"`
Mode string `json:"mode"`
SupportsPromptCaching bool `json:"supports_prompt_caching"`
OutputCostPerImage float64 `json:"output_cost_per_image"` // 图片生成模型每张图片价格
}
// PricingRemoteClient 远程价格数据获取接口
@@ -51,6 +52,7 @@ type LiteLLMRawEntry struct {
LiteLLMProvider string `json:"litellm_provider"`
Mode string `json:"mode"`
SupportsPromptCaching bool `json:"supports_prompt_caching"`
OutputCostPerImage *float64 `json:"output_cost_per_image"`
}
// PricingService 动态价格服务
@@ -319,6 +321,9 @@ func (s *PricingService) parsePricingData(body []byte) (map[string]*LiteLLMModel
if entry.CacheReadInputTokenCost != nil {
pricing.CacheReadInputTokenCost = *entry.CacheReadInputTokenCost
}
if entry.OutputCostPerImage != nil {
pricing.OutputCostPerImage = *entry.OutputCostPerImage
}
result[modelName] = pricing
}

4
backend/internal/service/usage_log.go
View File

@@ -39,6 +39,10 @@ type UsageLog struct {
DurationMs *int
FirstTokenMs *int
// 图片生成字段
ImageCount int
ImageSize *string
CreatedAt time.Time
User *User