diff --git a/server/internal/infrastructure/piagent/image_generation.go b/server/internal/infrastructure/piagent/image_generation.go index 2dd2d49..78d70d2 100644 --- a/server/internal/infrastructure/piagent/image_generation.go +++ b/server/internal/infrastructure/piagent/image_generation.go @@ -17,6 +17,7 @@ import ( "path/filepath" "strconv" "strings" + "sync" "time" "unicode/utf8" @@ -215,9 +216,83 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq break } } + if count > 1 && isImageGenerationCountFallbackError(lastErr) { + logx.Errorf( + "gpt image generation count fallback model=%s size=%s count=%d prompt_runes=%d elapsed=%s idempotency_key=%s err=%v", + model, + size, + count, + promptRunes, + time.Since(startedAt).Round(time.Millisecond), + idempotencyKey, + lastErr, + ) + return c.generateImageCountFallback(ctx, prompt, model, size, count, images, idempotencyKey) + } return GeneratedImage{}, lastErr } +func (c *ImageClient) generateImageCountFallback(ctx context.Context, prompt string, model string, size string, count int, images []string, idempotencyKey string) (GeneratedImage, error) { + fallbackCtx, cancel := context.WithCancel(ctx) + defer cancel() + + urls := make([]string, count) + var wg sync.WaitGroup + var mu sync.Mutex + var firstErr error + setErr := func(err error) { + if err == nil { + return + } + mu.Lock() + defer mu.Unlock() + if firstErr == nil { + firstErr = err + cancel() + } + } + + for index := 0; index < count; index++ { + index := index + wg.Add(1) + go func() { + defer wg.Done() + fallbackOpts := ImageRequestOptions{ + Model: model, + Size: size, + Count: 1, + Images: append([]string(nil), images...), + } + if idempotencyKey != "" { + fallbackOpts.IdempotencyKey = fmt.Sprintf("%s-%d", idempotencyKey, index+1) + } + image, err := c.Generate(fallbackCtx, prompt, fallbackOpts) + if err != nil { + setErr(err) + return + } + url := image.URL + if url == "" && len(image.URLs) > 0 { + url = image.URLs[0] + } + if url == "" { + setErr(fmt.Errorf("image generation count fallback returned empty image")) + return + } + urls[index] = url + }() + } + wg.Wait() + + mu.Lock() + err := firstErr + mu.Unlock() + if err != nil { + return GeneratedImage{}, err + } + return GeneratedImage{URL: urls[0], URLs: urls}, nil +} + func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) { if len(images) > 0 { if c.inputImageTransport == "url" && imageInputsCanUseURLTransport(images) { @@ -787,6 +862,21 @@ func isImageGenerationStreamFallbackError(err error) bool { strings.Contains(message, "text/event-stream") } +func isImageGenerationCountFallbackError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + if !strings.Contains(message, "unknown parameter") { + return false + } + return strings.Contains(message, "tools[0].n") || + strings.Contains(message, "param:n") || + strings.Contains(message, "param: n") || + strings.Contains(message, "parameter: 'n'") || + strings.Contains(message, `parameter: "n"`) +} + func shouldFallbackImageGenerationStream(err error, idempotencyKey string) bool { if isImageGenerationStreamFallbackError(err) { return true diff --git a/server/internal/infrastructure/piagent/image_generation_test.go b/server/internal/infrastructure/piagent/image_generation_test.go index 082a948..0ec63f3 100644 --- a/server/internal/infrastructure/piagent/image_generation_test.go +++ b/server/internal/infrastructure/piagent/image_generation_test.go @@ -3,10 +3,12 @@ package piagent import ( "context" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" "strings" + "sync" "testing" "time" ) @@ -170,6 +172,89 @@ func TestImageClientSendsImageCount(t *testing.T) { } } +func TestImageClientFallsBackToConcurrentSingleImageRequestsWhenCountUnsupported(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 + maxActiveFallbacks := 0 + releaseFallbacks := make(chan struct{}) + + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + var requestBody map[string]any + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Errorf("decode request body: %v", err) + 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) + fallbackKeys = append(fallbackKeys, r.Header.Get("Idempotency-Key")) + index := len(fallbackBodies) + activeFallbacks++ + if activeFallbacks > maxActiveFallbacks { + maxActiveFallbacks = activeFallbacks + } + if index == 3 { + close(releaseFallbacks) + } + mu.Unlock() + + select { + case <-releaseFallbacks: + case <-time.After(2 * time.Second): + } + + mu.Lock() + activeFallbacks-- + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"data":[{"url":"https://example.com/fallback-%d.png"}]}`, index) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"}) + image, err := client.Generate(context.Background(), "create three images", ImageRequestOptions{Model: "gpt-image-2", Count: 3, IdempotencyKey: "batch-key"}) + if err != nil { + t.Fatal(err) + } + 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) + } + if len(fallbackBodies) != 3 { + t.Fatalf("expected three single-image fallback requests, got %d", len(fallbackBodies)) + } + if maxActiveFallbacks < 2 { + t.Fatalf("expected concurrent fallback 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) + } + keySet[fallbackKeys[index]] = true + } + for _, key := range []string{"batch-key-1", "batch-key-2", "batch-key-3"} { + if !keySet[key] { + t.Fatalf("expected fallback idempotency key %q, got %#v", key, fallbackKeys) + } + } +} + func TestImageClientUploadsDataReferenceImagesAsMultipart(t *testing.T) { var requestPath string var requestContentType string