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) {