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
@@ -0,0 +1,87 @@
+package httpapi
+
+import "testing"
+
+func TestGeminiImageTaskBodyMapsTextOnlyToImageGenerate(t *testing.T) {
+	mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
+		"contents": []any{
+			map[string]any{"parts": []any{
+				map[string]any{"text": "draw"},
+				map[string]any{"text": "a cat"},
+			}},
+		},
+		"generationConfig": map[string]any{
+			"candidateCount": float64(2),
+			"imageConfig": map[string]any{
+				"aspectRatio": "16:9",
+				"imageSize":   "2K",
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("map Gemini request: %v", err)
+	}
+	if mapping.Kind != "images.generations" || mapping.Body["modelType"] != "image_generate" {
+		t.Fatalf("unexpected mapping kind/body: %+v", mapping)
+	}
+	if mapping.Body["prompt"] != "draw\na cat" || mapping.Body["n"] != 2 || mapping.Body["aspect_ratio"] != "16:9" || mapping.Body["resolution"] != "2K" {
+		t.Fatalf("unexpected mapped body: %+v", mapping.Body)
+	}
+}
+
+func TestGeminiImageTaskBodyMapsInlineDataToImageEdit(t *testing.T) {
+	mapping, err := geminiImageTaskBody("gemini-image", map[string]any{
+		"contents": []any{
+			map[string]any{"parts": []any{
+				map[string]any{"text": "make it brighter"},
+				map[string]any{"inlineData": map[string]any{
+					"mimeType": "image/png",
+					"data":     "aW1hZ2U=",
+				}},
+			}},
+		},
+	})
+	if err != nil {
+		t.Fatalf("map Gemini request: %v", err)
+	}
+	if mapping.Kind != "images.edits" || mapping.Body["modelType"] != "image_edit" {
+		t.Fatalf("unexpected mapping kind/body: %+v", mapping)
+	}
+	if mapping.Body["prompt"] != "make it brighter" || mapping.Body["image"] != "data:image/png;base64,aW1hZ2U=" {
+		t.Fatalf("unexpected mapped body: %+v", mapping.Body)
+	}
+}
+
+func TestGeminiGenerateContentResponseMapsURLAndUsage(t *testing.T) {
+	response := geminiGenerateContentResponse(map[string]any{
+		"data": []any{map[string]any{
+			"type": "image",
+			"url":  "https://cdn.example/out.webp",
+		}},
+		"usage": map[string]any{
+			"prompt_tokens":     3,
+			"completion_tokens": 4,
+			"total_tokens":      7,
+		},
+	}, "gemini-image")
+
+	candidates := response["candidates"].([]any)
+	content := candidates[0].(map[string]any)["content"].(map[string]any)
+	parts := content["parts"].([]any)
+	fileData := parts[0].(map[string]any)["fileData"].(map[string]any)
+	if fileData["fileUri"] != "https://cdn.example/out.webp" || fileData["mimeType"] != "image/webp" {
+		t.Fatalf("unexpected fileData part: %+v", fileData)
+	}
+	usage := response["usageMetadata"].(map[string]any)
+	if usage["promptTokenCount"] != 3 || usage["candidatesTokenCount"] != 4 || usage["totalTokenCount"] != 7 {
+		t.Fatalf("unexpected usage metadata: %+v", usage)
+	}
+}
+
+func TestGeminiFileResponseShape(t *testing.T) {
+	response := geminiFileResponse("files/test", "https://cdn.example/input.png", "image/png", 12)
+	file := response["file"].(map[string]any)
+	if file["name"] != "files/test" || file["uri"] != "https://cdn.example/input.png" || file["mimeType"] != "image/png" || file["sizeBytes"] != "12" || file["state"] != "ACTIVE" {
+		t.Fatalf("unexpected file response: %+v", file)
+	}
+}