From 7679854b5de33d1efc09ba7a4493c316f3da01ea Mon Sep 17 00:00:00 2001 From: liangxu Date: Wed, 8 Jul 2026 12:06:39 +0800 Subject: [PATCH] feat(image): stream text-to-image generation with SSE and fallback Send stream=true for text-to-image requests and consume the SSE response, falling back to a non-streamed request when the stream errors in a way that warrants it. Raises the scanner line limit to handle large base64 image payloads. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../piagent/image_generation.go | 339 +++++++++++++++--- .../piagent/image_generation_test.go | 109 ++++++ 2 files changed, 399 insertions(+), 49 deletions(-) diff --git a/server/internal/infrastructure/piagent/image_generation.go b/server/internal/infrastructure/piagent/image_generation.go index c4e6373..2996274 100644 --- a/server/internal/infrastructure/piagent/image_generation.go +++ b/server/internal/infrastructure/piagent/image_generation.go @@ -1,6 +1,7 @@ package piagent import ( + "bufio" "bytes" "context" "encoding/base64" @@ -28,6 +29,7 @@ const ( defaultImageGenerationTimeout = 10 * time.Minute defaultImageGenerationResponseHeaderTimeout = 4 * time.Minute maxImageGenerationTimeout = 10 * time.Minute + imageGenerationStreamMaxLineBytes = 64 << 20 ) type ImageGenerator interface { @@ -132,70 +134,95 @@ func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageReq images := normalizeImageInputs(opts.Images) idempotencyKey := strings.TrimSpace(opts.IdempotencyKey) - body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images) - if err != nil { - return GeneratedImage{}, err + streamPlans := []bool{false} + if len(images) == 0 { + streamPlans = []bool{true, false} } var lastErr error startedAt := time.Now() promptRunes := utf8.RuneCountInString(prompt) - for attempt := 0; attempt < 3; attempt++ { - image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0) - if err == nil { - logx.Infof( - "gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s images=%d idempotency_key=%s", + for _, stream := range streamPlans { + body, contentType, err := c.imageRequestBody(ctx, model, prompt, size, count, images, stream) + if err != nil { + return GeneratedImage{}, err + } + for attempt := 0; attempt < 3; attempt++ { + image, err := c.generateOnce(ctx, body, contentType, idempotencyKey, len(images) > 0, stream, count) + if err == nil { + logx.Infof( + "gpt image generation succeeded model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s images=%d idempotency_key=%s", + model, + size, + count, + promptRunes, + attempt+1, + stream, + time.Since(startedAt).Round(time.Millisecond), + c.client.Timeout, + len(image.URLs), + idempotencyKey, + ) + return image, nil + } + if ctx.Err() != nil { + return GeneratedImage{}, ctx.Err() + } + lastErr = err + if stream && isImageGenerationStreamFallbackError(err) { + logx.Errorf( + "gpt image generation stream fallback model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s idempotency_key=%s err=%v", + model, + size, + count, + promptRunes, + attempt+1, + time.Since(startedAt).Round(time.Millisecond), + idempotencyKey, + err, + ) + break + } + timeoutRetryable := idempotencyKey != "" && isImageGenerationTimeoutError(err) + retryable := isRetryableImageGenerationError(err) || timeoutRetryable + terminal := isImageGenerationTerminalError(err) && !timeoutRetryable + logx.Errorf( + "gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d stream=%t elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v", model, size, count, promptRunes, attempt+1, + stream, time.Since(startedAt).Round(time.Millisecond), c.client.Timeout, - len(image.URLs), + retryable, + terminal, idempotencyKey, + err, ) - return image, nil + if !retryable || attempt == 2 { + break + } + select { + case <-ctx.Done(): + return GeneratedImage{}, ctx.Err() + case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond): + } } - if ctx.Err() != nil { - return GeneratedImage{}, ctx.Err() - } - lastErr = err - timeoutRetryable := idempotencyKey != "" && isImageGenerationTimeoutError(err) - retryable := isRetryableImageGenerationError(err) || timeoutRetryable - terminal := isImageGenerationTerminalError(err) && !timeoutRetryable - logx.Errorf( - "gpt image generation failed model=%s size=%s count=%d prompt_runes=%d attempt=%d elapsed=%s timeout=%s retryable=%t terminal=%t idempotency_key=%s err=%v", - model, - size, - count, - promptRunes, - attempt+1, - time.Since(startedAt).Round(time.Millisecond), - c.client.Timeout, - retryable, - terminal, - idempotencyKey, - err, - ) - if !retryable || attempt == 2 { + if !(stream && shouldFallbackImageGenerationStream(lastErr, idempotencyKey)) { break } - select { - case <-ctx.Done(): - return GeneratedImage{}, ctx.Err() - case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond): - } } return GeneratedImage{}, lastErr } -func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) { +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" { - return jsonImageEditBody(model, prompt, size, count, images) + return jsonImageEditBody(model, prompt, size, count, images, stream) } - return c.multipartImageEditBody(ctx, model, prompt, size, count, images) + return c.multipartImageEditBody(ctx, model, prompt, size, count, images, stream) } requestPayload := map[string]any{ @@ -208,11 +235,14 @@ func (c *ImageClient) imageRequestBody(ctx context.Context, model string, prompt if size != "" { requestPayload["size"] = size } + if stream { + requestPayload["stream"] = true + } body, err := json.Marshal(requestPayload) return body, "application/json", err } -func jsonImageEditBody(model string, prompt string, size string, count int, images []string) ([]byte, string, error) { +func jsonImageEditBody(model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) { requestPayload := map[string]any{ "model": model, "prompt": prompt, @@ -224,11 +254,14 @@ func jsonImageEditBody(model string, prompt string, size string, count int, imag if size != "" { requestPayload["size"] = size } + if stream { + requestPayload["stream"] = true + } body, err := json.Marshal(requestPayload) return body, "application/json", err } -func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string) ([]byte, string, error) { +func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, prompt string, size string, count int, images []string, stream bool) ([]byte, string, error) { var body bytes.Buffer writer := multipart.NewWriter(&body) fields := map[string]string{ @@ -241,6 +274,9 @@ func (c *ImageClient) multipartImageEditBody(ctx context.Context, model string, if count > 1 { fields["n"] = strconv.Itoa(count) } + if stream { + fields["stream"] = "true" + } for key, value := range fields { if err := writer.WriteField(key, value); err != nil { return nil, "", err @@ -413,7 +449,7 @@ func imageExtension(contentType string, fallback string) string { } } -func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool) (GeneratedImage, error) { +func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType string, idempotencyKey string, edit bool, stream bool, expectedCount int) (GeneratedImage, error) { endpoint := "/v1/images/generations" if edit { endpoint = "/v1/images/edits" @@ -425,6 +461,9 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType req.Header.Set("Authorization", "Bearer "+c.apiKey) req.Header.Set("Content-Type", contentType) req.Header.Set("Connection", "close") + if stream { + req.Header.Set("Accept", "text/event-stream") + } if idempotencyKey != "" { req.Header.Set("Idempotency-Key", idempotencyKey) } @@ -437,16 +476,21 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { - var detail struct { - Error any `json:"error"` + respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) + if readErr != nil { + return GeneratedImage{}, readErr } - _ = json.NewDecoder(resp.Body).Decode(&detail) - if detail.Error != nil { - return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %v", resp.Status, detail.Error) + detail := imageErrorBody(respBody) + if detail != "" { + return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %s", resp.Status, detail) } return GeneratedImage{}, fmt.Errorf("image generation failed: %s", resp.Status) } + if stream { + return decodeImageGenerationStreamOrJSON(resp.Body, expectedCount) + } + var responsePayload any if err := json.NewDecoder(resp.Body).Decode(&responsePayload); err != nil { return GeneratedImage{}, err @@ -459,6 +503,180 @@ func (c *ImageClient) generateOnce(ctx context.Context, body []byte, contentType return GeneratedImage{URL: urls[0], URLs: urls}, nil } +type imageStreamEvent struct { + name string + dataLines []string + hasData bool +} + +func decodeImageGenerationStreamOrJSON(body io.Reader, expectedCount int) (GeneratedImage, error) { + if expectedCount <= 0 { + expectedCount = 1 + } + scanner := bufio.NewScanner(body) + scanner.Buffer(make([]byte, 0, 64*1024), imageGenerationStreamMaxLineBytes) + var raw strings.Builder + event := imageStreamEvent{name: "message"} + sawEvents := false + urls := make([]string, 0, expectedCount) + seen := make(map[string]struct{}) + appendURLs := func(values []string) { + for _, value := range values { + if value == "" { + continue + } + if _, ok := seen[value]; ok { + continue + } + seen[value] = struct{}{} + urls = append(urls, value) + } + } + flush := func() (bool, error) { + if !event.hasData { + event = imageStreamEvent{name: "message"} + return false, nil + } + sawEvents = true + name := event.name + dataText := strings.TrimSpace(strings.Join(event.dataLines, "\n")) + event = imageStreamEvent{name: "message"} + if dataText == "" { + return false, nil + } + if dataText == "[DONE]" { + return true, nil + } + if strings.Contains(strings.ToLower(name), "partial_image") { + return false, nil + } + var payload any + if err := json.Unmarshal([]byte(dataText), &payload); err != nil { + return false, nil + } + if isPartialImagePayload(payload) { + return false, nil + } + if message := imageStreamFailureMessage(payload); message != "" { + return true, fmt.Errorf("image generation stream failed: %s", message) + } + appendURLs(collectGeneratedImageURLs(payload)) + return len(urls) >= expectedCount, nil + } + + for scanner.Scan() { + line := strings.TrimSuffix(scanner.Text(), "\r") + trimmed := strings.TrimSpace(line) + switch { + case line == "": + done, err := flush() + if err != nil { + return GeneratedImage{}, err + } + if done && len(urls) > 0 { + return GeneratedImage{URL: urls[0], URLs: urls}, nil + } + case strings.HasPrefix(line, ":"): + continue + case strings.HasPrefix(line, "event:"): + event.name = strings.TrimSpace(strings.TrimPrefix(line, "event:")) + case strings.HasPrefix(line, "data:"): + event.hasData = true + event.dataLines = append(event.dataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:"))) + default: + if !sawEvents && trimmed != "" { + raw.WriteString(line) + raw.WriteByte('\n') + } + } + } + if err := scanner.Err(); err != nil { + return GeneratedImage{}, err + } + done, err := flush() + if err != nil { + return GeneratedImage{}, err + } + if len(urls) > 0 { + return GeneratedImage{URL: urls[0], URLs: urls}, nil + } + if sawEvents || done { + return GeneratedImage{}, fmt.Errorf("image generation stream returned no complete image") + } + + var responsePayload any + if err := json.Unmarshal([]byte(raw.String()), &responsePayload); err != nil { + return GeneratedImage{}, err + } + urls = collectGeneratedImageURLs(responsePayload) + if len(urls) == 0 { + return GeneratedImage{}, fmt.Errorf("image generation returned empty image") + } + return GeneratedImage{URL: urls[0], URLs: urls}, nil +} + +func imageErrorBody(body []byte) string { + body = bytes.TrimSpace(body) + if len(body) == 0 { + return "" + } + var detail struct { + Error any `json:"error"` + } + if err := json.Unmarshal(body, &detail); err == nil && detail.Error != nil { + return fmt.Sprint(detail.Error) + } + return trimErrorBody(body) +} + +func isPartialImagePayload(payload any) bool { + item, ok := payload.(map[string]any) + if !ok { + return false + } + if value, ok := item["type"].(string); ok && strings.Contains(strings.ToLower(value), "partial_image") { + return true + } + return false +} + +func imageStreamFailureMessage(payload any) string { + item, ok := payload.(map[string]any) + if !ok { + return "" + } + if message := imageErrorMessage(item["error"]); message != "" { + return message + } + typeValue, _ := item["type"].(string) + statusValue, _ := item["status"].(string) + if strings.Contains(strings.ToLower(typeValue), "failed") || strings.EqualFold(statusValue, "failed") { + if message, ok := item["message"].(string); ok && strings.TrimSpace(message) != "" { + return strings.TrimSpace(message) + } + return "upstream reported failed status" + } + return "" +} + +func imageErrorMessage(value any) string { + switch typed := value.(type) { + case nil: + return "" + case string: + return strings.TrimSpace(typed) + case map[string]any: + for _, key := range []string{"message", "detail", "type", "code"} { + if message, ok := typed[key].(string); ok && strings.TrimSpace(message) != "" { + return strings.TrimSpace(message) + } + } + return fmt.Sprint(typed) + default: + return fmt.Sprint(typed) + } +} + func collectGeneratedImageURLs(payload any) []string { urls := make([]string, 0) seen := make(map[string]struct{}) @@ -516,13 +734,36 @@ func isGeneratedImageURLKey(key string) bool { func isGeneratedImageBase64Key(key string) bool { switch key { - case "b64_json", "b64json", "base64", "image_base64", "imagebase64": + case "b64_json", "b64json", "base64", "image_base64", "imagebase64", "result": return true default: return false } } +func isImageGenerationStreamFallbackError(err error) bool { + if err == nil { + return false + } + message := strings.ToLower(err.Error()) + return strings.Contains(message, "unknown parameter") && strings.Contains(message, "stream") || + strings.Contains(message, "unsupported") && strings.Contains(message, "stream") || + strings.Contains(message, "not support") && strings.Contains(message, "stream") || + strings.Contains(message, "stream returned no complete image") || + strings.Contains(message, "incomplete stream") || + strings.Contains(message, "text/event-stream") +} + +func shouldFallbackImageGenerationStream(err error, idempotencyKey string) bool { + if isImageGenerationStreamFallbackError(err) { + return true + } + if strings.TrimSpace(idempotencyKey) == "" { + return false + } + return isRetryableImageGenerationError(err) || isImageGenerationTimeoutError(err) +} + func isRetryableImageGenerationError(err error) bool { if err == nil { return false diff --git a/server/internal/infrastructure/piagent/image_generation_test.go b/server/internal/infrastructure/piagent/image_generation_test.go index bd40e5f..a9b392f 100644 --- a/server/internal/infrastructure/piagent/image_generation_test.go +++ b/server/internal/infrastructure/piagent/image_generation_test.go @@ -97,6 +97,9 @@ func TestImageClientSendsTransportDefaultSizeWhenUnset(t *testing.T) { if requestBody["size"] != "1024x1024" { t.Fatalf("expected transport default size, got %#v", requestBody["size"]) } + if requestBody["stream"] != true { + t.Fatalf("expected image generations to use stream transport, got %#v", requestBody["stream"]) + } if _, ok := requestBody["n"]; ok { t.Fatalf("expected single image request to omit n for compatible transports, got %#v", requestBody["n"]) } @@ -306,6 +309,112 @@ func TestImageClientParsesAlternateImageResponseShapes(t *testing.T) { } } +func TestImageClientParsesStreamingImageResponse(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", "text/event-stream") + _, _ = w.Write([]byte("event: image_generation.completed\n")) + _, _ = w.Write([]byte(`data: {"type":"image_generation.completed","result":"Z2VuZXJhdGVkIGJ5dGVz"}` + "\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"}) + image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{Model: "gpt-image-2"}) + if err != nil { + t.Fatal(err) + } + if requestBody["stream"] != true { + t.Fatalf("expected stream request body, got %#v", requestBody) + } + if image.URL != "data:image/png;base64,Z2VuZXJhdGVkIGJ5dGVz" { + t.Fatalf("expected streamed base64 image, got %#v", image) + } +} + +func TestImageClientFallsBackWhenStreamIsUnsupported(t *testing.T) { + calls := 0 + streamValues := make([]any, 0, 2) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + var requestBody map[string]any + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Fatal(err) + } + streamValues = append(streamValues, requestBody["stream"]) + if calls == 1 { + http.Error(w, `{"error":{"message":"unknown parameter: stream"}}`, http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = 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 an image", ImageRequestOptions{Model: "gpt-image-2"}) + if err != nil { + t.Fatal(err) + } + if image.URL != "https://example.com/image.png" { + t.Fatalf("expected fallback image, got %#v", image) + } + if calls != 2 { + t.Fatalf("expected stream then json fallback, got %d calls", calls) + } + if len(streamValues) != 2 || streamValues[0] != true || streamValues[1] != nil { + t.Fatalf("expected first request stream=true and fallback without stream, got %#v", streamValues) + } +} + +func TestImageClientFallsBackAfterIdempotentStreamFailures(t *testing.T) { + calls := 0 + streamValues := make([]any, 0, 4) + keys := make([]string, 0, 4) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + calls++ + keys = append(keys, r.Header.Get("Idempotency-Key")) + var requestBody map[string]any + if err := json.NewDecoder(r.Body).Decode(&requestBody); err != nil { + t.Fatal(err) + } + streamValues = append(streamValues, requestBody["stream"]) + if calls <= 3 { + http.Error(w, `{"error":{"message":"temporary stream failure"}}`, http.StatusServiceUnavailable) + return + } + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"url":"https://example.com/recovered.png"}]}`)) + })) + defer server.Close() + + client := NewImageClient(ImageClientOptions{BaseURL: server.URL, APIKey: "test-key"}) + image, err := client.Generate(context.Background(), "create an image", ImageRequestOptions{ + Model: "gpt-image-2", + IdempotencyKey: "same-generation", + }) + if err != nil { + t.Fatal(err) + } + if image.URL != "https://example.com/recovered.png" { + t.Fatalf("expected fallback image, got %#v", image) + } + if calls != 4 { + t.Fatalf("expected three stream attempts then json fallback, got %d calls", calls) + } + if len(streamValues) != 4 || streamValues[0] != true || streamValues[1] != true || streamValues[2] != true || streamValues[3] != nil { + t.Fatalf("expected stream retries then json fallback, got %#v", streamValues) + } + for _, key := range keys { + if key != "same-generation" { + t.Fatalf("expected idempotency key to be preserved, got %#v", keys) + } + } +} + func TestImageClientRetriesTransientFailures(t *testing.T) { calls := 0 server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {