slogan/server/styleagent/agent/imagegen_wanx_client.go

package agent

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"
)

// wanxClient 通义万相图像生成（wan2.7-image-pro：image-generation 异步接口 + 轮询）
type wanxClient struct {
	apiKey   string
	model    string
	base     string // 任务提交端点
	taskBase string // 任务查询端点
}

type wanxSubmitReq struct {
	Model      string         `json:"model"`
	Input      wanxInput      `json:"input"`
	Parameters map[string]any `json:"parameters"`
}

type wanxInput struct {
	Messages []wanxMessage `json:"messages"`
}

type wanxMessage struct {
	Role    string             `json:"role"`
	Content []wanxInputContent `json:"content"`
}

// 请求侧 content 元素：wan2.7-image-pro 原生格式直接放 text / image 字段，
// 不能用 OpenAI 兼容的 {"type":"image_url"} 形式（会报 "Either 'text' or 'image' must be provided, but not both"）
type wanxInputContent struct {
	Text  string `json:"text,omitempty"`
	Image string `json:"image,omitempty"`
}

// 响应侧 content 元素（图片在 image 字段）
type wanxContent struct {
	Type  string `json:"type"`
	Text  string `json:"text,omitempty"`
	Image string `json:"image,omitempty"`
}

type wanxTaskResp struct {
	Output struct {
		TaskStatus string `json:"task_status"`
		Message    string `json:"message"`
		Code       string `json:"code"`
		Choices    []struct {
			Message struct {
				Content []wanxContent `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	} `json:"output"`
}

// Generate 文生图或图生图（BaseImageURL 本地路径转 data URI，http(s) 直传），异步任务 + 轮询
// 注：wan2.7-image-pro 图生图需同一条 user 消息中并列 {"text"} 与 {"image"} 两个 content 元素
func (c *wanxClient) Generate(ctx context.Context, req *GenerateReq) (string, error) {
	content := []wanxInputContent{{Text: buildPrompt(req.Prompt, "", "", req.Angle)}}
	if req.BaseImageURL != "" {
		imgURL, err := resolveImageURL(req.BaseImageURL)
		if err != nil {
			return "", err
		}
		content = append(content, wanxInputContent{Image: imgURL})
	}
	messages := []wanxMessage{{Role: "user", Content: content}}

	body, err := json.Marshal(wanxSubmitReq{
		Model:      c.model,
		Input:      wanxInput{Messages: messages},
		Parameters: map[string]any{"n": 1, "size": "768*1024", "seed": req.Seed},
	})
	if err != nil {
		return "", err
	}

	taskID, err := c.submit(ctx, body)
	if err != nil {
		return "", err
	}
	url, err := c.poll(ctx, taskID)
	if err != nil {
		return "", err
	}
	return url, nil
}

// resolveImageURL 本地 /workspace 路径转 data URI（dashscope 无法访问相对路径），http(s) 原样返回
func resolveImageURL(raw string) (string, error) {
	if strings.HasPrefix(raw, "http://") || strings.HasPrefix(raw, "https://") {
		return raw, nil
	}
	path := strings.TrimPrefix(raw, "/")
	data, err := os.ReadFile(path)
	if err != nil {
		return "", fmt.Errorf("读取参考图失败 %s: %w", raw, err)
	}
	ext := "png"
	if i := strings.LastIndex(path, "."); i >= 0 {
		ext = strings.TrimPrefix(path[i+1:], ".")
	}
	return fmt.Sprintf("data:image/%s;base64,%s", ext, base64.StdEncoding.EncodeToString(data)), nil
}

func (c *wanxClient) submit(ctx context.Context, body []byte) (string, error) {
	req, err := http.NewRequestWithContext(ctx, "POST", c.base, bytes.NewReader(body))
	if err != nil {
		return "", err
	}
	req.Header.Set("Authorization", "Bearer "+c.apiKey)
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("X-DashScope-Async", "enable")
	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer func() { _ = resp.Body.Close() }()
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("读取万相提交响应失败: %w", err)
	}
	var r struct {
		Output struct {
			TaskID string `json:"task_id"`
		} `json:"output"`
		Code    string `json:"code"`
		Message string `json:"message"`
	}
	if err := json.Unmarshal(data, &r); err != nil {
		return "", fmt.Errorf("万相提交响应解析失败: %s", string(data))
	}
	if r.Output.TaskID == "" {
		return "", fmt.Errorf("万相提交失败 code=%s msg=%s", r.Code, r.Message)
	}
	return r.Output.TaskID, nil
}

func (c *wanxClient) poll(ctx context.Context, taskID string) (string, error) {
	client := &http.Client{Timeout: 30 * time.Second}
	for i := 0; i < 120; i++ {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(5 * time.Second):
		}
		req, err := http.NewRequestWithContext(ctx, "GET", c.taskBase+"/"+taskID, nil)
		if err != nil {
			return "", err
		}
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
		resp, err := client.Do(req)
		if err != nil {
			return "", err
		}
		data, err := io.ReadAll(resp.Body)
		resp.Body.Close()
		if err != nil {
			return "", fmt.Errorf("读取万相任务响应失败: %w", err)
		}
		var r wanxTaskResp
		if err := json.Unmarshal(data, &r); err != nil {
			return "", fmt.Errorf("万相任务查询解析失败: %s", string(data))
		}
		switch r.Output.TaskStatus {
		case "SUCCEEDED":
			for _, ch := range r.Output.Choices {
				for _, ct := range ch.Message.Content {
					if ct.Type == "image" && ct.Image != "" {
						return ct.Image, nil
					}
				}
			}
			return "", fmt.Errorf("万相任务成功但无结果")
		case "FAILED":
			return "", fmt.Errorf("万相任务失败: %s", r.Output.Message)
		}
	}
	return "", fmt.Errorf("万相任务超时")
}