1

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,9 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant
 			endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI}
 		}
 	}
-	if IsImageGenerationModel(modelName) {
+	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,53 @@
+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 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/model.go

@@ -13,10 +13,15 @@ var (
 		"dall-e-3",
 		"dall-e-2",
 		"gpt-image-1",
+		"exact:grok-imagine-1.0",
+		"exact:grok-imagine-1.0-fast",
 		"prefix:imagen-",
 		"flux-",
 		"flux.1-",
 	}
+	ImageEditModels = []string{
+		"exact:grok-imagine-1.0-edit",
+	}
 	OpenAITextModels = []string{
 		"gpt-",
 		"o1",
@@ -38,6 +43,25 @@ 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
 		}

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"

controller/playground.go

@@ -35,22 +35,91 @@ func Playground(c *gin.Context) {
 		return
 	}
 
-	userId := c.GetInt("id")
+	if newAPIError = setupPlaygroundTokenContext(c, fmt.Sprintf("playground-%s", relayInfo.UsingGroup), relayInfo.UsingGroup); newAPIError != nil {
+		return
+	}
+
+	Relay(c, types.RelayFormatOpenAI)
+}
 
-	// Write user context to ensure acceptUnsetRatio is available
+func PlaygroundVideoSubmit(c *gin.Context) {
+	var newAPIError *types.NewAPIError
+	defer func() {
+		if newAPIError != nil {
+			c.JSON(newAPIError.StatusCode, gin.H{
+				"error": newAPIError.ToOpenAIError(),
+			})
+		}
+	}()
+	if newAPIError = setupPlaygroundTokenContext(c, "playground-video", c.GetString("group")); newAPIError != nil {
+		return
+	}
+	RelayTask(c)
+}
+
+func PlaygroundImageGenerations(c *gin.Context) {
+	var newAPIError *types.NewAPIError
+	defer func() {
+		if newAPIError != nil {
+			c.JSON(newAPIError.StatusCode, gin.H{
+				"error": newAPIError.ToOpenAIError(),
+			})
+		}
+	}()
+	if newAPIError = setupPlaygroundTokenContext(c, "playground-image", c.GetString("group")); newAPIError != nil {
+		return
+	}
+	Relay(c, types.RelayFormatOpenAIImage)
+}
+
+func PlaygroundImageEdits(c *gin.Context) {
+	var newAPIError *types.NewAPIError
+	defer func() {
+		if newAPIError != nil {
+			c.JSON(newAPIError.StatusCode, gin.H{
+				"error": newAPIError.ToOpenAIError(),
+			})
+		}
+	}()
+	if newAPIError = setupPlaygroundTokenContext(c, "playground-image-edit", c.GetString("group")); newAPIError != nil {
+		return
+	}
+	Relay(c, types.RelayFormatOpenAIImage)
+}
+
+func PlaygroundVideoFetch(c *gin.Context) {
+	var newAPIError *types.NewAPIError
+	defer func() {
+		if newAPIError != nil {
+			c.JSON(newAPIError.StatusCode, gin.H{
+				"error": newAPIError.ToOpenAIError(),
+			})
+		}
+	}()
+	if newAPIError = setupPlaygroundTokenContext(c, "playground-video-fetch", c.GetString("group")); newAPIError != nil {
+		return
+	}
+	RelayTaskFetch(c)
+}
+
+func setupPlaygroundTokenContext(c *gin.Context, tokenName string, tokenGroup string) *types.NewAPIError {
+	userId := c.GetInt("id")
 	userCache, err := model.GetUserCache(userId)
 	if err != nil {
-		newAPIError = types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
-		return
+		return types.NewError(err, types.ErrorCodeQueryDataError, types.ErrOptionWithSkipRetry())
 	}
 	userCache.WriteContext(c)
-
+	if tokenGroup == "" {
+		tokenGroup = c.GetString("group")
+	}
+	if tokenGroup == "" {
+		tokenGroup = userCache.Group
+	}
 	tempToken := &model.Token{
 		UserId: userId,
-		Name:   fmt.Sprintf("playground-%s", relayInfo.UsingGroup),
-		Group:  relayInfo.UsingGroup,
+		Name:   tokenName,
+		Group:  tokenGroup,
 	}
 	_ = middleware.SetupContextForToken(c, tempToken)
-
-	Relay(c, types.RelayFormatOpenAI)
+	return nil
 }

controller/relay.go

@@ -16,6 +16,7 @@ import (
 	"github.com/QuantumNous/new-api/middleware"
 	"github.com/QuantumNous/new-api/model"
 	"github.com/QuantumNous/new-api/relay"
+	taskcommon "github.com/QuantumNous/new-api/relay/channel/task/taskcommon"
 	relaycommon "github.com/QuantumNous/new-api/relay/common"
 	relayconstant "github.com/QuantumNous/new-api/relay/constant"
 	"github.com/QuantumNous/new-api/relay/helper"
@@ -586,6 +587,42 @@ func RelayTask(c *gin.Context) {
 		task.Quota = result.Quota
 		task.Data = result.TaskData
 		task.Action = relayInfo.Action
+		if adaptor := relay.GetTaskAdaptor(result.Platform); adaptor != nil && len(result.TaskData) > 0 {
+			if taskInfo, err := adaptor.ParseTaskResult(result.TaskData); err == nil && taskInfo != nil && taskInfo.Status != "" {
+				now := time.Now().Unix()
+				task.Status = model.TaskStatus(taskInfo.Status)
+				switch task.Status {
+				case model.TaskStatusSubmitted:
+					task.Progress = taskcommon.ProgressSubmitted
+				case model.TaskStatusQueued:
+					task.Progress = taskcommon.ProgressQueued
+				case model.TaskStatusInProgress:
+					task.Progress = taskcommon.ProgressInProgress
+					if task.StartTime == 0 {
+						task.StartTime = now
+					}
+				case model.TaskStatusSuccess:
+					task.Progress = taskcommon.ProgressComplete
+					if task.StartTime == 0 {
+						task.StartTime = now
+					}
+					if task.FinishTime == 0 {
+						task.FinishTime = now
+					}
+					task.PrivateData.ResultURL = taskInfo.Url
+				case model.TaskStatusFailure:
+					task.Progress = taskcommon.ProgressComplete
+					if task.FinishTime == 0 {
+						task.FinishTime = now
+					}
+					task.FailReason = taskInfo.Reason
+					task.PrivateData.ResultURL = taskInfo.Url
+				}
+				if taskInfo.Progress != "" {
+					task.Progress = taskInfo.Progress
+				}
+			}
+		}
 		if insertErr := task.Insert(); insertErr != nil {
 			common.SysError("insert task error: " + insertErr.Error())
 		}

dto/openai_request.go

@@ -46,6 +46,8 @@ type GeneralOpenAIRequest struct {
 	Input               any               `json:"input,omitempty"`
 	Instruction         string            `json:"instruction,omitempty"`
 	Size                string            `json:"size,omitempty"`
+	Seconds             *string           `json:"seconds,omitempty"`
+	Quality             *string           `json:"quality,omitempty"`
 	Functions           json.RawMessage   `json:"functions,omitempty"`
 	FrequencyPenalty    *float64          `json:"frequency_penalty,omitempty"`
 	PresencePenalty     *float64          `json:"presence_penalty,omitempty"`

middleware/distributor.go

@@ -81,8 +81,8 @@ func Distribute() func(c *gin.Context) {
 				}
 				var selectGroup string
 				usingGroup := common.GetContextKeyString(c, constant.ContextKeyUsingGroup)
-				// check path is /pg/chat/completions
-				if strings.HasPrefix(c.Request.URL.Path, "/pg/chat/completions") {
+				// playground requests may override group in body
+				if strings.HasPrefix(c.Request.URL.Path, "/pg/") {
 					playgroundRequest := &dto.PlayGroundRequest{}
 					err = common.UnmarshalBodyReusable(c, playgroundRequest)
 					if err != nil {