From ae1e68792a367e2d7e2235080269e507b4d687cb Mon Sep 17 00:00:00 2001 From: fengsilin Date: Wed, 17 Jun 2026 17:35:04 +0800 Subject: [PATCH] feat(task-pricing): resolve pricing dimensions from task requests --- relay/helper/derived_funcs.go | 126 ++++++++++++ relay/helper/derived_funcs_test.go | 41 ++++ relay/helper/dimension_resolver.go | 251 ++++++++++++++++++++++++ relay/helper/dimension_resolver_test.go | 127 ++++++++++++ 4 files changed, 545 insertions(+) create mode 100644 relay/helper/derived_funcs.go create mode 100644 relay/helper/derived_funcs_test.go create mode 100644 relay/helper/dimension_resolver.go create mode 100644 relay/helper/dimension_resolver_test.go diff --git a/relay/helper/derived_funcs.go b/relay/helper/derived_funcs.go new file mode 100644 index 0000000..7cc1c8b --- /dev/null +++ b/relay/helper/derived_funcs.go @@ -0,0 +1,126 @@ +package helper + +import ( + "fmt" + "math" + "strconv" + "strings" +) + +func init() { + registerDerivedPricingFunc("megapixels", derivedMegapixels) + registerDerivedPricingFunc("aspect_ratio", derivedAspectRatio) + registerDerivedPricingFunc("input_mode", derivedInputMode) +} + +func derivedMegapixels(req, resp map[string]any) (any, error) { + resolution, ok, _ := getNestedValue(req, "resolution") + if !ok || !hasNonEmptyValue(resolution) { + resolution, ok, _ = getNestedValue(req, "options.size") + } + if !ok || !hasNonEmptyValue(resolution) { + return nil, nil + } + width, height, err := parseResolution(fmt.Sprintf("%v", resolution)) + if err != nil { + return nil, err + } + return float64(width*height) / 1_000_000, nil +} + +func derivedAspectRatio(req, resp map[string]any) (any, error) { + resolution, ok, _ := getNestedValue(req, "resolution") + if !ok || !hasNonEmptyValue(resolution) { + resolution, ok, _ = getNestedValue(req, "options.size") + } + if !ok || !hasNonEmptyValue(resolution) { + return nil, nil + } + width, height, err := parseResolution(fmt.Sprintf("%v", resolution)) + if err != nil { + return nil, err + } + if height == 0 { + return nil, fmt.Errorf("resolution height must be greater than zero") + } + gcd := intGCD(width, height) + return fmt.Sprintf("%d:%d", width/gcd, height/gcd), nil +} + +func derivedInputMode(req, resp map[string]any) (any, error) { + if hasAnyNonEmpty(req, "video", "video_url", "reference_video", "first_frame_video") { + return "v2v", nil + } + if hasAnyNonEmpty(req, "image", "image_url", "image[]", "images", "first_frame_image", "last_frame_image") { + return "i2v", nil + } + return "t2v", nil +} + +func hasNonEmptyValue(v any) bool { + switch value := v.(type) { + case nil: + return false + case string: + return strings.TrimSpace(value) != "" + case []any: + return len(value) > 0 + case []string: + return len(value) > 0 + case map[string]any: + return len(value) > 0 + default: + return true + } +} + +func hasAnyNonEmpty(req map[string]any, keys ...string) bool { + for _, key := range keys { + value, ok, _ := getNestedValue(req, key) + if ok && hasNonEmptyValue(value) { + return true + } + } + return false +} + +func parseResolution(raw string) (int, int, error) { + normalized := strings.ToLower(strings.TrimSpace(raw)) + if strings.HasSuffix(normalized, "p") { + height, err := strconv.Atoi(strings.TrimSuffix(normalized, "p")) + if err != nil || height <= 0 { + return 0, 0, fmt.Errorf("invalid resolution %q", raw) + } + width := int(math.Round(float64(height) * 16 / 9)) + return width, height, nil + } + + parts := strings.FieldsFunc(normalized, func(r rune) bool { + return r == 'x' || r == '*' || r == '×' + }) + if len(parts) != 2 { + return 0, 0, fmt.Errorf("invalid resolution %q", raw) + } + width, err := strconv.Atoi(strings.TrimSpace(parts[0])) + if err != nil || width <= 0 { + return 0, 0, fmt.Errorf("invalid resolution width %q", raw) + } + height, err := strconv.Atoi(strings.TrimSpace(parts[1])) + if err != nil || height <= 0 { + return 0, 0, fmt.Errorf("invalid resolution height %q", raw) + } + return width, height, nil +} + +func intGCD(a, b int) int { + for b != 0 { + a, b = b, a%b + } + if a < 0 { + return -a + } + if a == 0 { + return 1 + } + return a +} diff --git a/relay/helper/derived_funcs_test.go b/relay/helper/derived_funcs_test.go new file mode 100644 index 0000000..fad16de --- /dev/null +++ b/relay/helper/derived_funcs_test.go @@ -0,0 +1,41 @@ +package helper + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDerivedInputMode(t *testing.T) { + mode, err := derivedInputMode(map[string]any{}, nil) + require.NoError(t, err) + require.Equal(t, "t2v", mode) + + mode, err = derivedInputMode(map[string]any{"image": "frame.png"}, nil) + require.NoError(t, err) + require.Equal(t, "i2v", mode) + + mode, err = derivedInputMode(map[string]any{"image": "frame.png", "video": "clip.mp4"}, nil) + require.NoError(t, err) + require.Equal(t, "v2v", mode) +} + +func TestDerivedResolutionHelpers(t *testing.T) { + mp, err := derivedMegapixels(map[string]any{"resolution": "1280x720"}, nil) + require.NoError(t, err) + require.InDelta(t, 0.9216, mp, 0.0001) + + ratio, err := derivedAspectRatio(map[string]any{"options": map[string]any{"size": "1280x720"}}, nil) + require.NoError(t, err) + require.Equal(t, "16:9", ratio) +} + +func TestHasNonEmptyValue(t *testing.T) { + require.False(t, hasNonEmptyValue("")) + require.False(t, hasNonEmptyValue(" ")) + require.False(t, hasNonEmptyValue(nil)) + require.False(t, hasNonEmptyValue([]any{})) + require.True(t, hasNonEmptyValue([]any{"x"})) + require.True(t, hasNonEmptyValue("x")) + require.True(t, hasNonEmptyValue(0)) +} diff --git a/relay/helper/dimension_resolver.go b/relay/helper/dimension_resolver.go new file mode 100644 index 0000000..fb523bd --- /dev/null +++ b/relay/helper/dimension_resolver.go @@ -0,0 +1,251 @@ +package helper + +import ( + "fmt" + "io" + "mime/multipart" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" +) + +type derivedPricingFunc func(req, resp map[string]any) (any, error) + +var derivedPricingFuncs = map[string]derivedPricingFunc{} + +type ResolvedDimensionValues struct { + Values map[string]any + Missing map[string]bool + RawInput map[string]any +} + +type DimensionResolver struct { + c *gin.Context + request map[string]any + response map[string]any +} + +func NewDimensionResolver(c *gin.Context) (*DimensionResolver, error) { + contentType := c.Request.Header.Get("Content-Type") + req := map[string]any{} + + switch { + case strings.HasPrefix(contentType, "application/json"): + storage, err := common.GetBodyStorage(c) + if err != nil { + return nil, err + } + body, err := storage.Bytes() + if err != nil { + return nil, err + } + if len(strings.TrimSpace(string(body))) > 0 { + if err := common.Unmarshal(body, &req); err != nil { + return nil, err + } + } + if err := resetRequestBody(c, storage); err != nil { + return nil, err + } + case strings.Contains(contentType, gin.MIMEMultipartPOSTForm): + form, err := common.ParseMultipartFormReusable(c) + if err != nil { + return nil, err + } + req = requestMapFromMultipart(form) + default: + return nil, fmt.Errorf("unsupported content-type for pricing dimensions: %s", contentType) + } + + return &DimensionResolver{ + c: c, + request: req, + response: map[string]any{}, + }, nil +} + +func (r *DimensionResolver) ResolveDimensions(dims []types.PricingDimension) (*ResolvedDimensionValues, error) { + resolved := &ResolvedDimensionValues{ + Values: make(map[string]any, len(dims)), + Missing: make(map[string]bool), + RawInput: types.CloneMapAny(r.request), + } + + for _, dim := range dims { + value, exists, err := r.resolveSource(dim.Source) + if err != nil { + return nil, err + } + if !exists || !hasNonEmptyValue(value) { + if dim.Optional { + resolved.Values[dim.Key] = dim.Default + continue + } + resolved.Values[dim.Key] = nil + resolved.Missing[dim.Key] = true + continue + } + converted, err := convertPricingDimensionValue(value, dim.Type) + if err != nil { + return nil, fmt.Errorf("dimension %q: %w", dim.Key, err) + } + resolved.Values[dim.Key] = converted + } + + return resolved, nil +} + +func IsAllowedPricingSource(source string) bool { + if strings.HasPrefix(source, "response.") { + return false + } + if strings.HasPrefix(source, "request.") { + return len(strings.TrimPrefix(source, "request.")) > 0 + } + if _, ok := derivedPricingFuncs[strings.TrimPrefix(source, "derived.")]; ok && strings.HasPrefix(source, "derived.") { + return true + } + return false +} + +func registerDerivedPricingFunc(name string, fn derivedPricingFunc) { + derivedPricingFuncs[name] = fn +} + +func (r *DimensionResolver) resolveSource(source string) (any, bool, error) { + if strings.HasPrefix(source, "response.") { + return nil, false, fmt.Errorf("pricing source %q is unsupported in schema v1", source) + } + if strings.HasPrefix(source, "request.") { + return getNestedValue(r.request, strings.TrimPrefix(source, "request.")) + } + if strings.HasPrefix(source, "derived.") { + name := strings.TrimPrefix(source, "derived.") + fn, ok := derivedPricingFuncs[name] + if !ok { + return nil, false, fmt.Errorf("derived pricing source %q is not registered", source) + } + value, err := fn(r.request, r.response) + if err != nil { + return nil, false, err + } + return value, hasNonEmptyValue(value), nil + } + return nil, false, fmt.Errorf("pricing source %q is not allowed", source) +} + +func getNestedValue(src map[string]any, path string) (any, bool, error) { + if path == "" { + return nil, false, nil + } + var cur any = src + for _, part := range strings.Split(path, ".") { + m, ok := cur.(map[string]any) + if !ok { + return nil, false, nil + } + cur, ok = m[part] + if !ok { + return nil, false, nil + } + } + return cur, true, nil +} + +func convertPricingDimensionValue(value any, typ string) (any, error) { + switch typ { + case "string": + if s, ok := value.(string); ok { + return s, nil + } + return fmt.Sprintf("%v", value), nil + case "number": + switch v := value.(type) { + case float64: + return v, nil + case float32: + return float64(v), nil + case int: + return float64(v), nil + case int64: + return float64(v), nil + case int32: + return float64(v), nil + case string: + f, err := strconv.ParseFloat(strings.TrimSpace(v), 64) + if err != nil { + return nil, fmt.Errorf("value %q is not a number", v) + } + return f, nil + default: + return nil, fmt.Errorf("value is not a number") + } + case "boolean": + switch v := value.(type) { + case bool: + return v, nil + case string: + b, err := strconv.ParseBool(strings.TrimSpace(v)) + if err != nil { + return nil, fmt.Errorf("value %q is not a boolean", v) + } + return b, nil + default: + return nil, fmt.Errorf("value is not a boolean") + } + default: + return nil, fmt.Errorf("unknown dimension type %q", typ) + } +} + +func requestMapFromMultipart(form *multipart.Form) map[string]any { + req := map[string]any{} + for key, vals := range form.Value { + if len(vals) == 1 { + assignNestedValue(req, key, vals[0]) + continue + } + copied := append([]string(nil), vals...) + assignNestedValue(req, key, copied) + } + for key, files := range form.File { + if len(files) == 1 { + assignNestedValue(req, key, files[0].Filename) + continue + } + names := make([]string, 0, len(files)) + for _, file := range files { + names = append(names, file.Filename) + } + assignNestedValue(req, key, names) + } + return req +} + +func assignNestedValue(dst map[string]any, key string, value any) { + parts := strings.Split(key, ".") + cur := dst + for i, part := range parts { + if i == len(parts)-1 { + cur[part] = value + return + } + next, ok := cur[part].(map[string]any) + if !ok { + next = map[string]any{} + cur[part] = next + } + cur = next + } +} + +func resetRequestBody(c *gin.Context, storage common.BodyStorage) error { + if _, err := storage.Seek(0, io.SeekStart); err != nil { + return err + } + c.Request.Body = io.NopCloser(storage) + return nil +} diff --git a/relay/helper/dimension_resolver_test.go b/relay/helper/dimension_resolver_test.go new file mode 100644 index 0000000..05cab01 --- /dev/null +++ b/relay/helper/dimension_resolver_test.go @@ -0,0 +1,127 @@ +package helper + +import ( + "bytes" + "mime/multipart" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/types" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func TestDimensionResolver_JSONBody(t *testing.T) { + c := newPricingJSONContext(`{"resolution":"720P","duration":"10","options":{"size":"1280x720"},"image_url":"https://example.test/a.png"}`) + + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + values, err := resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "resolution", Source: "request.resolution", Type: "string"}, + {Key: "size", Source: "request.options.size", Type: "string"}, + {Key: "duration", Source: "request.duration", Type: "number"}, + {Key: "mode", Source: "derived.input_mode", Type: "string"}, + }) + require.NoError(t, err) + + require.Equal(t, "720P", values.Values["resolution"]) + require.Equal(t, "1280x720", values.Values["size"]) + require.Equal(t, 10.0, values.Values["duration"]) + require.Equal(t, "i2v", values.Values["mode"]) + require.Empty(t, values.Missing) +} + +func TestDimensionResolver_DerivedInputModePriority(t *testing.T) { + c := newPricingJSONContext(`{"image_url":"https://example.test/a.png","video_url":"https://example.test/a.mp4"}`) + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + values, err := resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "mode", Source: "derived.input_mode", Type: "string"}, + }) + require.NoError(t, err) + require.Equal(t, "v2v", values.Values["mode"]) +} + +func TestDimensionResolver_MissingRequiredValue(t *testing.T) { + c := newPricingJSONContext(`{"resolution":"720P"}`) + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + values, err := resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "duration", Source: "request.duration", Type: "number"}, + }) + require.NoError(t, err) + require.True(t, values.Missing["duration"]) + require.Nil(t, values.Values["duration"]) +} + +func TestDimensionResolver_OptionalDefault(t *testing.T) { + c := newPricingJSONContext(`{"resolution":"720P"}`) + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + values, err := resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "duration", Source: "request.duration", Type: "number", Optional: true, Default: 5.0}, + }) + require.NoError(t, err) + require.Equal(t, 5.0, values.Values["duration"]) + require.Empty(t, values.Missing) +} + +func TestDimensionResolver_MultipartFileInputMode(t *testing.T) { + body := &bytes.Buffer{} + writer := multipart.NewWriter(body) + require.NoError(t, writer.WriteField("resolution", "720P")) + part, err := writer.CreateFormFile("image", "frame.png") + require.NoError(t, err) + _, err = part.Write([]byte("fake")) + require.NoError(t, err) + require.NoError(t, writer.Close()) + + c := newPricingContext(body.String(), writer.FormDataContentType()) + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + values, err := resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "resolution", Source: "request.resolution", Type: "string"}, + {Key: "mode", Source: "derived.input_mode", Type: "string"}, + }) + require.NoError(t, err) + require.Equal(t, "720P", values.Values["resolution"]) + require.Equal(t, "i2v", values.Values["mode"]) +} + +func TestDimensionResolver_UnsupportedContentType(t *testing.T) { + c := newPricingContext("resolution=720P", "application/x-www-form-urlencoded") + _, err := NewDimensionResolver(c) + require.ErrorContains(t, err, "unsupported content-type") +} + +func TestDimensionResolver_ResponseSourceUnsupported(t *testing.T) { + c := newPricingJSONContext(`{"resolution":"720P"}`) + resolver, err := NewDimensionResolver(c) + require.NoError(t, err) + + _, err = resolver.ResolveDimensions([]types.PricingDimension{ + {Key: "tokens", Source: "response.usage.total_tokens", Type: "number"}, + }) + require.ErrorContains(t, err, "unsupported in schema v1") +} + +func newPricingJSONContext(body string) *gin.Context { + return newPricingContext(body, "application/json") +} + +func newPricingContext(body string, contentType string) *gin.Context { + gin.SetMode(gin.TestMode) + w := httptest.NewRecorder() + c, _ := gin.CreateTestContext(w) + req, _ := http.NewRequest(http.MethodPost, "/v1/video", strings.NewReader(body)) + req.Header.Set("Content-Type", contentType) + c.Request = req + return c +}