Codex/group media fixed pricing

common/billing_model.go

@@ -0,0 +1,50 @@
+package common
+
+import "strings"
+
+func normalizeBillingModelName(modelName string) string {
+	return strings.ToLower(strings.TrimSpace(modelName))
+}
+
+func IsDurationOnlyBillingModel(modelName string) bool {
+	name := normalizeBillingModelName(modelName)
+	if name == "" {
+		return false
+	}
+	if strings.HasPrefix(name, "veo") {
+		return true
+	}
+	return strings.Contains(name, "grok-imagine") && strings.Contains(name, "video")
+}
+
+func IsResolutionOnlyBillingModel(modelName string) bool {
+	name := normalizeBillingModelName(modelName)
+	if name == "" {
+		return false
+	}
+	return strings.Contains(name, "banana")
+}
+
+func FilterOtherRatiosForBillingModel(modelName string, ratios map[string]float64) map[string]float64 {
+	if len(ratios) == 0 {
+		return map[string]float64{}
+	}
+
+	filtered := make(map[string]float64, len(ratios))
+	switch {
+	case IsDurationOnlyBillingModel(modelName):
+		if ratio, ok := ratios["seconds"]; ok && ratio > 0 {
+			filtered["seconds"] = ratio
+		}
+	case IsResolutionOnlyBillingModel(modelName):
+		return filtered
+	default:
+		for key, ratio := range ratios {
+			if ratio > 0 {
+				filtered[key] = ratio
+			}
+		}
+	}
+
+	return filtered
+}

common/billing_model_test.go

@@ -0,0 +1,30 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFilterOtherRatiosForDurationOnlyModel(t *testing.T) {
+	filtered := FilterOtherRatiosForBillingModel("grok-imagine-1.0-video", map[string]float64{
+		"seconds":    6,
+		"size":       1.666667,
+		"resolution": 1.5,
+	})
+
+	assert.Equal(t, map[string]float64{
+		"seconds": 6,
+	}, filtered)
+}
+
+func TestFilterOtherRatiosForResolutionOnlyModel(t *testing.T) {
+	filtered := FilterOtherRatiosForBillingModel("nano-banana-pro", map[string]float64{
+		"resolution":        2,
+		"quality":           1.5,
+		"output_resolution": 4,
+		"n":                 3,
+	})
+
+	assert.Empty(t, filtered)
+}

common/endpoint_defaults.go

@@ -24,7 +24,9 @@ var defaultEndpointInfoMap = map[constant.EndpointType]EndpointInfo{
 	constant.EndpointTypeGemini:                {Path: "/v1beta/models/{model}:generateContent", Method: "POST"},
 	constant.EndpointTypeJinaRerank:            {Path: "/v1/rerank", Method: "POST"},
 	constant.EndpointTypeImageGeneration:       {Path: "/v1/images/generations", Method: "POST"},
+	constant.EndpointTypeImageEdit:             {Path: "/v1/images/edits", Method: "POST"},
 	constant.EndpointTypeEmbeddings:            {Path: "/v1/embeddings", Method: "POST"},
+	constant.EndpointTypeOpenAIVideo:           {Path: "/v1/video/generations", Method: "POST"},
 }
 
 // GetDefaultEndpointInfo 返回指定端点类型的默认信息以及是否存在

common/endpoint_type.go

@@ -37,7 +37,11 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant
 			endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI}
 		}
 	}
-	if IsImageGenerationModel(modelName) {
+	if IsOpenAIVideoModel(modelName) {
+		endpointTypes = append([]constant.EndpointType{constant.EndpointTypeOpenAIVideo}, endpointTypes...)
+	} else if IsImageEditModel(modelName) {
+		endpointTypes = append([]constant.EndpointType{constant.EndpointTypeImageEdit}, endpointTypes...)
+	} else if IsImageGenerationModel(modelName) {
 		// add to first
 		endpointTypes = append([]constant.EndpointType{constant.EndpointTypeImageGeneration}, endpointTypes...)
 	}

common/endpoint_type_test.go

@@ -0,0 +1,59 @@
+package common
+
+import (
+	"testing"
+
+	"github.com/QuantumNous/new-api/constant"
+)
+
+func TestGetEndpointTypesByChannelTypeRecognizesGrokImagineImageModels(t *testing.T) {
+	tests := []struct {
+		name     string
+		channel  int
+		model    string
+		expected constant.EndpointType
+	}{
+		{
+			name:     "grok imagine 1.0 generation",
+			channel:  constant.ChannelTypeXai,
+			model:    "grok-imagine-1.0",
+			expected: constant.EndpointTypeImageGeneration,
+		},
+		{
+			name:     "grok imagine 1.0 fast generation",
+			channel:  constant.ChannelTypeXai,
+			model:    "grok-imagine-1.0-fast",
+			expected: constant.EndpointTypeImageGeneration,
+		},
+		{
+			name:     "grok imagine 1.0 edit",
+			channel:  constant.ChannelTypeXai,
+			model:    "grok-imagine-1.0-edit",
+			expected: constant.EndpointTypeImageEdit,
+		},
+		{
+			name:     "grok imagine 1.0 video on xai",
+			channel:  constant.ChannelTypeXai,
+			model:    "grok-imagine-1.0-video",
+			expected: constant.EndpointTypeOpenAIVideo,
+		},
+		{
+			name:     "grok imagine video stays video",
+			channel:  constant.ChannelTypeSora,
+			model:    "grok-imagine-1.0-video",
+			expected: constant.EndpointTypeOpenAIVideo,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := GetEndpointTypesByChannelType(tt.channel, tt.model)
+			if len(got) == 0 {
+				t.Fatalf("expected endpoint types for %s", tt.model)
+			}
+			if got[0] != tt.expected {
+				t.Fatalf("expected first endpoint %s, got %s", tt.expected, got[0])
+			}
+		})
+	}
+}

common/init.go

@@ -150,7 +150,9 @@ func initConstantEnv() {
 	// 任务轮询时查询的最大数量
 	constant.TaskQueryLimit = GetEnvOrDefault("TASK_QUERY_LIMIT", 1000)
 	// 异步任务超时时间（分钟），超过此时间未完成的任务将被标记为失败并退款。0 表示禁用。
-	constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 1440)
+	constant.TaskTimeoutMinutes = GetEnvOrDefault("TASK_TIMEOUT_MINUTES", 20)
+	// 视频轮询遇到泛化 404 Not Found 时的宽限时间（分钟）。0 表示立即失败。
+	constant.TaskNotFoundGraceMinutes = GetEnvOrDefault("TASK_NOT_FOUND_GRACE_MINUTES", 10)
 
 	soraPatchStr := GetEnvOrDefaultString("TASK_PRICE_PATCH", "")
 	if soraPatchStr != "" {

common/model.go

@@ -13,10 +13,20 @@ var (
 		"dall-e-3",
 		"dall-e-2",
 		"gpt-image-1",
+		"gpt-image2",
+		"exact:grok-imagine-1.0",
+		"exact:grok-imagine-1.0-fast",
 		"prefix:imagen-",
 		"flux-",
 		"flux.1-",
 	}
+	ImageEditModels = []string{
+		"exact:grok-imagine-1.0-edit",
+	}
+	OpenAIVideoModels = []string{
+		"exact:grok-imagine-1.0-video",
+		"exact:grok-imagine-video",
+	}
 	OpenAITextModels = []string{
 		"gpt-",
 		"o1",
@@ -38,6 +48,41 @@ func IsOpenAIResponseOnlyModel(modelName string) bool {
 func IsImageGenerationModel(modelName string) bool {
 	modelName = strings.ToLower(modelName)
 	for _, m := range ImageGenerationModels {
+		if strings.HasPrefix(m, "exact:") && modelName == strings.TrimPrefix(m, "exact:") {
+			return true
+		}
+		if strings.Contains(modelName, m) {
+			return true
+		}
+		if strings.HasPrefix(m, "prefix:") && strings.HasPrefix(modelName, strings.TrimPrefix(m, "prefix:")) {
+			return true
+		}
+	}
+	return false
+}
+
+func IsImageEditModel(modelName string) bool {
+	modelName = strings.ToLower(modelName)
+	for _, m := range ImageEditModels {
+		if strings.HasPrefix(m, "exact:") && modelName == strings.TrimPrefix(m, "exact:") {
+			return true
+		}
+		if strings.Contains(modelName, m) {
+			return true
+		}
+		if strings.HasPrefix(m, "prefix:") && strings.HasPrefix(modelName, strings.TrimPrefix(m, "prefix:")) {
+			return true
+		}
+	}
+	return false
+}
+
+func IsOpenAIVideoModel(modelName string) bool {
+	modelName = strings.ToLower(modelName)
+	for _, m := range OpenAIVideoModels {
+		if strings.HasPrefix(m, "exact:") && modelName == strings.TrimPrefix(m, "exact:") {
+			return true
+		}
 		if strings.Contains(modelName, m) {
 			return true
 		}

constant/endpoint_type.go

@@ -10,6 +10,7 @@ const (
 	EndpointTypeGemini                EndpointType = "gemini"
 	EndpointTypeJinaRerank            EndpointType = "jina-rerank"
 	EndpointTypeImageGeneration       EndpointType = "image-generation"
+	EndpointTypeImageEdit             EndpointType = "image-edit"
 	EndpointTypeEmbeddings            EndpointType = "embeddings"
 	EndpointTypeOpenAIVideo           EndpointType = "openai-video"
 	//EndpointTypeMidjourney     EndpointType = "midjourney-proxy"

constant/env.go

@@ -17,6 +17,7 @@ var GenerateDefaultToken bool
 var ErrorLogEnabled bool
 var TaskQueryLimit int
 var TaskTimeoutMinutes int
+var TaskNotFoundGraceMinutes int
 
 // temporary variable for sora patch, will be removed in future
 var TaskPricePatches []string

constant/task.go

@@ -11,6 +11,8 @@ const (
 	SunoActionMusic  = "MUSIC"
 	SunoActionLyrics = "LYRICS"
 
+	TaskActionImageGenerate     = "imageGenerate"
+	TaskActionImageEdit         = "imageEdit"
 	TaskActionGenerate          = "generate"
 	TaskActionTextGenerate      = "textGenerate"
 	TaskActionFirstTailGenerate = "firstTailGenerate"