feat(gateway): 补齐桌面端高级媒体直连接口
ci / verify (pull_request) Successful in 15m34s
ci / verify (pull_request) Successful in 15m34s
新增图片矢量化、视频超分、每日用量、计价与任务隔离能力,并通过环境变量解析平台凭据。 已通过 Go 全量门禁、迁移检查、镜像构建以及 Vectorizer 五格式和 Topaz 3 秒视频真实 DEV 验收。
This commit is contained in:
@@ -486,8 +486,12 @@ func modelTypeAliases(value string) []string {
|
||||
return []string{"image_generate"}
|
||||
case "images.edits":
|
||||
return []string{"image_edit"}
|
||||
case "images.vectorize", "vectorize":
|
||||
return []string{"image_vectorize"}
|
||||
case "video", "videos.generations":
|
||||
return []string{"video_generate"}
|
||||
case "videos.upscales", "video_upscale":
|
||||
return []string{"video_enhance"}
|
||||
case "omni_video":
|
||||
return []string{"video_generate", "image_to_video", "omni_video"}
|
||||
case "song", "music", "song.generations", "music.generations", "music_generate":
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
)
|
||||
|
||||
type DailyTokenUsage struct {
|
||||
Date string `json:"date"`
|
||||
TotalTokens int64 `json:"totalTokens"`
|
||||
InputTokens int64 `json:"inputTokens"`
|
||||
OutputTokens int64 `json:"outputTokens"`
|
||||
CachedInputTokens int64 `json:"cachedInputTokens"`
|
||||
ResourcePoints float64 `json:"resourcePoints"`
|
||||
TaskCount int64 `json:"taskCount"`
|
||||
}
|
||||
|
||||
func (s *Store) ListDailyTokenUsage(ctx context.Context, user *auth.User, from, to time.Time, timezone string) ([]DailyTokenUsage, error) {
|
||||
gatewayUserID := localGatewayUserID(user)
|
||||
apiKeyID := ""
|
||||
userID := ""
|
||||
if user != nil {
|
||||
apiKeyID = strings.TrimSpace(user.APIKeyID)
|
||||
userID = strings.TrimSpace(user.ID)
|
||||
}
|
||||
if gatewayUserID == "" && userID == "" {
|
||||
return nil, ErrLocalUserRequired
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT to_char(created_at AT TIME ZONE $6, 'YYYY-MM-DD') AS usage_day,
|
||||
COALESCE(sum(
|
||||
CASE WHEN jsonb_typeof(usage->'totalTokens') = 'number' THEN (usage->>'totalTokens')::numeric
|
||||
WHEN jsonb_typeof(usage->'total_tokens') = 'number' THEN (usage->>'total_tokens')::numeric
|
||||
ELSE COALESCE(CASE WHEN jsonb_typeof(usage->'inputTokens') = 'number' THEN (usage->>'inputTokens')::numeric ELSE 0 END, 0)
|
||||
+ COALESCE(CASE WHEN jsonb_typeof(usage->'outputTokens') = 'number' THEN (usage->>'outputTokens')::numeric ELSE 0 END, 0)
|
||||
END
|
||||
), 0)::bigint AS total_tokens,
|
||||
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'inputTokens') = 'number' THEN (usage->>'inputTokens')::numeric WHEN jsonb_typeof(usage->'input_tokens') = 'number' THEN (usage->>'input_tokens')::numeric ELSE 0 END), 0)::bigint,
|
||||
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'outputTokens') = 'number' THEN (usage->>'outputTokens')::numeric WHEN jsonb_typeof(usage->'output_tokens') = 'number' THEN (usage->>'output_tokens')::numeric ELSE 0 END), 0)::bigint,
|
||||
COALESCE(sum(CASE WHEN jsonb_typeof(usage->'cachedInputTokens') = 'number' THEN (usage->>'cachedInputTokens')::numeric WHEN jsonb_typeof(usage->'cached_input_tokens') = 'number' THEN (usage->>'cached_input_tokens')::numeric ELSE 0 END), 0)::bigint,
|
||||
COALESCE(sum(final_charge_amount), 0)::float8,
|
||||
count(*)::bigint
|
||||
FROM gateway_tasks
|
||||
WHERE status = 'succeeded'
|
||||
AND ((NULLIF($1, '')::uuid IS NOT NULL AND gateway_user_id = NULLIF($1, '')::uuid)
|
||||
OR (NULLIF($1, '')::uuid IS NULL AND NULLIF($2, '') IS NOT NULL AND user_id = $2))
|
||||
AND (NULLIF($3, '') IS NULL OR api_key_id = $3)
|
||||
AND created_at >= ($4::date::timestamp AT TIME ZONE $6)
|
||||
AND created_at < (($5::date + 1)::timestamp AT TIME ZONE $6)
|
||||
GROUP BY usage_day
|
||||
ORDER BY usage_day`, gatewayUserID, userID, apiKeyID, from.Format("2006-01-02"), to.Format("2006-01-02"), timezone)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := make([]DailyTokenUsage, 0)
|
||||
for rows.Next() {
|
||||
var item DailyTokenUsage
|
||||
if err := rows.Scan(&item.Date, &item.TotalTokens, &item.InputTokens, &item.OutputTokens, &item.CachedInputTokens, &item.ResourcePoints, &item.TaskCount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, item)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
)
|
||||
|
||||
func TestListDailyTokenUsageScopesAPIKeyAndAppliesIANATimezone(t *testing.T) {
|
||||
databaseURL := strings.TrimSpace(os.Getenv("AI_GATEWAY_TEST_DATABASE_URL"))
|
||||
if databaseURL == "" {
|
||||
t.Skip("set AI_GATEWAY_TEST_DATABASE_URL to run daily usage PostgreSQL integration tests")
|
||||
}
|
||||
ctx := context.Background()
|
||||
db, err := Connect(ctx, databaseURL)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
var databaseName string
|
||||
if err := db.pool.QueryRow(ctx, `SELECT current_database()`).Scan(&databaseName); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !strings.Contains(strings.ToLower(databaseName), "test") {
|
||||
t.Fatalf("refusing to use non-test database %q", databaseName)
|
||||
}
|
||||
|
||||
suffix := fmt.Sprint(time.Now().UnixNano())
|
||||
var userID, otherUserID string
|
||||
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-user-"+suffix, "daily-user-"+suffix).Scan(&userID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.pool.QueryRow(ctx, `INSERT INTO gateway_users (user_key, username) VALUES ($1, $2) RETURNING id::text`, "daily-other-"+suffix, "daily-other-"+suffix).Scan(&otherUserID); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_tasks WHERE gateway_user_id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
|
||||
_, _ = db.pool.Exec(context.Background(), `DELETE FROM gateway_users WHERE id IN ($1::uuid, $2::uuid)`, userID, otherUserID)
|
||||
})
|
||||
insertTask := func(ownerID, keyID, createdAt string, totalTokens int, points float64) {
|
||||
t.Helper()
|
||||
_, err := db.pool.Exec(ctx, `
|
||||
INSERT INTO gateway_tasks (kind, user_id, gateway_user_id, api_key_id, model, status, usage, final_charge_amount, created_at, finished_at)
|
||||
VALUES ('chat.completions', $1::text, ($1::text)::uuid, $2::text, 'daily-test-model', 'succeeded', jsonb_build_object('totalTokens', $3::bigint, 'inputTokens', $3::bigint - 1, 'outputTokens', 1), $4::numeric, $5::timestamptz, $5::timestamptz)`, ownerID, keyID, totalTokens, points, createdAt)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
insertTask(userID, "daily-key-a-"+suffix, "2026-07-20T15:30:00Z", 10, 1)
|
||||
insertTask(userID, "daily-key-a-"+suffix, "2026-07-21T16:30:00Z", 20, 2)
|
||||
insertTask(userID, "daily-key-b-"+suffix, "2026-07-20T16:30:00Z", 100, 10)
|
||||
insertTask(otherUserID, "daily-key-a-"+suffix, "2026-07-20T16:30:00Z", 999, 99)
|
||||
|
||||
location, err := time.LoadLocation("Asia/Shanghai")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
from := time.Date(2026, 7, 20, 0, 0, 0, 0, location)
|
||||
to := time.Date(2026, 7, 22, 0, 0, 0, 0, location)
|
||||
keyItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID, APIKeyID: "daily-key-a-" + suffix}, from, to, "Asia/Shanghai")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(keyItems) != 2 || keyItems[0].Date != "2026-07-20" || keyItems[0].TotalTokens != 10 || keyItems[1].Date != "2026-07-22" || keyItems[1].TotalTokens != 20 {
|
||||
t.Fatalf("API key usage leaked another key/user or ignored timezone: %+v", keyItems)
|
||||
}
|
||||
userItems, err := db.ListDailyTokenUsage(ctx, &auth.User{GatewayUserID: userID}, from, to, "Asia/Shanghai")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(userItems) != 3 || userItems[1].Date != "2026-07-21" || userItems[1].TotalTokens != 100 {
|
||||
t.Fatalf("JWT user usage did not include all owned API keys: %+v", userItems)
|
||||
}
|
||||
}
|
||||
@@ -54,9 +54,13 @@ func billingResourcesForModelTypes(modelTypes []string) map[string]bool {
|
||||
resources["image"] = true
|
||||
case "images.edits", "image_edit":
|
||||
resources["image_edit"] = true
|
||||
case "image_vectorize", "images.vectorize":
|
||||
resources["image_vectorize"] = true
|
||||
case "video", "videos.generations", "video_generate", "image_to_video", "text_to_video",
|
||||
"video_edit", "omni_video", "video_reference", "video_first_last_frame":
|
||||
resources["video"] = true
|
||||
case "video_enhance", "videos.upscales", "video_upscale":
|
||||
resources["video_enhance"] = true
|
||||
case "audio", "text_to_speech", "speech", "voice_clone":
|
||||
resources["audio"] = true
|
||||
case "music", "music_generate", "audio_generate":
|
||||
@@ -100,8 +104,12 @@ func billingConfigKeyAllowed(key string, resources map[string]bool) bool {
|
||||
return resources["image"]
|
||||
case "image_edit", "imageedit", "editbase":
|
||||
return resources["image_edit"]
|
||||
case "image_vectorize", "imagevectorize", "vectorizebase":
|
||||
return resources["image_vectorize"]
|
||||
case "video", "videobase":
|
||||
return resources["video"]
|
||||
case "video_enhance", "videoenhance", "videoenhancebase":
|
||||
return resources["video_enhance"]
|
||||
case "audio", "audiobase":
|
||||
return resources["audio"]
|
||||
case "music", "musicbase":
|
||||
|
||||
@@ -29,7 +29,7 @@ const (
|
||||
)
|
||||
|
||||
func defaultAPIKeyScopes() []string {
|
||||
return []string{"chat", "embedding", "rerank", "image", "video", "music", "audio", "voice_clone"}
|
||||
return []string{"chat", "embedding", "rerank", "image", "image_vectorize", "video", "video_enhance", "music", "audio", "voice_clone"}
|
||||
}
|
||||
|
||||
func normalizeAPIKeyScopes(scopes []string) []string {
|
||||
@@ -517,7 +517,7 @@ type GatewayTask struct {
|
||||
Message string `json:"message,omitempty"`
|
||||
AttemptCount int `json:"attemptCount"`
|
||||
RemoteTaskID string `json:"remoteTaskId,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"remoteTaskPayload,omitempty"`
|
||||
RemoteTaskPayload map[string]any `json:"-"`
|
||||
Result map[string]any `json:"result,omitempty"`
|
||||
Billings []any `json:"billings,omitempty"`
|
||||
Usage map[string]any `json:"usage"`
|
||||
|
||||
@@ -331,7 +331,7 @@ ORDER BY priority ASC, resource_type ASC, rule_key ASC`, id)
|
||||
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s currency does not match its rule set", ruleKey)
|
||||
}
|
||||
switch calculatorType {
|
||||
case "token_usage", "unit_weight", "duration_weight":
|
||||
case "token_usage", "unit_weight", "duration_weight", "transition_matrix":
|
||||
default:
|
||||
return EffectivePricingConfig{}, fmt.Errorf("pricing rule %s uses unsupported calculator %s", ruleKey, calculatorType)
|
||||
}
|
||||
@@ -386,8 +386,12 @@ func ValidateEffectivePricingRuleShape(resourceType string, unit string, calcula
|
||||
allowed = unit == "1k_tokens" && calculatorType == "token_usage"
|
||||
case "image", "image_edit":
|
||||
allowed = unit == "image" && calculatorType == "unit_weight"
|
||||
case "image_vectorize":
|
||||
allowed = unit == "conversion" && calculatorType == "unit_weight"
|
||||
case "video":
|
||||
allowed = unit == "5s" && calculatorType == "duration_weight"
|
||||
case "video_enhance":
|
||||
allowed = unit == "5s" && calculatorType == "transition_matrix"
|
||||
case "music":
|
||||
allowed = (unit == "song" || unit == "item") && calculatorType == "unit_weight"
|
||||
case "audio":
|
||||
@@ -415,6 +419,8 @@ func DefaultEffectivePricingCalculator(resourceType string) string {
|
||||
return "token_usage"
|
||||
case "video":
|
||||
return "duration_weight"
|
||||
case "video_enhance":
|
||||
return "transition_matrix"
|
||||
default:
|
||||
return "unit_weight"
|
||||
}
|
||||
|
||||
@@ -13,6 +13,8 @@ func TestValidateEffectivePricingRuleShape(t *testing.T) {
|
||||
{name: "text", resource: "text_input", unit: "1k_tokens", calculator: "token_usage", valid: true},
|
||||
{name: "image", resource: "image", unit: "image", calculator: "unit_weight", valid: true},
|
||||
{name: "video", resource: "video", unit: "5s", calculator: "duration_weight", valid: true},
|
||||
{name: "video enhance", resource: "video_enhance", unit: "5s", calculator: "transition_matrix", valid: true},
|
||||
{name: "image vectorize", resource: "image_vectorize", unit: "conversion", calculator: "unit_weight", valid: true},
|
||||
{name: "wrong video unit", resource: "video", unit: "second", calculator: "duration_weight"},
|
||||
{name: "wrong image calculator", resource: "image", unit: "image", calculator: "token_usage"},
|
||||
}
|
||||
|
||||
@@ -0,0 +1,25 @@
|
||||
package store
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestGatewayTaskJSONOmitsPrivateRemoteTaskPayload(t *testing.T) {
|
||||
raw, err := json.Marshal(GatewayTask{
|
||||
ID: "task-private-state",
|
||||
RemoteTaskID: "safe-upstream-id",
|
||||
RemoteTaskPayload: map[string]any{"imageToken": "private-image-token", "receipt": "private-receipt"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
serialized := strings.ToLower(string(raw))
|
||||
if strings.Contains(serialized, "private-image-token") || strings.Contains(serialized, "private-receipt") || strings.Contains(serialized, "remotetaskpayload") {
|
||||
t.Fatalf("private provider recovery state leaked into task JSON: %s", serialized)
|
||||
}
|
||||
if !strings.Contains(serialized, "safe-upstream-id") {
|
||||
t.Fatalf("public upstream task id should remain available: %s", serialized)
|
||||
}
|
||||
}
|
||||
@@ -528,13 +528,11 @@ WHERE id = $1::uuid
|
||||
}
|
||||
_, err = tx.Exec(ctx, `
|
||||
UPDATE gateway_task_attempts
|
||||
SET remote_task_id = NULLIF($2::text, ''),
|
||||
response_snapshot = COALESCE(response_snapshot, '{}'::jsonb) || jsonb_build_object('remote_task_payload', $3::jsonb)
|
||||
SET remote_task_id = NULLIF($2::text, '')
|
||||
WHERE id = $1::uuid
|
||||
AND status = 'running'`,
|
||||
attemptID,
|
||||
remoteTaskID,
|
||||
string(payloadJSON),
|
||||
)
|
||||
return err
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user