137 lines
3.6 KiB
Go
137 lines
3.6 KiB
Go
package agent
|
|
|
|
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("万相任务超时")
|
|
}
|