root/moteva
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(image): stream text-to-image generation with SSE and fallback

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
liangxu
2026-07-08 12:06:39 +08:00
parent 7e05a8d47b
commit 7679854b5d
2 changed files with 399 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server/internal/infrastructure/piagent/image_generation.go
+290 -49
View File
@@ -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
server/internal/infrastructure/piagent/image_generation_test.go
+109
View File
@@ -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) {