feat(piagent): support image quality and output_format in image client

Add quality and output_format to ImageRequestOptions, plumb them through the JSON/multipart request bodies, generalize asset URL collection to non-image formats, and add GenerateLayeredDocument for PSD output. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-09 20:42:55 +08:00
parent cb24ffd185
commit 41208e566c
3 changed files with 288 additions and 59 deletions
@@ -11,6 +11,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"img_infinite_canvas/internal/domain/design"
 )

 func TestNewImageClientDefaultsToTenMinuteTimeout(t *testing.T) {
@@ -145,36 +147,140 @@ func TestImageClientSendsExplicitSize(t *testing.T) {
 	}
 }

-func TestImageClientSendsImageCount(t *testing.T) {
+func TestImageClientSendsQuality(t *testing.T) {
 	var requestBody map[string]any
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
 			t.Fatal(err)
 		}
 		w.Header().Set("Content-Type", "application/json")
-		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image-1.png"},{"url":"https://example.com/image-2.png"}]}`))
+		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/image.png"}]}`))
 	}))
 	defer server.Close()

 	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
-	image, err := client.Generate(context.Background(), "create two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2})
-	if err != nil {
+	if _, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2", Quality: "high"}); err != nil {
 		t.Fatal(err)
 	}
-	if requestBody["n"] != float64(2) {
-		t.Fatalf("expected image request n=2, got %#v", requestBody["n"])
-	}
-	if requestBody["stream"] != nil {
-		t.Fatalf("expected multi-image request to avoid stream, got %#v", requestBody["stream"])
-	}
-	if len(image.URLs) != 2 {
-		t.Fatalf("expected two generated images, got %#v", image.URLs)
+	if requestBody["quality"] != "high" {
+		t.Fatalf("expected quality=high, got %#v", requestBody["quality"])
 	}
 }

-func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported(t *testing.T) {
+func TestImageClientSendsOutputFormatAndParsesPSDBase64(t *testing.T) {
+	var requestBody map[string]any
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"b64_json":"cHNkLWJ5dGVz"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create layered file", ImageRequestOptions{Model: "gpt-image-2", OutputFormat: "psd"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestBody["output_format"] != "psd" {
+		t.Fatalf("expected output_format=psd, got %#v", requestBody["output_format"])
+	}
+	if requestBody["stream"] != nil {
+		t.Fatalf("expected psd output to avoid streaming, got %#v", requestBody["stream"])
+	}
+	if image.URL != "data:image/vnd.adobe.photoshop;base64,cHNkLWJ5dGVz" {
+		t.Fatalf("expected psd data url, got %#v", image)
+	}
+}
+
+func TestImageClientGenerateLayeredDocumentUsesPSDOutput(t *testing.T) {
+	var requestBody map[string]any
+	var requestPath string
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		requestPath = r.URL.Path
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/layers.psd"}]}`))
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key", InputImageTransport: "url"})
+	document, err := client.GenerateLayeredDocument(context.Background(), design.LayeredDocumentRequest{
+		ImageURL: "https://example.com/source.png",
+		Prompt:   "separate layers",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if requestPath != "/v1/images/edits" {
+		t.Fatalf("expected image edit endpoint, got %s", requestPath)
+	}
+	if requestBody["output_format"] != "psd" || requestBody["size"] != "auto" {
+		t.Fatalf("expected psd auto request, got %#v", requestBody)
+	}
+	if document.Content != "https://example.com/layers.psd" || document.ContentType != "image/vnd.adobe.photoshop" {
+		t.Fatalf("unexpected layered document: %#v", document)
+	}
+}
+
+func TestImageClientSplitsImageCountIntoStreamedSingleRequests(t *testing.T) {
+	var mu sync.Mutex
+	requestBodies := make([]map[string]any, 0, 2)
+	requestKeys := make([]string, 0, 2)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		key := r.Header.Get("Idempotency-Key")
+		imageIndex := strings.TrimPrefix(key, "batch-key-")
+		if imageIndex == key {
+			imageIndex = "unknown"
+		}
+		var requestBody map[string]any
+		if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil {
+			t.Fatal(err)
+		}
+		mu.Lock()
+		requestBodies = append(requestBodies, requestBody)
+		requestKeys = append(requestKeys, key)
+		mu.Unlock()
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = fmt.Fprintf(w, `{"data":[{"url":"https://example.com/image-%s.png"}]}`, imageIndex)
+	}))
+	defer server.Close()
+
+	client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"})
+	image, err := client.Generate(context.Background(), "create two images", ImageRequestOptions{Model: "gpt-image-2", Count: 2, IdempotencyKey: "batch-key"})
+	if err != nil {
+		t.Fatal(err)
+	}
+	mu.Lock()
+	defer mu.Unlock()
+	if len(requestBodies) != 2 {
+		t.Fatalf("expected two single-image requests, got %d", len(requestBodies))
+	}
+	for index, body := range requestBodies {
+		if body["n"] != nil {
+			t.Fatalf("expected split request %d to omit n, got %#v", index, body)
+		}
+		if body["stream"] != true {
+			t.Fatalf("expected split request %d to keep stream=true, got %#v", index, body)
+		}
+	}
+	keySet := map[string]bool{}
+	for _, key := range requestKeys {
+		keySet[key] = true
+	}
+	if !keySet["batch-key-1"] || !keySet["batch-key-2"] {
+		t.Fatalf("expected suffixed idempotency keys, got %#v", requestKeys)
+	}
+	if len(image.URLs) != 2 || image.URLs[0] != "https://example.com/image-1.png" || image.URLs[1] != "https://example.com/image-2.png" {
+		t.Fatalf("expected ordered generated images, got %#v", image.URLs)
+	}
+}
+
+func TestImageClientSplitsImageCountRequestsConcurrently(t *testing.T) {
 	var mu sync.Mutex
-	var initialBody map[string]any
 	fallbackBodies := make([]map[string]any, 0, 3)
 	fallbackKeys := make([]string, 0, 3)
 	activeFallbacks := 0
@@ -188,13 +294,6 @@ func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported
 			http.Error(w, err.Error(), http.StatusBadRequest)
 			return
 		}
-		if requestBody["n"] != nil {
-			mu.Lock()
-			initialBody = requestBody
-			mu.Unlock()
-			http.Error(w, `{"error":{"message":"Unknown parameter: 'tools[0].n'.","param":"tools[0].n","type":"invalid_request_error"}}`, http.StatusBadRequest)
-			return
-		}

 		mu.Lock()
 		fallbackBodies = append(fallbackBodies, requestBody)
@@ -229,22 +328,19 @@ func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported
 	}
 	mu.Lock()
 	defer mu.Unlock()
-	if initialBody["n"] != float64(3) {
-		t.Fatalf("expected first request n=3, got %#v", initialBody["n"])
-	}
 	if len(image.URLs) != 3 {
-		t.Fatalf("expected three fallback images, got %#v", image.URLs)
+		t.Fatalf("expected three generated images, got %#v", image.URLs)
 	}
 	if len(fallbackBodies) != 3 {
-		t.Fatalf("expected three single-image fallback requests, got %d", len(fallbackBodies))
+		t.Fatalf("expected three single-image requests, got %d", len(fallbackBodies))
 	}
 	if maxActiveFallbacks < 2 {
-		t.Fatalf("expected concurrent fallback requests, max active was %d", maxActiveFallbacks)
+		t.Fatalf("expected concurrent split requests, max active was %d", maxActiveFallbacks)
 	}
 	keySet := map[string]bool{}
 	for index, body := range fallbackBodies {
 		if body["n"] != nil {
-			t.Fatalf("expected fallback request %d to omit n, got %#v", index, body)
+			t.Fatalf("expected split request %d to omit n, got %#v", index, body)
 		}
 		keySet[fallbackKeys[index]] = true
 	}