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