observer/server/common/imagegen.go

package common

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"github.com/gogf/gf/v2/frame/g"
)

// ImageGenProvider 图像生成 provider 抽象：管理端「AI 生成图片」按 config.yml imageGen.provider 选择实现。
type ImageGenProvider interface {
	// Generate 生成一张图片，返回图片字节；prompt 由调用方保证不含位置描述（项目提示词规范）。
	Generate(ctx context.Context, prompt, size string) ([]byte, error)
}

// ImageGen 当前配置的图像生成 provider 单例（未配置返回 nil，调用方判 CodeImageGenNotConfigured）。
func ImageGen(ctx context.Context) ImageGenProvider {
	if g.Cfg().MustGet(ctx, "imageGen.apiKey").String() == "" {
		return nil
	}
	switch g.Cfg().MustGet(ctx, "imageGen.provider", "dashscope").String() {
	case "dashscope":
		return &dashScopeImageGen{apiKey: g.Cfg().MustGet(ctx, "imageGen.apiKey").String()}
	default:
		return nil
	}
}

// dashScopeImageGen 通义万相（DashScope）实现：提交异步任务 → 轮询 task 状态 → 下载产物图片。
// 文档：https://help.aliyun.com/zh/model-studio/text-to-image-api-reference
type dashScopeImageGen struct {
	apiKey string
}

const (
	dashScopeSynthUrl   = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text2image/image-synthesis"
	dashScopeTaskUrlFmt = "https://dashscope.aliyuncs.com/api/v1/tasks/%s"
)

func (p *dashScopeImageGen) Generate(ctx context.Context, prompt, size string) ([]byte, error) {
	model := g.Cfg().MustGet(ctx, "imageGen.model", "qwen-image-3.0").String()
	// 前端尺寸格式 1152x2048 → API 规格 1152*2048
	apiSize := strings.ReplaceAll(size, "x", "*")
	body, err := json.Marshal(map[string]any{
		"model":      model,
		"input":      map[string]string{"prompt": prompt, "size": apiSize},
		"parameters": map[string]any{"n": 1, "watermark": false},
	})
	if err != nil {
		return nil, err
	}
	resp, err := p.post(ctx, dashScopeSynthUrl, body)
	if err != nil {
		return nil, err
	}
	taskId, ok := resp["task_id"].(string)
	if !ok || taskId == "" {
		return nil, fmt.Errorf("DashScope 提交失败: %v", resp)
	}
	// 轮询任务结果：异步生成通常 10~60s，上限 120s（与设计一致：超时 2min/张）
	deadline := time.Now().Add(120 * time.Second)
	for {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("DashScope 生成超时(120s)")
		}
		task, err := p.post(ctx, fmt.Sprintf(dashScopeTaskUrlFmt, taskId), nil)
		if err != nil {
			return nil, err
		}
		status, _ := task["task_status"].(string)
		switch status {
		case "SUCCEEDED":
			results, _ := task["results"].([]any)
			if len(results) == 0 {
				return nil, fmt.Errorf("DashScope 成功但无产物图片")
			}
			url, _ := results[0].(map[string]any)["url"].(string)
			if url == "" {
				return nil, fmt.Errorf("DashScope 成功但无图片 URL")
			}
			return p.download(ctx, url)
		case "FAILED", "CANCELED":
			msg, _ := task["message"].(string)
			return nil, fmt.Errorf("DashScope 生成失败: %s", msg)
		}
		select {
		case <-time.After(2 * time.Second):
		case <-ctx.Done():
			return nil, ctx.Err()
		}
	}
}

// post 调用 DashScope HTTP 接口并解析统一 JSON（无 body 时为空 GET，轮询任务用）
func (p *dashScopeImageGen) post(ctx context.Context, url string, body []byte) (map[string]any, error) {
	method := http.MethodPost
	var rd io.Reader
	if body == nil {
		method = http.MethodGet
	} else {
		rd = strings.NewReader(string(body))
	}
	req, err := http.NewRequestWithContext(ctx, method, url, rd)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+p.apiKey)
	if body != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() { _ = resp.Body.Close() }()
	raw, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	if err != nil {
		return nil, err
	}
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("DashScope HTTP %d: %s", resp.StatusCode, truncateStr(string(raw), 200))
	}
	var m map[string]any
	if err := json.Unmarshal(raw, &m); err != nil {
		return nil, err
	}
	return m, nil
}

// download 下载生成产物图片（存于阿里云 OSS，无需鉴权）
func (p *dashScopeImageGen) download(ctx context.Context, url string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
	if err != nil {
		return nil, err
	}
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return nil, err
	}
	defer func() { _ = resp.Body.Close() }()
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("下载生成图片失败: HTTP %d", resp.StatusCode)
	}
	return io.ReadAll(io.LimitReader(resp.Body, 16<<20))
}

func truncateStr(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}