feat(task-pricing): resolve pricing dimensions from task requests

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
+}

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))
+}

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
+}

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
+}