slogan/styleagent/imagegen/wanx_client.go

package imagegen

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

// wanxClient 通义万相人像写真（image-synthesis 异步接口 + 轮询）
type wanxClient struct {
	apiKey string
	model  string
	base   string
}

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

type wanxInput struct {
	Prompt        string `json:"prompt"`
	BaseImageURL  string `json:"base_image_url,omitempty"`
	BaseImagePath string `json:"base_image_path,omitempty"`
}

type wanxResp struct {
	Output struct {
		TaskID    string `json:"task_id"`
		TaskStatus string `json:"task_status"`
		Results   []struct {
			URL string `json:"url"`
		} `json:"results"`
	} `json:"output"`
	Code      string `json:"code"`
	Message   string `json:"message"`
}

type wanxTaskResp struct {
	Output struct {
		TaskStatus string `json:"task_status"`
		Results    []struct {
			URL string `json:"url"`
		} `json:"results"`
	} `json:"output"`
	Code    string `json:"code"`
	Message string `json:"message"`
}

// Generate 提交任务并轮询直到完成，失败返回错误（由上层降级 mock）
func (c *wanxClient) Generate(ctx context.Context, req *GenerateReq) (string, error) {
	body, err := json.Marshal(wanxSubmitReq{
		Model: c.model,
		Input: wanxInput{Prompt: buildPrompt(req.Prompt, "", "", req.Angle), BaseImageURL: req.BaseImageURL},
		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
}

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 resp.Body.Close()
	data, _ := io.ReadAll(resp.Body)
	var r wanxResp
	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) {
	taskURL := c.base + "?task_id=" + taskID
	client := &http.Client{Timeout: 30 * time.Second}
	for i := 0; i < 60; i++ {
		select {
		case <-ctx.Done():
			return "", ctx.Err()
		case <-time.After(5 * time.Second):
		}
		req, err := http.NewRequestWithContext(ctx, "GET", taskURL, nil)
		if err != nil {
			return "", err
		}
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
		resp, err := client.Do(req)
		if err != nil {
			return "", err
		}
		data, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		var r wanxTaskResp
		if err := json.Unmarshal(data, &r); err != nil {
			return "", fmt.Errorf("万相任务查询解析失败: %s", string(data))
		}
		switch r.Output.TaskStatus {
		case "SUCCEEDED":
			if len(r.Output.Results) > 0 && r.Output.Results[0].URL != "" {
				return r.Output.Results[0].URL, nil
			}
			return "", fmt.Errorf("万相任务成功但无结果")
		case "FAILED":
			return "", fmt.Errorf("万相任务失败: %s", r.Message)
		}
	}
	return "", fmt.Errorf("万相任务超时")
}