初始化
This commit is contained in:
Ran
@@ -0,0 +1,64 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type BalanceResult struct {
|
||||
RemainBalance float64 `json:"remain_balance"`
|
||||
RemainCredits float64 `json:"remain_credits"`
|
||||
UsedBalance float64 `json:"used_balance"`
|
||||
UsedCredits float64 `json:"used_credits"`
|
||||
UnlimitedQuota bool `json:"unlimited_quota"`
|
||||
}
|
||||
|
||||
func (c Client) Balance(ctx context.Context, baseURL, apiKey string) (BalanceResult, error) {
|
||||
var result BalanceResult
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(baseURL, "/")+"/user/balance", nil)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
|
||||
client := c.HTTPClient
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
response, err := client.Do(request)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 1024*1024))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return result, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
|
||||
var decoded struct {
|
||||
Success bool `json:"success"`
|
||||
Message string `json:"message"`
|
||||
BalanceResult
|
||||
}
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return result, fmt.Errorf("余额响应不是有效 JSON: %w", err)
|
||||
}
|
||||
if !decoded.Success {
|
||||
if decoded.Message == "" {
|
||||
return result, errors.New("获取渠道余额失败")
|
||||
}
|
||||
return result, errors.New(decoded.Message)
|
||||
}
|
||||
return decoded.BalanceResult, nil
|
||||
}
|
||||
@@ -0,0 +1,42 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestBalanceUsesConfiguredBaseURLAndAPIKey(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/v1/user/balance" {
|
||||
t.Fatalf("unexpected path: %s", request.URL.Path)
|
||||
}
|
||||
if request.Header.Get("Authorization") != "Bearer test-key" {
|
||||
t.Fatalf("unexpected authorization header: %s", request.Header.Get("Authorization"))
|
||||
}
|
||||
fmt.Fprint(response, `{"success":true,"remain_balance":10.5,"remain_credits":105.25,"used_balance":2.3,"used_credits":2005.21388,"unlimited_quota":false}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).Balance(context.Background(), server.URL+"/v1", "test-key")
|
||||
if err != nil {
|
||||
t.Fatalf("Balance returned error: %v", err)
|
||||
}
|
||||
if result.RemainBalance != 10.5 || result.RemainCredits != 105.25 || result.UsedBalance != 2.3 || result.UsedCredits != 2005.21388 || result.UnlimitedQuota {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBalanceReturnsUpstreamMessage(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
|
||||
fmt.Fprint(response, `{"success":false,"message":"record not found"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
_, err := NewClient(server.Client()).Balance(context.Background(), server.URL, "test-key")
|
||||
if err == nil || err.Error() != "record not found" {
|
||||
t.Fatalf("expected upstream message, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type ChatResult struct {
|
||||
Content string
|
||||
Raw json.RawMessage
|
||||
FinishReason string
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
UsageRaw json.RawMessage
|
||||
}
|
||||
|
||||
func (c Client) Chat(ctx context.Context, baseURL, apiKey string, payload map[string]any) (string, json.RawMessage, error) {
|
||||
result, err := c.ChatWithUsage(ctx, baseURL, apiKey, payload)
|
||||
return result.Content, result.Raw, err
|
||||
}
|
||||
|
||||
func (c Client) ChatWithUsage(ctx context.Context, baseURL, apiKey string, payload map[string]any) (ChatResult, error) {
|
||||
var result ChatResult
|
||||
requestPayload := make(map[string]any, len(payload)+1)
|
||||
for key, value := range payload {
|
||||
requestPayload[key] = value
|
||||
}
|
||||
requestPayload["stream"] = false
|
||||
body, err := json.Marshal(requestPayload)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(baseURL, "/")+"/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
response, err := c.HTTPClient.Do(request)
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 16*1024*1024))
|
||||
if err != nil {
|
||||
return result, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
return result, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
type chatChoice struct {
|
||||
Message struct {
|
||||
Content any `json:"content"`
|
||||
} `json:"message"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
}
|
||||
var decoded struct {
|
||||
Choices []chatChoice `json:"choices"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
Data struct {
|
||||
Choices []chatChoice `json:"choices"`
|
||||
Usage json.RawMessage `json:"usage"`
|
||||
} `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
return result, fmt.Errorf("文本模型响应不是有效 JSON: %w；上游响应: %s", err, safeBody(data))
|
||||
}
|
||||
result.Raw = append(json.RawMessage(nil), data...)
|
||||
result.UsageRaw = decoded.Usage
|
||||
choices := decoded.Choices
|
||||
if len(choices) == 0 {
|
||||
choices = decoded.Data.Choices
|
||||
if len(result.UsageRaw) == 0 {
|
||||
result.UsageRaw = decoded.Data.Usage
|
||||
}
|
||||
}
|
||||
if len(choices) == 0 {
|
||||
return result, fmt.Errorf("文本模型响应缺少 choices；上游响应: %s", safeBody(data))
|
||||
}
|
||||
result.FinishReason = choices[0].FinishReason
|
||||
if len(result.UsageRaw) > 0 {
|
||||
var usage map[string]any
|
||||
if json.Unmarshal(result.UsageRaw, &usage) == nil {
|
||||
result.InputTokens = intValue(usage["prompt_tokens"], usage["input_tokens"])
|
||||
result.OutputTokens = intValue(usage["completion_tokens"], usage["output_tokens"])
|
||||
result.TotalTokens = intValue(usage["total_tokens"])
|
||||
if result.TotalTokens == 0 {
|
||||
result.TotalTokens = result.InputTokens + result.OutputTokens
|
||||
}
|
||||
}
|
||||
}
|
||||
switch content := choices[0].Message.Content.(type) {
|
||||
case string:
|
||||
result.Content = content
|
||||
return result, nil
|
||||
case []any:
|
||||
parts := make([]string, 0)
|
||||
for _, item := range content {
|
||||
if block, ok := item.(map[string]any); ok {
|
||||
if text, ok := block["text"].(string); ok {
|
||||
parts = append(parts, text)
|
||||
}
|
||||
}
|
||||
}
|
||||
result.Content = strings.Join(parts, "\n")
|
||||
return result, nil
|
||||
default:
|
||||
return result, errors.New("文本模型响应内容无效")
|
||||
}
|
||||
}
|
||||
|
||||
func intValue(values ...any) int64 {
|
||||
for _, value := range values {
|
||||
switch number := value.(type) {
|
||||
case float64:
|
||||
if number > 0 {
|
||||
return int64(number)
|
||||
}
|
||||
case int64:
|
||||
if number > 0 {
|
||||
return number
|
||||
}
|
||||
case json.Number:
|
||||
if parsed, err := number.Int64(); err == nil && parsed > 0 {
|
||||
return parsed
|
||||
}
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -0,0 +1,58 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
func TestChatWithUsageReadsOpenAICompatibleAndNestedUsage(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
response string
|
||||
content string
|
||||
input int64
|
||||
output int64
|
||||
total int64
|
||||
finish string
|
||||
}{
|
||||
{name: "top level", response: `{"choices":[{"message":{"content":"result"},"finish_reason":"stop"}],"usage":{"prompt_tokens":120,"completion_tokens":30,"total_tokens":150}}`, content: "result", input: 120, output: 30, total: 150, finish: "stop"},
|
||||
{name: "nested alternate names", response: `{"data":{"choices":[{"message":{"content":[{"type":"text","text":"part one"},{"type":"text","text":"part two"}]}}],"usage":{"input_tokens":44,"output_tokens":11}}}`, content: "part one\npart two", input: 44, output: 11, total: 55},
|
||||
}
|
||||
for _, test := range cases {
|
||||
t.Run(test.name, func(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.URL.Path != "/chat/completions" {
|
||||
t.Fatalf("unexpected path: %s", request.URL.Path)
|
||||
}
|
||||
response.Header().Set("Content-Type", "application/json")
|
||||
fmt.Fprint(response, test.response)
|
||||
}))
|
||||
defer server.Close()
|
||||
result, err := NewClient(server.Client()).ChatWithUsage(context.Background(), server.URL, "key", map[string]any{"model": "text-model"})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatWithUsage returned error: %v", err)
|
||||
}
|
||||
if result.Content != test.content || result.InputTokens != test.input || result.OutputTokens != test.output || result.TotalTokens != test.total || result.FinishReason != test.finish {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChatWithUsageReturnsExplicitUpstreamFailure(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, _ *http.Request) {
|
||||
response.WriteHeader(http.StatusBadGateway)
|
||||
fmt.Fprint(response, `{"error":"relay unavailable"}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
_, err := NewClient(server.Client()).ChatWithUsage(context.Background(), server.URL, "key", map[string]any{"model": "text-model"})
|
||||
httpErr, ok := err.(*provider.HTTPError)
|
||||
if !ok || httpErr.StatusCode != http.StatusBadGateway {
|
||||
t.Fatalf("expected HTTPError with relay status, got %T %v", err, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,64 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type Client struct {
|
||||
HTTPClient *http.Client
|
||||
}
|
||||
|
||||
func NewClient(client *http.Client) Client {
|
||||
return Client{HTTPClient: client}
|
||||
}
|
||||
|
||||
func (c Client) requestJSON(ctx context.Context, method, baseURL, path, apiKey, requestID string, payload map[string]any) ([]byte, error) {
|
||||
var body io.Reader
|
||||
if payload != nil {
|
||||
encoded, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body = bytes.NewReader(encoded)
|
||||
}
|
||||
request, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(baseURL, "/")+path, body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
request.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
if payload != nil {
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
if requestID != "" {
|
||||
request.Header.Set("Idempotency-Key", requestID)
|
||||
request.Header.Set("X-Request-ID", requestID)
|
||||
}
|
||||
response, err := c.HTTPClient.Do(request)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer response.Body.Close()
|
||||
data, err := io.ReadAll(io.LimitReader(response.Body, 4*1024*1024))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if response.StatusCode < 200 || response.StatusCode >= 300 {
|
||||
return nil, &provider.HTTPError{StatusCode: response.StatusCode, Body: safeBody(data), RetryAfter: response.Header.Get("Retry-After")}
|
||||
}
|
||||
return data, nil
|
||||
}
|
||||
|
||||
func safeBody(data []byte) string {
|
||||
value := strings.TrimSpace(string(data))
|
||||
if len(value) > 2000 {
|
||||
value = value[:2000]
|
||||
}
|
||||
return value
|
||||
}
|
||||
@@ -0,0 +1,103 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
func (c Client) SubmitImage(ctx context.Context, baseURL, apiKey, requestID string, payload map[string]any) (provider.SubmitResult, error) {
|
||||
data, err := c.requestJSON(ctx, http.MethodPost, baseURL, "/images/generations", apiKey, requestID, payload)
|
||||
if err != nil {
|
||||
return provider.SubmitResult{}, err
|
||||
}
|
||||
result := provider.SubmitResult{Raw: append(json.RawMessage(nil), data...)}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.SubmitResult{}, errors.New("中转站提交响应不是有效 JSON")
|
||||
}
|
||||
result.TaskID = findString(decoded, "task_id", "id", "request_id")
|
||||
result.URL = findURL(decoded)
|
||||
if result.TaskID == "" && result.URL == "" {
|
||||
return provider.SubmitResult{}, errors.New("中转站提交响应缺少任务 ID 或结果 URL")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
// PollImage 通过 APIMart 统一任务接口查询图片生成状态和结果。
|
||||
func (c Client) PollImage(ctx context.Context, baseURL, apiKey, taskID string) (provider.PollResult, error) {
|
||||
data, err := c.requestJSON(ctx, http.MethodGet, baseURL, "/tasks/"+taskID, apiKey, "", nil)
|
||||
if err != nil {
|
||||
return provider.PollResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.PollResult{}, errors.New("中转站轮询响应不是有效 JSON")
|
||||
}
|
||||
status := strings.ToLower(findString(decoded, "task_status", "status", "state"))
|
||||
if status == "" {
|
||||
status = "processing"
|
||||
}
|
||||
return provider.PollResult{Status: status, URL: findURL(decoded), Error: findString(decoded, "error_message", "error", "message", "detail"), Raw: append(json.RawMessage(nil), data...)}, nil
|
||||
}
|
||||
|
||||
func findURL(value any) string {
|
||||
var walk func(any, string) string
|
||||
walk = func(node any, parentKey string) string {
|
||||
switch data := node.(type) {
|
||||
case string:
|
||||
if (strings.Contains(strings.ToLower(parentKey), "url") || strings.HasPrefix(data, "http://") || strings.HasPrefix(data, "https://")) && (strings.HasPrefix(data, "https://") || strings.HasPrefix(data, "http://")) {
|
||||
return data
|
||||
}
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if found := walk(child, key); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if found := walk(child, parentKey); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
return walk(value, "")
|
||||
}
|
||||
|
||||
func findString(value any, keys ...string) string {
|
||||
wanted := map[string]bool{}
|
||||
for _, key := range keys {
|
||||
wanted[strings.ToLower(key)] = true
|
||||
}
|
||||
return findMatchingString(value, func(key, _ string) bool { return wanted[strings.ToLower(key)] })
|
||||
}
|
||||
|
||||
func findMatchingString(value any, match func(key, value string) bool) string {
|
||||
switch data := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if text, ok := child.(string); ok && match(key, text) {
|
||||
return text
|
||||
}
|
||||
}
|
||||
for _, child := range data {
|
||||
if found := findMatchingString(child, match); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if found := findMatchingString(child, match); found != "" {
|
||||
return found
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
// 图片生成中转站适配器测试，验证任务提交、状态轮询及成功结果解析契约。
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestSubmitImageUsesGenerationEndpoint 验证图片生成使用配置渠道下的标准提交接口。
|
||||
func TestSubmitImageUsesGenerationEndpoint(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.Method != http.MethodPost || request.URL.Path != "/v1/images/generations" {
|
||||
t.Fatalf("unexpected request: %s %s", request.Method, request.URL.Path)
|
||||
}
|
||||
fmt.Fprint(response, `{"code":200,"data":[{"status":"submitted","task_id":"task-image-1"}]}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).SubmitImage(context.Background(), server.URL+"/v1", "test-key", "request-1", map[string]any{"model": "gpt-image-2"})
|
||||
if err != nil {
|
||||
t.Fatalf("SubmitImage returned error: %v", err)
|
||||
}
|
||||
if result.TaskID != "task-image-1" {
|
||||
t.Fatalf("unexpected task id: %s", result.TaskID)
|
||||
}
|
||||
}
|
||||
|
||||
// TestPollImageUsesTaskEndpoint 验证图片任务使用 APIMart 统一任务接口查询并解析完成结果。
|
||||
func TestPollImageUsesTaskEndpoint(t *testing.T) {
|
||||
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
|
||||
if request.Method != http.MethodGet || request.URL.Path != "/v1/tasks/task-image-1" {
|
||||
t.Fatalf("unexpected request: %s %s", request.Method, request.URL.Path)
|
||||
}
|
||||
fmt.Fprint(response, `{"code":200,"data":{"id":"task-image-1","status":"completed","result":{"images":[{"url":["https://example.com/result.png"]}]}}}`)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
result, err := NewClient(server.Client()).PollImage(context.Background(), server.URL+"/v1", "test-key", "task-image-1")
|
||||
if err != nil {
|
||||
t.Fatalf("PollImage returned error: %v", err)
|
||||
}
|
||||
if result.Status != "completed" || result.URL != "https://example.com/result.png" {
|
||||
t.Fatalf("unexpected result: %#v", result)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"juhe-factory/api/internal/provider"
|
||||
)
|
||||
|
||||
type Seedance2 struct {
|
||||
Gateway Client
|
||||
}
|
||||
|
||||
type ReferenceCapability struct {
|
||||
Supported bool `json:"supported"`
|
||||
MaxCount int `json:"max_count"`
|
||||
MinTotalDurationMS int64 `json:"min_total_duration_ms,omitempty"`
|
||||
MaxTotalDurationMS int64 `json:"max_total_duration_ms,omitempty"`
|
||||
MinShortSidePixels int `json:"min_short_side_pixels,omitempty"`
|
||||
MaxShortSidePixels int `json:"max_short_side_pixels,omitempty"`
|
||||
}
|
||||
|
||||
type VideoModelCapabilities struct {
|
||||
// DurationMinSeconds / DurationMaxSeconds 表示视频模型支持的生成时长范围（秒），供前端可视化选择与后端校验。
|
||||
DurationMinSeconds int `json:"duration_min_seconds"`
|
||||
DurationMaxSeconds int `json:"duration_max_seconds"`
|
||||
ReferenceImages ReferenceCapability `json:"reference_images"`
|
||||
ReferenceAudios ReferenceCapability `json:"reference_audios"`
|
||||
ReferenceVideos ReferenceCapability `json:"reference_videos"`
|
||||
}
|
||||
|
||||
var seedance2Capabilities = map[string]VideoModelCapabilities{
|
||||
"doubao-seedance-2.0": seedance2ReferenceCapabilities(),
|
||||
"doubao-seedance-2.0-fast": seedance2ReferenceCapabilities(),
|
||||
"doubao-seedance-2.0-mini": seedance2ReferenceCapabilities(),
|
||||
}
|
||||
|
||||
func seedance2ReferenceCapabilities() VideoModelCapabilities {
|
||||
return VideoModelCapabilities{
|
||||
DurationMinSeconds: 5,
|
||||
DurationMaxSeconds: 15,
|
||||
ReferenceImages: ReferenceCapability{Supported: true, MaxCount: 9},
|
||||
ReferenceAudios: ReferenceCapability{Supported: true, MaxCount: 3, MaxTotalDurationMS: 15_000},
|
||||
ReferenceVideos: ReferenceCapability{Supported: true, MaxCount: 3, MinTotalDurationMS: 1_800, MaxTotalDurationMS: 15_200, MinShortSidePixels: 480, MaxShortSidePixels: 720},
|
||||
}
|
||||
}
|
||||
|
||||
func VideoCapabilities(modelName string) (VideoModelCapabilities, bool) {
|
||||
capabilities, ok := seedance2Capabilities[strings.TrimSpace(modelName)]
|
||||
return capabilities, ok
|
||||
}
|
||||
|
||||
func NewSeedance2(client *http.Client) Seedance2 {
|
||||
return Seedance2{Gateway: NewClient(client)}
|
||||
}
|
||||
|
||||
func BuildSeedance2Payload(modelName string, inputData json.RawMessage) (map[string]any, error) {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
capabilities, supported := VideoCapabilities(modelName)
|
||||
if !supported {
|
||||
return nil, fmt.Errorf("暂不支持视频模型 %q", modelName)
|
||||
}
|
||||
input := map[string]any{}
|
||||
if err := json.Unmarshal(inputData, &input); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
prompt := strings.TrimSpace(fmt.Sprint(input["prompt"]))
|
||||
if prompt == "" || prompt == "<nil>" {
|
||||
return nil, errors.New("生成提示词不能为空")
|
||||
}
|
||||
size := strings.TrimSpace(fmt.Sprint(input["size"]))
|
||||
if size != "16:9" && size != "9:16" {
|
||||
return nil, errors.New("项目画面比例不受 Seedance 2 支持")
|
||||
}
|
||||
duration, err := strconv.Atoi(fmt.Sprint(input["duration"]))
|
||||
if err != nil || duration < 5 || duration > 15 {
|
||||
return nil, errors.New("Seedance 2 视频时长必须为 5 到 15 秒")
|
||||
}
|
||||
images, err := seedanceURLList(input["image_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考图无效: %w", err)
|
||||
}
|
||||
if len(images) > capabilities.ReferenceImages.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 张参考图", capabilities.ReferenceImages.MaxCount)
|
||||
}
|
||||
audios, err := seedanceURLList(input["audio_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考音频无效: %w", err)
|
||||
}
|
||||
if len(audios) > capabilities.ReferenceAudios.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考音频", capabilities.ReferenceAudios.MaxCount)
|
||||
}
|
||||
videos, err := seedanceURLList(input["video_urls"])
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("Seedance 2 参考视频无效: %w", err)
|
||||
}
|
||||
if len(videos) > capabilities.ReferenceVideos.MaxCount {
|
||||
return nil, fmt.Errorf("Seedance 2 最多支持 %d 个参考视频", capabilities.ReferenceVideos.MaxCount)
|
||||
}
|
||||
if len(audios) > 0 && len(images) == 0 && len(videos) == 0 {
|
||||
return nil, errors.New("Seedance 2 参考音频必须与参考图片或参考视频同时使用")
|
||||
}
|
||||
payload := map[string]any{
|
||||
"model": modelName, "prompt": prompt, "size": size, "duration": duration,
|
||||
"resolution": seedanceString(input["resolution"], "480p"),
|
||||
}
|
||||
if len(images) > 0 {
|
||||
payload["image_urls"] = images
|
||||
}
|
||||
if len(audios) > 0 {
|
||||
payload["audio_urls"] = audios
|
||||
}
|
||||
if len(videos) > 0 {
|
||||
payload["video_urls"] = videos
|
||||
}
|
||||
if value, ok := input["generate_audio"].(bool); ok {
|
||||
payload["generate_audio"] = value
|
||||
}
|
||||
if value, ok := input["return_last_frame"].(bool); ok {
|
||||
payload["return_last_frame"] = value
|
||||
}
|
||||
return payload, nil
|
||||
}
|
||||
|
||||
func (s Seedance2) Submit(ctx context.Context, baseURL, apiKey, requestID string, payload map[string]any) (provider.SubmitResult, error) {
|
||||
data, err := s.Gateway.requestJSON(ctx, http.MethodPost, baseURL, "/videos/generations", apiKey, requestID, payload)
|
||||
if err != nil {
|
||||
return provider.SubmitResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.SubmitResult{}, errors.New("Seedance 2 提交响应不是有效 JSON")
|
||||
}
|
||||
result := provider.SubmitResult{TaskID: findString(decoded, "task_id", "id"), URL: findSeedanceVideoURL(decoded), Raw: append(json.RawMessage(nil), data...)}
|
||||
if result.TaskID == "" && result.URL == "" {
|
||||
return provider.SubmitResult{}, errors.New("Seedance 2 提交响应缺少任务 ID 或视频地址")
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (s Seedance2) Poll(ctx context.Context, baseURL, apiKey, taskID string) (provider.PollResult, error) {
|
||||
data, err := s.Gateway.requestJSON(ctx, http.MethodGet, baseURL, "/tasks/"+taskID, apiKey, "", nil)
|
||||
if err != nil {
|
||||
return provider.PollResult{}, err
|
||||
}
|
||||
var decoded any
|
||||
if err := json.Unmarshal(data, &decoded); err != nil {
|
||||
return provider.PollResult{}, errors.New("Seedance 2 轮询响应不是有效 JSON")
|
||||
}
|
||||
status := strings.ToLower(findString(decoded, "task_status", "status", "state"))
|
||||
if status == "" {
|
||||
status = "processing"
|
||||
}
|
||||
return provider.PollResult{
|
||||
Status: status, URL: findSeedanceVideoURL(decoded),
|
||||
Error: findString(decoded, "error_message", "message", "detail"), Raw: append(json.RawMessage(nil), data...),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func seedanceURLList(value any) ([]string, error) {
|
||||
if value == nil {
|
||||
return nil, nil
|
||||
}
|
||||
raw, ok := value.([]any)
|
||||
if !ok {
|
||||
return nil, errors.New("必须是 URL 数组")
|
||||
}
|
||||
result := make([]string, 0, len(raw))
|
||||
for _, item := range raw {
|
||||
url := strings.TrimSpace(fmt.Sprint(item))
|
||||
if !strings.HasPrefix(url, "https://") && !strings.HasPrefix(url, "http://") && !strings.HasPrefix(url, "asset://") {
|
||||
return nil, errors.New("包含不可访问的 URL")
|
||||
}
|
||||
result = append(result, url)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func seedanceString(value any, fallback string) string {
|
||||
text := strings.TrimSpace(fmt.Sprint(value))
|
||||
if text == "" || text == "<nil>" {
|
||||
return fallback
|
||||
}
|
||||
return text
|
||||
}
|
||||
|
||||
func findSeedanceVideoURL(value any) string {
|
||||
if url := findURLInSeedanceVideoField(value); url != "" {
|
||||
return url
|
||||
}
|
||||
return findURL(value)
|
||||
}
|
||||
|
||||
func findURLInSeedanceVideoField(value any) string {
|
||||
switch data := value.(type) {
|
||||
case map[string]any:
|
||||
for key, child := range data {
|
||||
if strings.Contains(strings.ToLower(key), "video") {
|
||||
if text, ok := child.(string); ok && (strings.HasPrefix(text, "https://") || strings.HasPrefix(text, "http://")) {
|
||||
return text
|
||||
}
|
||||
if url := findURL(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, child := range data {
|
||||
if url := findURLInSeedanceVideoField(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
case []any:
|
||||
for _, child := range data {
|
||||
if url := findURLInSeedanceVideoField(child); url != "" {
|
||||
return url
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
package apimart
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestVideoCapabilities(t *testing.T) {
|
||||
capabilities, ok := VideoCapabilities("doubao-seedance-2.0")
|
||||
if !ok {
|
||||
t.Fatal("expected Seedance 2 capabilities")
|
||||
}
|
||||
if capabilities.ReferenceImages.MaxCount != 9 || capabilities.ReferenceAudios.MaxCount != 3 || capabilities.ReferenceVideos.MaxCount != 3 {
|
||||
t.Fatalf("unexpected reference limits: %+v", capabilities)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSeedance2PayloadSupportsVideoAndAudioReferences(t *testing.T) {
|
||||
input := json.RawMessage(`{
|
||||
"prompt":"test",
|
||||
"size":"16:9",
|
||||
"duration":5,
|
||||
"video_urls":["https://example.com/reference.mp4"],
|
||||
"audio_urls":["https://example.com/reference.mp3"]
|
||||
}`)
|
||||
payload, err := BuildSeedance2Payload("doubao-seedance-2.0", input)
|
||||
if err != nil {
|
||||
t.Fatalf("BuildSeedance2Payload() error = %v", err)
|
||||
}
|
||||
if videos, ok := payload["video_urls"].([]string); !ok || len(videos) != 1 {
|
||||
t.Fatalf("video_urls = %#v", payload["video_urls"])
|
||||
}
|
||||
if audios, ok := payload["audio_urls"].([]string); !ok || len(audios) != 1 {
|
||||
t.Fatalf("audio_urls = %#v", payload["audio_urls"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSeedance2PayloadRejectsTooManyVideos(t *testing.T) {
|
||||
input := json.RawMessage(`{
|
||||
"prompt":"test",
|
||||
"size":"16:9",
|
||||
"duration":5,
|
||||
"video_urls":["https://example.com/1.mp4","https://example.com/2.mp4","https://example.com/3.mp4","https://example.com/4.mp4"]
|
||||
}`)
|
||||
if _, err := BuildSeedance2Payload("doubao-seedance-2.0", input); err == nil {
|
||||
t.Fatal("expected too many reference videos to be rejected")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user