Codex/fix grok video url

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

@@ -57,6 +57,36 @@ func PlaygroundVideoSubmit(c *gin.Context) {
 	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() {

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

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 {