瀏覽代碼

✨ feat(endpoint types): add support for image generation models in endpoint type handling

CaIon 5 月之前
父節點
當前提交
30fb349d91
共有 3 個文件被更改,包括 32 次插入6 次删除
  1. 4 0
      common/endpoint_type.go
  2. 22 1
      common/model.go
  3. 6 5
      constant/endpoint_type.go

+ 4 - 0
common/endpoint_type.go

@@ -33,5 +33,9 @@ func GetEndpointTypesByChannelType(channelType int, modelName string) []constant
 			endpointTypes = []constant.EndpointType{constant.EndpointTypeOpenAI}
 		}
 	}
+	if IsImageGenerationModel(modelName) {
+		// add to first
+		endpointTypes = append([]constant.EndpointType{constant.EndpointTypeImageGeneration}, endpointTypes...)
+	}
 	return endpointTypes
 }

+ 22 - 1
common/model.go

@@ -9,11 +9,32 @@ var (
 		"o3-deep-research",
 		"o4-mini-deep-research",
 	}
+	ImageGenerationModels = []string{
+		"dall-e-3",
+		"dall-e-2",
+		"gpt-image-1",
+		"prefix:imagen-",
+		"flux-",
+		"flux.1-",
+	}
 )
 
 func IsOpenAIResponseOnlyModel(modelName string) bool {
 	for _, m := range OpenAIResponseOnlyModels {
-		if strings.Contains(m, modelName) {
+		if strings.Contains(modelName, m) {
+			return true
+		}
+	}
+	return false
+}
+
+func IsImageGenerationModel(modelName string) bool {
+	modelName = strings.ToLower(modelName)
+	for _, m := range ImageGenerationModels {
+		if strings.Contains(modelName, m) {
+			return true
+		}
+		if strings.HasPrefix(m, "prefix:") && strings.HasPrefix(modelName, strings.TrimPrefix(m, "prefix:")) {
 			return true
 		}
 	}

+ 6 - 5
constant/endpoint_type.go

@@ -3,11 +3,12 @@ package constant
 type EndpointType string
 
 const (
-	EndpointTypeOpenAI         EndpointType = "openai"
-	EndpointTypeOpenAIResponse EndpointType = "openai-response"
-	EndpointTypeAnthropic      EndpointType = "anthropic"
-	EndpointTypeGemini         EndpointType = "gemini"
-	EndpointTypeJinaRerank     EndpointType = "jina-rerank"
+	EndpointTypeOpenAI          EndpointType = "openai"
+	EndpointTypeOpenAIResponse  EndpointType = "openai-response"
+	EndpointTypeAnthropic       EndpointType = "anthropic"
+	EndpointTypeGemini          EndpointType = "gemini"
+	EndpointTypeJinaRerank      EndpointType = "jina-rerank"
+	EndpointTypeImageGeneration EndpointType = "image-generation"
 	//EndpointTypeMidjourney     EndpointType = "midjourney-proxy"
 	//EndpointTypeSuno           EndpointType = "suno-proxy"
 	//EndpointTypeKling          EndpointType = "kling"