Merge pull request #4175 from router-for-me/image

This commit is contained in:
Luis Pater
2026-07-10 12:05:45 +08:00
committed by GitHub
2 changed files with 53 additions and 1 deletions

View File

@@ -1746,6 +1746,27 @@ func isCodexFreePlanAuth(auth *cliproxyauth.Auth) bool {
return strings.EqualFold(strings.TrimSpace(auth.Attributes["plan_type"]), "free")
}
func isImageGenerationFunctionTool(tool gjson.Result) bool {
switch tool.Get("type").String() {
case "function":
return tool.Get("name").String() == "image_gen.imagegen"
case "namespace":
if tool.Get("name").String() != "image_gen" {
return false
}
tools := tool.Get("tools")
if !tools.IsArray() {
return false
}
for _, nestedTool := range tools.Array() {
if nestedTool.Get("type").String() == "function" && nestedTool.Get("name").String() == "imagegen" {
return true
}
}
}
return false
}
func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth.Auth) []byte {
if strings.HasSuffix(baseModel, "spark") {
return body
@@ -1760,7 +1781,7 @@ func ensureImageGenerationTool(body []byte, baseModel string, auth *cliproxyauth
return body
}
for _, t := range tools.Array() {
if t.Get("type").String() == "image_generation" {
if t.Get("type").String() == "image_generation" || isImageGenerationFunctionTool(t) {
return body
}
}

View File

@@ -58,6 +58,37 @@ func TestEnsureImageGenerationTool_AlreadyPresent(t *testing.T) {
}
}
func TestEnsureImageGenerationTool_ImageGenNamespaceDoesNotInjectTool(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"namespace","name":"image_gen","tools":[{"type":"function","name":"imagegen","parameters":{}}]}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
if string(result) != string(body) {
t.Fatalf("expected body to be unchanged, got %s", string(result))
}
}
func TestEnsureImageGenerationTool_FlattenedImageGenFunctionDoesNotInjectTool(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"function","name":"image_gen.imagegen","parameters":{}}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
if string(result) != string(body) {
t.Fatalf("expected body to be unchanged, got %s", string(result))
}
}
func TestEnsureImageGenerationTool_SimilarNamespaceStillInjectsTool(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[{"type":"namespace","name":"image_tools","tools":[{"type":"function","name":"imagegen","parameters":{}}]}]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)
tools := gjson.GetBytes(result, "tools").Array()
if len(tools) != 2 {
t.Fatalf("expected 2 tools, got %d", len(tools))
}
if tools[1].Get("type").String() != "image_generation" {
t.Fatalf("expected second tool type=image_generation, got %s", tools[1].Get("type").String())
}
}
func TestEnsureImageGenerationTool_EmptyToolsArray(t *testing.T) {
body := []byte(`{"model":"gpt-5.4","tools":[]}`)
result := ensureImageGenerationTool(body, "gpt-5.4", nil)