package piagent import ( "bufio" "bytes" "context" "encoding/base64" "encoding/json" "errors" "fmt" "io" "mime/multipart" "net" "net/http" "net/textproto" "net/url" "path/filepath" "strconv" "strings" "time" "unicode/utf8" "img_infinite_canvas/internal/domain/design" "img_infinite_canvas/internal/infrastructure/netutil" "github.com/zeromicro/go-zero/core/logx" ) const ( defaultImageGenerationTimeout = 10 * time.Minute defaultImageGenerationResponseHeaderTimeout = 4 * time.Minute maxImageGenerationTimeout = 10 * time.Minute imageGenerationStreamMaxLineBytes = 64 << 20 ) type ImageGenerator interface { Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) } type GeneratedImage struct { URL string URLs []string } type ImageRequestOptions struct { Model string Size string Count int Images []string IdempotencyKey string } type ImageClientOptions struct { BaseURL string APIKey string InputImageTransport string TimeoutSeconds int ResponseHeaderTimeoutSeconds int } type ImageClient struct { baseURL string apiKey string inputImageTransport string client *http.Client } func NewImageClient(opts ImageClientOptions) *ImageClient { if strings.TrimSpace(opts.BaseURL) == "" || strings.TrimSpace(opts.APIKey) == "" { return nil } timeout := time.Duration(opts.TimeoutSeconds) * time.Second if timeout <= 0 { timeout = defaultImageGenerationTimeout } if timeout > maxImageGenerationTimeout { timeout = maxImageGenerationTimeout } responseHeaderTimeout := time.Duration(opts.ResponseHeaderTimeoutSeconds) * time.Second if responseHeaderTimeout <= 0 { responseHeaderTimeout = defaultImageGenerationResponseHeaderTimeout } if responseHeaderTimeout > timeout { responseHeaderTimeout = timeout } return &ImageClient{ baseURL: strings.TrimRight(opts.BaseURL, "/"), apiKey: opts.APIKey, inputImageTransport: normalizeInputImageTransport(opts.InputImageTransport), client: &http.Client{ Timeout: timeout, Transport: &http.Transport{ Proxy: http.ProxyFromEnvironment, DialContext: (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext, DisableKeepAlives: true, ForceAttemptHTTP2: false, ResponseHeaderTimeout: responseHeaderTimeout, TLSHandshakeTimeout: 10 * time.Second, ExpectContinueTimeout: time.Second, }, }, } } func normalizeInputImageTransport(value string) string { switch strings.ToLower(strings.TrimSpace(value)) { case "url", "image_url", "image-url": return "url" default: return "file" } } func (c *ImageClient) Generate(ctx context.Context, prompt string, opts ImageRequestOptions) (GeneratedImage, error) { prompt = strings.TrimSpace(prompt) if prompt == "" { return GeneratedImage{}, fmt.Errorf("image prompt is required") } model := strings.TrimSpace(opts.Model) if model == "" { model = "gpt-image-2" } size := strings.TrimSpace(opts.Size) if size == "" { size = design.DefaultGeneratedImageSize } count := opts.Count if count < 1 { count = 1 } if count > 10 { count = 10 } images := normalizeImageInputs(opts.Images) idempotencyKey := strings.TrimSpace(opts.IdempotencyKey) streamPlans := []bool{false} if len(images) == 0 { streamPlans = []bool{true, false} } var lastErr error startedAt := time.Now() promptRunes := utf8.RuneCountInString(prompt) 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, retryable, terminal, idempotencyKey, err, ) if !retryable || attempt == 2 { break } select { case <-ctx.Done(): return GeneratedImage{}, ctx.Err() case <-time.After(time.Duration(attempt+1) * 1500 * time.Millisecond): } } if !(stream && shouldFallbackImageGenerationStream(lastErr, idempotencyKey)) { break } } return GeneratedImage{}, lastErr } 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) { return jsonImageEditBody(model, prompt, size, count, images, stream) } return c.multipartImageEditBody(ctx, model, prompt, size, count, images, stream) } requestPayload := map[string]any{ "model": model, "prompt": prompt, } if count > 1 { requestPayload["n"] = count } if size != "" { requestPayload["size"] = size } if stream { requestPayload["stream"] = true } body, err := json.Marshal(requestPayload) return body, "application/json", err } func imageInputsCanUseURLTransport(images []string) bool { for _, image := range images { if imageInputRequiresFileUpload(image) { return false } } return true } func imageInputRequiresFileUpload(value string) bool { value = strings.TrimSpace(value) if value == "" { return false } if strings.HasPrefix(value, "data:image/") || strings.HasPrefix(value, "/") { return true } parsed, err := url.Parse(value) if err != nil { return true } scheme := strings.ToLower(parsed.Scheme) if scheme != "http" && scheme != "https" { return true } host := strings.ToLower(parsed.Hostname()) if host == "" { return true } return netutil.IsLocalOrPrivateHost(host) } 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, "images": imageReferencePayload(images), } if count > 1 { requestPayload["n"] = count } 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, stream bool) ([]byte, string, error) { var body bytes.Buffer writer := multipart.NewWriter(&body) fields := map[string]string{ "model": model, "prompt": prompt, } if size != "" { fields["size"] = size } 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 } } fieldName := "image" if len(images) > 1 { fieldName = "image[]" } for index, image := range images { input, err := c.imageInputFile(ctx, image, index) if err != nil { _ = writer.Close() return nil, "", err } part, err := createImageFormFile(writer, fieldName, input) if err != nil { _ = writer.Close() return nil, "", err } if _, err := part.Write(input.Data); err != nil { _ = writer.Close() return nil, "", err } } if err := writer.Close(); err != nil { return nil, "", err } return body.Bytes(), writer.FormDataContentType(), nil } func createImageFormFile(writer *multipart.Writer, fieldName string, input imageInputFile) (io.Writer, error) { header := make(textproto.MIMEHeader) header.Set("Content-Disposition", fmt.Sprintf(`form-data; name="%s"; filename="%s"`, escapeMultipartValue(fieldName), escapeMultipartValue(input.FileName))) contentType := strings.TrimSpace(input.ContentType) if contentType == "" { contentType = "application/octet-stream" } header.Set("Content-Type", contentType) return writer.CreatePart(header) } func escapeMultipartValue(value string) string { return strings.NewReplacer("\\", "\\\\", `"`, "\\\"").Replace(value) } type imageInputFile struct { FileName string ContentType string Data []byte } func (c *ImageClient) imageInputFile(ctx context.Context, value string, index int) (imageInputFile, error) { value = strings.TrimSpace(value) if strings.HasPrefix(value, "data:image/") { header, payload, ok := strings.Cut(value, ",") if !ok { return imageInputFile{}, fmt.Errorf("invalid image data url") } contentType := strings.TrimPrefix(strings.Split(header, ";")[0], "data:") data, err := base64.StdEncoding.DecodeString(payload) if err != nil { return imageInputFile{}, err } return imageInputFile{ FileName: fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, "")), ContentType: contentType, Data: data, }, nil } if !strings.HasPrefix(value, "http://") && !strings.HasPrefix(value, "https://") { return imageInputFile{}, fmt.Errorf("unsupported image input %q", value) } req, err := http.NewRequestWithContext(ctx, http.MethodGet, value, nil) if err != nil { return imageInputFile{}, err } req.Header.Set("User-Agent", "img-infinite-canvas/1.0") resp, err := c.client.Do(req) if err != nil { return imageInputFile{}, err } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { return imageInputFile{}, fmt.Errorf("download reference image failed: %s", resp.Status) } contentType := strings.TrimSpace(strings.Split(resp.Header.Get("Content-Type"), ";")[0]) data, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20)) if err != nil { return imageInputFile{}, err } if len(data) == 0 { return imageInputFile{}, fmt.Errorf("download reference image returned empty body") } if contentType == "" || !strings.HasPrefix(contentType, "image/") { contentType = http.DetectContentType(data) } if !strings.HasPrefix(contentType, "image/") { return imageInputFile{}, fmt.Errorf("reference input is not an image: %s", contentType) } return imageInputFile{ FileName: imageFileName(value, contentType, index), ContentType: contentType, Data: data, }, nil } func normalizeImageInputs(images []string) []string { if len(images) == 0 { return nil } normalized := make([]string, 0, len(images)) seen := make(map[string]struct{}) for _, image := range images { image = strings.TrimSpace(image) if image == "" { continue } if _, ok := seen[image]; ok { continue } seen[image] = struct{}{} normalized = append(normalized, image) if len(normalized) >= 16 { break } } return normalized } func imageReferencePayload(images []string) []map[string]string { items := make([]map[string]string, 0, len(images)) for _, image := range images { items = append(items, map[string]string{"image_url": image}) } return items } func imageFileName(value string, contentType string, index int) string { name := "" if parsed, err := url.Parse(value); err == nil { name = filepath.Base(parsed.Path) } if name == "" || name == "." || name == "/" { name = fmt.Sprintf("reference-%d%s", index+1, imageExtension(contentType, "")) } if filepath.Ext(name) == "" { name += imageExtension(contentType, ".png") } return name } func imageExtension(contentType string, fallback string) string { switch strings.ToLower(strings.TrimSpace(contentType)) { case "image/png": return ".png" case "image/jpeg", "image/jpg": return ".jpg" case "image/webp": return ".webp" case "image/gif": return ".gif" default: if fallback != "" { return fallback } return ".png" } } 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" } req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+endpoint, bytes.NewReader(body)) if err != nil { return GeneratedImage{}, err } 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) } req.Close = true resp, err := c.client.Do(req) if err != nil { return GeneratedImage{}, err } defer resp.Body.Close() if resp.StatusCode < 200 || resp.StatusCode >= 300 { respBody, readErr := io.ReadAll(io.LimitReader(resp.Body, 4<<20)) if readErr != nil { return GeneratedImage{}, readErr } 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 } urls := collectGeneratedImageURLs(responsePayload) if len(urls) == 0 { return GeneratedImage{}, fmt.Errorf("image generation returned empty image") } 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{}) var walk func(any, string) walk = func(value any, key string) { switch typed := value.(type) { case map[string]any: for childKey, childValue := range typed { walk(childValue, childKey) } case []any: for _, childValue := range typed { walk(childValue, key) } case string: if imageURL := generatedImageString(key, typed); imageURL != "" { if _, ok := seen[imageURL]; ok { return } seen[imageURL] = struct{}{} urls = append(urls, imageURL) } } } walk(payload, "") return urls } func generatedImageString(key string, value string) string { value = strings.TrimSpace(value) if value == "" { return "" } if strings.HasPrefix(value, "data:image/") { return value } normalizedKey := strings.ToLower(strings.TrimSpace(key)) if isGeneratedImageURLKey(normalizedKey) && (strings.HasPrefix(value, "http://") || strings.HasPrefix(value, "https://")) { return value } if isGeneratedImageBase64Key(normalizedKey) { return "data:image/png;base64," + strings.TrimPrefix(value, "data:image/png;base64,") } return "" } func isGeneratedImageURLKey(key string) bool { switch key { case "url", "image_url", "imageurl", "image", "output_url", "public_url", "publicurl": return true default: return false } } func isGeneratedImageBase64Key(key string) bool { switch key { 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 } if isImageGenerationTerminalError(err) { return false } message := strings.ToLower(err.Error()) return strings.Contains(message, "eof") || strings.Contains(message, "connection reset") || strings.Contains(message, "unexpected end") || strings.Contains(message, "temporarily unavailable") || strings.Contains(message, "too many requests") || strings.Contains(message, "rate limit") || strings.Contains(message, "429") || strings.Contains(message, "500") || strings.Contains(message, "502") || strings.Contains(message, "503") || strings.Contains(message, "504") } func isImageGenerationTerminalError(err error) bool { if err == nil { return false } if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return true } return isImageGenerationTimeoutError(err) } func isImageGenerationTimeoutError(err error) bool { if err == nil { return false } var netErr net.Error if errors.As(err, &netErr) && netErr.Timeout() { return true } message := strings.ToLower(err.Error()) return strings.Contains(message, "client.timeout exceeded") || strings.Contains(message, "context deadline exceeded") || strings.Contains(message, "deadline exceeded") || strings.Contains(message, "request canceled") || strings.Contains(message, "request cancelled") || strings.Contains(message, "timeout awaiting response headers") || strings.Contains(message, "timeout while awaiting headers") }