Skip to content

Commit a3ca301

Browse files
luojun96luojun96
committed
add modal price guard for multimodal apis in aigateway
1 parent b7cb412 commit a3ca301

8 files changed

Lines changed: 310 additions & 0 deletions

File tree

‎aigateway/handler/chat_trace_test.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,38 @@ func TestPreflightTraceRecordsEarlyError(t *testing.T) {
118118
requireSpanAttrValue(t, spans[0].Attributes, "error.category", "gateway_error")
119119
}
120120

121+
func TestPreflightTraceRecordsModalPriceGuardError(t *testing.T) {
122+
exporter := tracetest.NewInMemoryExporter()
123+
provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))
124+
oldProvider := otel.GetTracerProvider()
125+
otel.SetTracerProvider(provider)
126+
t.Cleanup(func() {
127+
otel.SetTracerProvider(oldProvider)
128+
_ = provider.Shutdown(context.Background())
129+
})
130+
131+
_, preflight := startPreflightTrace(context.Background(), preflightTraceStart{
132+
API: "/v1/images/generations",
133+
RequestID: "req-1",
134+
UserID: "user-1",
135+
})
136+
preflight.SetTargetModel("model1", &resolvedModelTarget{
137+
Model: &types.Model{BaseModel: types.BaseModel{ID: "internal-model"}},
138+
ModelName: "resolved-model",
139+
})
140+
preflight.RecordError(newInvalidRequestModelTargetError(
141+
"invalid_request_error",
142+
"resolution exceeds the maximum supported resolution",
143+
modelTargetErrorOptions{},
144+
), "modal_price_guard")
145+
146+
spans := exporter.GetSpans()
147+
require.Len(t, spans, 1)
148+
require.Equal(t, "Error", spans[0].Status.Code.String())
149+
requireSpanAttrValue(t, spans[0].Attributes, "error.type", "modal_price_guard")
150+
requireSpanAttrValue(t, spans[0].Attributes, "aigateway.request.model", "model1")
151+
}
152+
121153
func TestPreflightTraceRecordsResolvedTargetModel(t *testing.T) {
122154
exporter := tracetest.NewInMemoryExporter()
123155
provider := sdktrace.NewTracerProvider(sdktrace.WithSyncer(exporter))

‎aigateway/handler/model_price_guard_ce.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,3 +7,7 @@ import "opencsg.com/csghub-server/aigateway/types"
77
func modelSKUPriceStatus(model *types.Model) (requiresSKUPrice bool, hasConfiguredSKUPrice bool) {
88
return false, false
99
}
10+
11+
func checkModalRequestAllowed(model *types.Model, size string) *modelTargetError {
12+
return nil
13+
}

‎aigateway/handler/openai_image.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,12 @@ func (h *OpenAIHandlerImpl) GenerateImage(c *gin.Context) {
7070
return
7171
}
7272
preflight.SetTargetModel(modelID, modelTarget)
73+
74+
if guardErr := checkModalRequestAllowed(modelTarget.Model, string(req.Size)); guardErr != nil {
75+
preflight.RecordError(guardErr, "modal_price_guard")
76+
handleModelTargetError(c, ctx, modelID, "modal price guard rejected request", guardErr)
77+
return
78+
}
7379
preflight.End()
7480

7581
traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

‎aigateway/handler/openai_image_edit.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -84,6 +84,12 @@ func (h *OpenAIHandlerImpl) EditImage(c *gin.Context) {
8484
return
8585
}
8686
preflight.SetTargetModel(modelID, modelTarget)
87+
88+
if guardErr := checkModalRequestAllowed(modelTarget.Model, c.PostForm("size")); guardErr != nil {
89+
preflight.RecordError(guardErr, "modal_price_guard")
90+
handleModelTargetError(c, ctx, modelID, "modal price guard rejected request", guardErr)
91+
return
92+
}
8793
preflight.End()
8894

8995
traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

‎aigateway/handler/openai_video.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,12 @@ func (h *OpenAIHandlerImpl) CreateVideo(c *gin.Context) {
7373
return
7474
}
7575
preflight.SetTargetModel(input.modelID, modelTarget)
76+
77+
if guardErr := checkModalRequestAllowed(modelTarget.Model, input.adapterReq.Size); guardErr != nil {
78+
preflight.RecordError(guardErr, "modal_price_guard")
79+
handleModelTargetError(c, ctx, input.modelID, "modal price guard rejected request", guardErr)
80+
return
81+
}
7682
preflight.End()
7783

7884
traceCtx, generationRecorder := h.startModalGenerationTrace(ctx, modalTraceStartInput{

‎aigateway/types/openai.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,6 +313,7 @@ type ModelModalPrice struct {
313313
SkuUnit int64 `json:"sku_unit,omitempty"`
314314
SkuUnitType string `json:"sku_unit_type,omitempty"`
315315
Resolution string `json:"resolution,omitempty"`
316+
SkuDesc string `json:"sku_desc,omitempty"`
316317
}
317318

318319
// ModelScenePrice is Metadata["pricing"]: serverless and external_llm use token and modal generation prices.

‎common/utils/common/resolution.go‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
package common
2+
3+
import (
4+
"fmt"
5+
"regexp"
6+
"strconv"
7+
"strings"
8+
)
9+
10+
// This file parses resolution strings in preset ("720p"), "WIDTHxHEIGHT",
11+
// or bare-integer format. ExtractEventResolutionMaxSide is lenient (used for
12+
// historical/billing data), while ParseResolutionMaxSide additionally
13+
// requires a strictly positive value (used for real-time request/SKU
14+
// validation), so gateway and accounting never interpret the same string
15+
// differently.
16+
var (
17+
resolutionPFormat = regexp.MustCompile(`^(\d+)\s*[pP]$`)
18+
resolutionWxHFormat = regexp.MustCompile(`^(\d+)\s*[xX×]\s*(\d+)$`)
19+
resolutionPureFormat = regexp.MustCompile(`^(\d+)$`)
20+
)
21+
22+
func longEdge(width, height int64) int64 {
23+
if width > height {
24+
return width
25+
}
26+
return height
27+
}
28+
29+
func parseResolutionPFormat(s string) (int64, bool) {
30+
matches := resolutionPFormat.FindStringSubmatch(s)
31+
if matches == nil {
32+
return 0, false
33+
}
34+
v, err := strconv.ParseInt(matches[1], 10, 64)
35+
if err != nil {
36+
return 0, false
37+
}
38+
return v, true
39+
}
40+
41+
func parseResolutionWxHPair(s string) (width, height int64, ok bool) {
42+
matches := resolutionWxHFormat.FindStringSubmatch(s)
43+
if matches == nil {
44+
return 0, 0, false
45+
}
46+
w, err1 := strconv.ParseInt(matches[1], 10, 64)
47+
h, err2 := strconv.ParseInt(matches[2], 10, 64)
48+
if err1 != nil || err2 != nil {
49+
return 0, 0, false
50+
}
51+
return w, h, true
52+
}
53+
54+
func parseResolutionWxHFormat(s string) (int64, bool) {
55+
w, h, ok := parseResolutionWxHPair(s)
56+
if !ok {
57+
return 0, false
58+
}
59+
return longEdge(w, h), true
60+
}
61+
62+
func parseResolutionPureFormat(s string) (int64, bool) {
63+
matches := resolutionPureFormat.FindStringSubmatch(s)
64+
if matches == nil {
65+
return 0, false
66+
}
67+
v, err := strconv.ParseInt(matches[1], 10, 64)
68+
if err != nil {
69+
return 0, false
70+
}
71+
return v, true
72+
}
73+
74+
func ExtractEventResolutionMaxSide(eventResolution string) (int64, error) {
75+
s := strings.TrimSpace(eventResolution)
76+
if len(s) < 1 {
77+
return 0, nil
78+
}
79+
80+
if v, ok := parseResolutionPFormat(s); ok {
81+
return v, nil
82+
}
83+
if v, ok := parseResolutionWxHFormat(s); ok {
84+
return v, nil
85+
}
86+
if v, ok := parseResolutionPureFormat(s); ok {
87+
return v, nil
88+
}
89+
90+
return 0, fmt.Errorf("unsupported resolution format: %s, must be 1080P or 720 or 1920x1080", s)
91+
}
92+
93+
// ParseResolutionMaxSide is like ExtractEventResolutionMaxSide but rejects
94+
// empty input and non-positive results.
95+
func ParseResolutionMaxSide(size string) (int64, bool) {
96+
v, err := ExtractEventResolutionMaxSide(size)
97+
if err != nil || v <= 0 {
98+
return 0, false
99+
}
100+
return v, true
101+
}

‎common/utils/common/resolution_test.go‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
package common
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
)
8+
9+
func TestParseResolutionPFormat(t *testing.T) {
10+
tests := []struct {
11+
input string
12+
want int64
13+
ok bool
14+
}{
15+
{"720P", 720, true},
16+
{"720p", 720, true},
17+
{"1080P", 1080, true},
18+
{"1080p", 1080, true},
19+
{"480P", 480, true},
20+
{"720 P", 720, true},
21+
{"P", 0, false},
22+
{"abcP", 0, false},
23+
{"720PX", 0, false},
24+
{"", 0, false},
25+
}
26+
27+
for _, tt := range tests {
28+
got, ok := parseResolutionPFormat(tt.input)
29+
require.Equal(t, tt.ok, ok, "parseResolutionPFormat(%q) ok", tt.input)
30+
if ok {
31+
require.Equal(t, tt.want, got, "parseResolutionPFormat(%q) value", tt.input)
32+
}
33+
}
34+
}
35+
36+
func TestParseResolutionWxHFormat(t *testing.T) {
37+
tests := []struct {
38+
input string
39+
want int64
40+
ok bool
41+
}{
42+
{"720x1280", 1280, true},
43+
{"1280x720", 1280, true},
44+
{"1920x1080", 1920, true},
45+
{"1080X1920", 1920, true},
46+
{"1920×1080", 1920, true},
47+
{"720 x 1280", 1280, true},
48+
{"0x0", 0, true},
49+
{"720x", 0, false},
50+
{"x720", 0, false},
51+
{"720x1280x3", 0, false},
52+
{"abcx720", 0, false},
53+
{"", 0, false},
54+
}
55+
56+
for _, tt := range tests {
57+
got, ok := parseResolutionWxHFormat(tt.input)
58+
require.Equal(t, tt.ok, ok, "parseResolutionWxHFormat(%q) ok", tt.input)
59+
if ok {
60+
require.Equal(t, tt.want, got, "parseResolutionWxHFormat(%q) value", tt.input)
61+
}
62+
}
63+
}
64+
65+
func TestParseResolutionPureFormat(t *testing.T) {
66+
tests := []struct {
67+
input string
68+
want int64
69+
ok bool
70+
}{
71+
{"720", 720, true},
72+
{"1080", 1080, true},
73+
{"0", 0, true},
74+
{"1", 1, true},
75+
{"720P", 0, false},
76+
{"720x1280", 0, false},
77+
{"abc", 0, false},
78+
{"", 0, false},
79+
}
80+
81+
for _, tt := range tests {
82+
got, ok := parseResolutionPureFormat(tt.input)
83+
require.Equal(t, tt.ok, ok, "parseResolutionPureFormat(%q) ok", tt.input)
84+
if ok {
85+
require.Equal(t, tt.want, got, "parseResolutionPureFormat(%q) value", tt.input)
86+
}
87+
}
88+
}
89+
90+
func TestExtractEventResolutionMaxSide(t *testing.T) {
91+
tests := []struct {
92+
input string
93+
want int64
94+
err bool
95+
}{
96+
{"", 0, false},
97+
{" ", 0, false},
98+
{"720p", 720, false},
99+
{"1080P", 1080, false},
100+
{"720", 720, false},
101+
{"1080", 1080, false},
102+
{"720x1280", 1280, false},
103+
{"1920x1080", 1920, false},
104+
{"1920×1080", 1920, false},
105+
{"invalid", 0, true},
106+
{"abcxdef", 0, true},
107+
}
108+
109+
for _, tt := range tests {
110+
got, err := ExtractEventResolutionMaxSide(tt.input)
111+
if tt.err {
112+
require.Error(t, err, "extractEventResolutionMaxSide(%q)", tt.input)
113+
} else {
114+
require.NoError(t, err, "extractEventResolutionMaxSide(%q)", tt.input)
115+
require.Equal(t, tt.want, got, "extractEventResolutionMaxSide(%q) value", tt.input)
116+
}
117+
}
118+
}
119+
120+
func TestParseResolutionMaxSide(t *testing.T) {
121+
tests := []struct {
122+
name string
123+
size string
124+
wantSide int64
125+
wantOK bool
126+
}{
127+
{"square", "1024x1024", 1024, true},
128+
{"landscape", "1280x720", 1280, true},
129+
{"portrait", "720x1280", 1280, true},
130+
{"whitespace wxh", " 1024 x 1536 ", 1536, true},
131+
{"uppercase wxh", "1024X1024", 1024, true},
132+
{"preset uppercase", "1080P", 1080, true},
133+
{"preset lowercase", "720p", 720, true},
134+
{"whitespace preset", " 1080p ", 1080, true},
135+
{"pure integer", "1024", 1024, true},
136+
{"whitespace integer", " 2048 ", 2048, true},
137+
{"empty", "", 0, false},
138+
{"degenerate wxh one side zero", "0x1024", 1024, true}, // matches ExtractEventResolutionMaxSide leniency for a single zero dimension
139+
{"wxh both sides zero", "0x0", 0, false},
140+
{"negative wxh", "-1024x1024", 0, false},
141+
{"zero preset", "0p", 0, false},
142+
{"zero integer", "0", 0, false},
143+
{"negative integer", "-1024", 0, false},
144+
{"non numeric", "abcxdef", 0, false},
145+
}
146+
147+
for _, tt := range tests {
148+
t.Run(tt.name, func(t *testing.T) {
149+
side, ok := ParseResolutionMaxSide(tt.size)
150+
require.Equal(t, tt.wantOK, ok)
151+
require.Equal(t, tt.wantSide, side)
152+
})
153+
}
154+
}

0 commit comments

Comments
 (0)