feat: add gateway billing estimate and rate limit details

2026-05-15 01:53:52 +08:00
parent bdc9be63d5
commit 37d0f919e5
14 changed files with 916 additions and 66 deletions
@@ -0,0 +1,132 @@
+package runner
+
+import (
+	"context"
+	"testing"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+func TestImageBillingEstimateUsesCountResolutionAndQuality(t *testing.T) {
+	service := &Service{}
+	candidate := store.RuntimeModelCandidate{
+		ModelName: "image-model",
+		BaseBillingConfig: map[string]any{
+			"image": map[string]any{
+				"basePrice": 10,
+				"dynamicWeight": map[string]any{
+					"resolutionFactors": map[string]any{"2K": 1.5},
+					"qualityFactors":    map[string]any{"high": 1.5},
+				},
+			},
+		},
+	}
+
+	items := service.billings(context.Background(), nil, "images.generations", map[string]any{
+		"count":      2,
+		"quality":    "high",
+		"resolution": "2K",
+	}, candidate, clients.Response{}, true)
+
+	line := firstBillingLine(t, items)
+	if got, want := floatFromAny(line["amount"]), 45.0; got != want {
+		t.Fatalf("image estimated amount = %v, want %v", got, want)
+	}
+	if got, want := line["quantity"], 2; got != want {
+		t.Fatalf("image quantity = %v, want %v", got, want)
+	}
+}
+
+func TestVideoBillingEstimateUsesFiveSecondUnitsAndDynamicWeights(t *testing.T) {
+	service := &Service{}
+	candidate := store.RuntimeModelCandidate{
+		ModelName: "video-model",
+		BaseBillingConfig: map[string]any{
+			"video": map[string]any{
+				"basePrice": 100,
+				"dynamicWeight": map[string]any{
+					"resolutionWeights":        map[string]any{"1080p": 1.5},
+					"audioWeights":             map[string]any{"true": 2},
+					"referenceVideoWeights":    map[string]any{"true": 1.5},
+					"voiceSpecifiedWeights":    map[string]any{"true": 1.2},
+					"unusedCompatibilityField": map[string]any{"true": 99},
+				},
+			},
+		},
+	}
+
+	items := service.billings(context.Background(), nil, "videos.generations", map[string]any{
+		"audio":      true,
+		"duration":   12,
+		"resolution": "1080p",
+		"voice_id":   "voice-a",
+		"content": []any{
+			map[string]any{"type": "video_url", "video_url": map[string]any{"url": "https://example.com/reference.mp4"}},
+		},
+	}, candidate, clients.Response{}, true)
+
+	line := firstBillingLine(t, items)
+	if got, want := floatFromAny(line["amount"]), 1620.0; got != want {
+		t.Fatalf("video estimated amount = %v, want %v", got, want)
+	}
+	if got, want := floatFromAny(line["durationUnitCount"]), 3.0; got != want {
+		t.Fatalf("video duration units = %v, want %v", got, want)
+	}
+	if got, want := line["quantity"], 3; got != want {
+		t.Fatalf("video quantity = %v, want %v", got, want)
+	}
+}
+
+func TestVideoBillingEstimateSupportsServerMainStyleDynamicKeys(t *testing.T) {
+	service := &Service{}
+	candidate := store.RuntimeModelCandidate{
+		ModelName: "legacy-video-model",
+		BaseBillingConfig: map[string]any{
+			"videoBase": 100,
+			"video": map[string]any{
+				"dynamicWeight": map[string]any{
+					"720p":                 1.25,
+					"audio-true":           2,
+					"reference-video-true": 1.5,
+				},
+			},
+		},
+	}
+
+	items := service.billings(context.Background(), nil, "videos.generations", map[string]any{
+		"audio":      "true",
+		"duration":   5,
+		"resolution": "720p",
+		"video_list": []any{map[string]any{"url": "https://example.com/reference.mp4"}},
+	}, candidate, clients.Response{}, true)
+
+	line := firstBillingLine(t, items)
+	if got, want := floatFromAny(line["amount"]), 375.0; got != want {
+		t.Fatalf("legacy video estimated amount = %v, want %v", got, want)
+	}
+}
+
+func TestVideoDurationEstimateSumsMultiShotDurations(t *testing.T) {
+	duration := requestDurationSeconds(map[string]any{
+		"multi_prompt": []any{
+			map[string]any{"prompt": "shot 1", "duration": 3},
+			map[string]any{"prompt": "shot 2", "duration": 7},
+		},
+	})
+	if duration != 10 {
+		t.Fatalf("multi-shot duration = %v, want 10", duration)
+	}
+}
+
+func firstBillingLine(t *testing.T, items []any) map[string]any {
+	t.Helper()
+	if len(items) != 1 {
+		t.Fatalf("items length = %d, want 1: %+v", len(items), items)
+	}
+	line, ok := items[0].(map[string]any)
+	if !ok {
+		t.Fatalf("item type = %T, want map[string]any", items[0])
+	}
+	return line
+}