feat: add Gemini image compatibility

2026-06-24 00:42:49 +08:00
parent 7f32446466
commit c15842a94d
6 changed files with 966 additions and 15 deletions
@@ -1026,6 +1026,130 @@ func TestGeminiClientChatConvertsMediaContentParts(t *testing.T) {
 	}
 }

+func TestGeminiClientImageGenerateBuildsNativeImageBody(t *testing.T) {
+	var captured map[string]any
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"candidates": []any{map[string]any{
+				"content": map[string]any{
+					"parts": []any{map[string]any{
+						"inlineData": map[string]any{
+							"mimeType": "image/png",
+							"data":     "aW1hZ2U=",
+						},
+					}},
+				},
+			}},
+		})
+	}))
+	defer server.Close()
+
+	response, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:      "images.generations",
+		ModelType: "image_generate",
+		Model:     "gemini-image",
+		Body: map[string]any{
+			"model":  "gemini-image",
+			"prompt": "draw a cat",
+			"generationConfig": map[string]any{
+				"temperature": 0.5,
+				"imageConfig": map[string]any{
+					"aspectRatio": "16:9",
+				},
+			},
+			"resolution": "2K",
+			"n":          2,
+		},
+		Candidate: store.RuntimeModelCandidate{
+			BaseURL:           server.URL,
+			ProviderModelName: "gemini-image",
+			ModelType:         "image_generate",
+			Credentials:       map[string]any{"apiKey": "gemini-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run gemini image generate: %v", err)
+	}
+	config := captured["generationConfig"].(map[string]any)
+	modalities := config["responseModalities"].([]any)
+	imageConfig := config["imageConfig"].(map[string]any)
+	if modalities[0] != "IMAGE" || config["temperature"] != 0.5 || numericValue(config["candidateCount"], 0) != 2 || imageConfig["aspectRatio"] != "16:9" || imageConfig["imageSize"] != "2K" {
+		t.Fatalf("unexpected generationConfig: %+v", config)
+	}
+	contents := captured["contents"].([]any)
+	parts := contents[0].(map[string]any)["parts"].([]any)
+	if parts[0].(map[string]any)["text"] != "draw a cat" {
+		t.Fatalf("unexpected image contents: %+v", captured)
+	}
+	data := response.Result["data"].([]any)
+	if data[0].(map[string]any)["b64_json"] != "aW1hZ2U=" {
+		t.Fatalf("unexpected image response: %+v", response.Result)
+	}
+}
+
+func TestGeminiClientImageEditPreservesNativeContentsAndFileData(t *testing.T) {
+	var captured map[string]any
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&captured); err != nil {
+			t.Fatalf("decode request: %v", err)
+		}
+		_ = json.NewEncoder(w).Encode(map[string]any{
+			"candidates": []any{map[string]any{
+				"content": map[string]any{
+					"parts": []any{map[string]any{
+						"fileData": map[string]any{
+							"mimeType": "image/webp",
+							"fileUri":  "https://cdn.example/out.webp",
+						},
+					}},
+				},
+			}},
+		})
+	}))
+	defer server.Close()
+
+	_, err := (GeminiClient{HTTPClient: server.Client()}).Run(context.Background(), Request{
+		Kind:      "images.edits",
+		ModelType: "image_edit",
+		Model:     "gemini-image",
+		Body: map[string]any{
+			"model": "gemini-image",
+			"contents": []any{map[string]any{
+				"parts": []any{
+					map[string]any{"fileData": map[string]any{
+						"mimeType": "image/png",
+						"fileUri":  "https://cdn.example/input.png",
+					}},
+					map[string]any{"text": "edit it"},
+				},
+			}},
+			"generationConfig": map[string]any{"temperature": 0.2},
+		},
+		Candidate: store.RuntimeModelCandidate{
+			BaseURL:           server.URL,
+			ProviderModelName: "gemini-image",
+			ModelType:         "image_edit",
+			Credentials:       map[string]any{"apiKey": "gemini-key"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("run gemini image edit: %v", err)
+	}
+	config := captured["generationConfig"].(map[string]any)
+	if config["temperature"] != 0.2 {
+		t.Fatalf("generationConfig should preserve caller settings: %+v", config)
+	}
+	contents := captured["contents"].([]any)
+	parts := contents[0].(map[string]any)["parts"].([]any)
+	fileData := parts[0].(map[string]any)["fileData"].(map[string]any)
+	if fileData["fileUri"] != "https://cdn.example/input.png" {
+		t.Fatalf("native fileData should be preserved: %+v", captured)
+	}
+}
+
 func TestGeminiClientChatRestoresToolContext(t *testing.T) {
 	var captured map[string]any
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -69,17 +69,20 @@ func geminiURL(baseURL string, model string, apiKey string) string {
 }

 func geminiBody(request Request) map[string]any {
+	if geminiImageModelType(request.ModelType) {
+		return geminiImageBody(request)
+	}
 	if contents, ok := request.Body["contents"]; ok {
-		return map[string]any{"contents": contents}
+		return geminiApplyRequestOptions(map[string]any{"contents": contents}, request, false)
 	}
 	prompt := firstNonEmptyPrompt(request.Body, "")
 	if prompt != "" {
-		return map[string]any{
+		return geminiApplyRequestOptions(map[string]any{
 			"contents": []any{map[string]any{
 				"role":  "user",
 				"parts": []any{map[string]any{"text": prompt}},
 			}},
-		}
+		}, request, false)
 	}
 	body := map[string]any{"contents": geminiContentsFromMessages(request.Body)}
 	if tools := geminiToolsFromOpenAITools(request.Body["tools"]); len(tools) > 0 {
@@ -87,12 +90,156 @@ func geminiBody(request Request) map[string]any {
 	}
 	contents, _ := body["contents"].([]any)
 	if len(contents) > 0 {
-		return body
+		return geminiApplyRequestOptions(body, request, false)
 	}
-	return map[string]any{"contents": []any{map[string]any{
+	return geminiApplyRequestOptions(map[string]any{"contents": []any{map[string]any{
 		"role":  "user",
 		"parts": []any{map[string]any{"text": textFromMessages(request.Body)}},
-	}}}
+	}}}, request, false)
+}
+
+func geminiImageBody(request Request) map[string]any {
+	body := map[string]any{}
+	if contents, ok := request.Body["contents"]; ok {
+		body["contents"] = contents
+	} else {
+		parts := make([]any, 0)
+		for _, image := range geminiImageInputValues(request.Body) {
+			if media := geminiMediaURLPart(image.URI, image.MimeType, "image"); media != nil {
+				parts = append(parts, media)
+			}
+		}
+		if prompt := firstNonEmptyPrompt(request.Body, ""); prompt != "" {
+			parts = append(parts, map[string]any{"text": prompt})
+		}
+		body["contents"] = []any{map[string]any{
+			"role":  "user",
+			"parts": parts,
+		}}
+	}
+	return geminiApplyRequestOptions(body, request, true)
+}
+
+type geminiImageInputValue struct {
+	URI      string
+	MimeType string
+}
+
+func geminiImageInputValues(body map[string]any) []geminiImageInputValue {
+	values := make([]geminiImageInputValue, 0)
+	for _, key := range []string{"image", "images", "image_url", "imageUrl", "image_urls", "imageUrls"} {
+		values = append(values, geminiImageInputValuesFromAny(body[key])...)
+	}
+	return values
+}
+
+func geminiImageInputValuesFromAny(value any) []geminiImageInputValue {
+	switch typed := value.(type) {
+	case string:
+		if uri := strings.TrimSpace(typed); uri != "" {
+			return []geminiImageInputValue{{URI: uri}}
+		}
+	case []any:
+		out := make([]geminiImageInputValue, 0, len(typed))
+		for _, item := range typed {
+			out = append(out, geminiImageInputValuesFromAny(item)...)
+		}
+		return out
+	case []string:
+		out := make([]geminiImageInputValue, 0, len(typed))
+		for _, item := range typed {
+			if uri := strings.TrimSpace(item); uri != "" {
+				out = append(out, geminiImageInputValue{URI: uri})
+			}
+		}
+		return out
+	case map[string]any:
+		if uri := firstNonEmptyString(typed["url"], typed["fileUri"], typed["file_uri"]); uri != "" {
+			return []geminiImageInputValue{{
+				URI:      uri,
+				MimeType: firstNonEmptyString(typed["mime_type"], typed["mimeType"]),
+			}}
+		}
+	}
+	return nil
+}
+
+func geminiApplyRequestOptions(body map[string]any, request Request, imageResponse bool) map[string]any {
+	if generationConfig := geminiGenerationConfig(request.Body, imageResponse); len(generationConfig) > 0 {
+		body["generationConfig"] = generationConfig
+	}
+	for _, item := range []struct {
+		from []string
+		to   string
+	}{
+		{from: []string{"systemInstruction", "system_instruction"}, to: "systemInstruction"},
+		{from: []string{"safetySettings", "safety_settings"}, to: "safetySettings"},
+		{from: []string{"toolConfig", "tool_config"}, to: "toolConfig"},
+		{from: []string{"cachedContent", "cached_content"}, to: "cachedContent"},
+	} {
+		for _, key := range item.from {
+			if value, ok := request.Body[key]; ok {
+				body[item.to] = value
+				break
+			}
+		}
+	}
+	if _, exists := body["tools"]; !exists {
+		if tools, ok := request.Body["tools"]; ok {
+			body["tools"] = tools
+		}
+	}
+	return body
+}
+
+func geminiGenerationConfig(body map[string]any, imageResponse bool) map[string]any {
+	source := mapFromAny(firstPresent(body["generationConfig"], body["generation_config"]))
+	out := cloneMapAny(source)
+	if aspectRatio := firstNonEmptyString(body["aspect_ratio"], body["aspectRatio"]); aspectRatio != "" {
+		imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
+		imageConfig["aspectRatio"] = aspectRatio
+		out["imageConfig"] = imageConfig
+		delete(out, "image_config")
+	}
+	if imageSize := firstNonEmptyString(body["resolution"], body["imageSize"], body["image_size"], body["size"]); imageSize != "" {
+		imageConfig := cloneMapAny(mapFromAny(firstPresent(out["imageConfig"], out["image_config"])))
+		imageConfig["imageSize"] = imageSize
+		out["imageConfig"] = imageConfig
+		delete(out, "image_config")
+	}
+	if count := intFromAny(firstPresent(body["n"], body["candidateCount"], body["candidate_count"])); count > 0 {
+		out["candidateCount"] = count
+	}
+	if imageResponse {
+		out["responseModalities"] = ensureGeminiResponseModalityImage(out["responseModalities"])
+	}
+	return out
+}
+
+func ensureGeminiResponseModalityImage(value any) []any {
+	items := []any{}
+	switch typed := value.(type) {
+	case []any:
+		items = append(items, typed...)
+	case []string:
+		for _, item := range typed {
+			items = append(items, item)
+		}
+	case string:
+		if strings.TrimSpace(typed) != "" {
+			items = append(items, typed)
+		}
+	}
+	for _, item := range items {
+		if strings.EqualFold(strings.TrimSpace(stringFromAny(item)), "IMAGE") {
+			return items
+		}
+	}
+	return append(items, "IMAGE")
+}
+
+func geminiImageModelType(modelType string) bool {
+	return modelType == "image_generate" || modelType == "image_edit"
 }

 func geminiContentsFromMessages(body map[string]any) []any {
@@ -387,7 +534,7 @@ func geminiToolsFromOpenAITools(value any) []any {
 }

 func geminiResult(request Request, raw map[string]any) map[string]any {
-	if request.ModelType == "image" {
+	if geminiImageModelType(request.ModelType) {
 		data := geminiImageData(raw)
 		if len(data) == 0 {
 			data = []any{map[string]any{"url": "/static/provider/gemini-image-placeholder.png"}}
@@ -507,7 +654,21 @@ func geminiImageData(raw map[string]any) []any {
 				inline, _ = partMap["inline_data"].(map[string]any)
 			}
 			if data, ok := inline["data"].(string); ok && data != "" {
-				out = append(out, map[string]any{"b64_json": data, "mime_type": inline["mimeType"]})
+				out = append(out, map[string]any{
+					"b64_json":  data,
+					"mime_type": firstNonEmptyString(inline["mimeType"], inline["mime_type"]),
+				})
+				continue
+			}
+			fileData, _ := partMap["fileData"].(map[string]any)
+			if fileData == nil {
+				fileData, _ = partMap["file_data"].(map[string]any)
+			}
+			if uri := firstNonEmptyString(fileData["fileUri"], fileData["file_uri"], fileData["uri"]); uri != "" {
+				out = append(out, map[string]any{
+					"url":       uri,
+					"mime_type": firstNonEmptyString(fileData["mimeType"], fileData["mime_type"], mimeFromURI(uri)),
+				})
 			}
 		}
 	}