moteva/server/internal/infrastructure/piagent/image_generation.go

package piagent

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"

	"img_infinite_canvas/internal/domain/design"

	"github.com/zeromicro/go-zero/core/logx"
)

const (
	defaultImageGenerationTimeout = 10 * time.Minute
	maxImageGenerationTimeout     = 10 * time.Minute
)

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
	IdempotencyKey string
}

type ImageClientOptions struct {
	BaseURL        string
	APIKey         string
	TimeoutSeconds int
}

type ImageClient struct {
	baseURL string
	apiKey  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
	}

	return &ImageClient{
		baseURL: strings.TrimRight(opts.BaseURL, "/"),
		apiKey:  opts.APIKey,
		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,
				TLSHandshakeTimeout:   10 * time.Second,
				ExpectContinueTimeout: time.Second,
			},
		},
	}
}

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
	}

	requestPayload := map[string]any{
		"model":  model,
		"prompt": prompt,
	}
	if count > 1 {
		requestPayload["n"] = count
	}
	if size != "" {
		requestPayload["size"] = size
	}
	idempotencyKey := strings.TrimSpace(opts.IdempotencyKey)

	body, err := json.Marshal(requestPayload)
	if err != nil {
		return GeneratedImage{}, err
	}

	var lastErr error
	startedAt := time.Now()
	promptRunes := utf8.RuneCountInString(prompt)
	for attempt := 0; attempt < 3; attempt++ {
		image, err := c.generateOnce(ctx, body, idempotencyKey)
		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",
				model,
				size,
				count,
				promptRunes,
				attempt+1,
				time.Since(startedAt).Round(time.Millisecond),
				c.client.Timeout,
				len(image.URLs),
				idempotencyKey,
			)
			return image, nil
		}
		lastErr = err
		retryable := isRetryableImageGenerationError(err)
		terminal := isImageGenerationTerminalError(err)
		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 {
			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) generateOnce(ctx context.Context, body []byte, idempotencyKey string) (GeneratedImage, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/v1/images/generations", bytes.NewReader(body))
	if err != nil {
		return GeneratedImage{}, err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Connection", "close")
	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 {
		var detail struct {
			Error any `json:"error"`
		}
		_ = json.NewDecoder(resp.Body).Decode(&detail)
		if detail.Error != nil {
			return GeneratedImage{}, fmt.Errorf("image generation failed: %s: %v", resp.Status, detail.Error)
		}
		return GeneratedImage{}, fmt.Errorf("image generation failed: %s", resp.Status)
	}

	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
}

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":
		return true
	default:
		return false
	}
}

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
	}
	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")
}