feat(api): migrate media clients and universal scripts

This commit is contained in:
2026-05-20 21:29:07 +08:00
parent 11a2c13e4a
commit af9b281d34
14 changed files with 2648 additions and 18 deletions
+232
View File
@@ -0,0 +1,232 @@
package clients
import (
"context"
"net/http"
"strings"
)
type JimengClient struct{ HTTPClient *http.Client }
type BlackforestClient struct{ HTTPClient *http.Client }
type HunyuanImageClient struct{ HTTPClient *http.Client }
type HunyuanVideoClient struct{ HTTPClient *http.Client }
type MinimaxClient struct{ HTTPClient *http.Client }
type MidjourneyClient struct{ HTTPClient *http.Client }
type ViduClient struct{ HTTPClient *http.Client }
type AliyunBailianClient struct{ HTTPClient *http.Client }
type NewAPIClient struct{ HTTPClient *http.Client }
func (c JimengClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: jimengSpec()}.Run(ctx, request)
}
func (c BlackforestClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: blackforestSpec()}.Run(ctx, request)
}
func (c HunyuanImageClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: hunyuanImageSpec()}.Run(ctx, request)
}
func (c HunyuanVideoClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: hunyuanVideoSpec()}.Run(ctx, request)
}
func (c MinimaxClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: minimaxSpec()}.Run(ctx, request)
}
func (c MidjourneyClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: midjourneySpec()}.Run(ctx, request)
}
func (c ViduClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: viduSpec()}.Run(ctx, request)
}
func (c AliyunBailianClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: aliyunBailianSpec()}.Run(ctx, request)
}
func (c NewAPIClient) Run(ctx context.Context, request Request) (Response, error) {
return providerTaskClient{HTTPClient: c.HTTPClient, Spec: newAPISpec()}.Run(ctx, request)
}
func jimengSpec() providerTaskSpec {
return providerTaskSpec{
Name: "jimeng",
SubmitPath: func(request Request, _ map[string]any) string {
return configuredPath(request, "?Action=CVSubmitTask&Version=2022-08-31", "submitPath", "submit_path")
},
PollPath: func(request Request, _ string, _ map[string]any) string {
return configuredPath(request, "?Action=CVSync2AsyncGetResult&Version=2022-08-31", "pollPath", "poll_path")
},
Auth: "bearer",
TaskIDPaths: []string{"data.task_id"},
StatusPaths: []string{"data.status"},
SuccessStatuses: []string{"done"},
DefaultSubmitBody: func(request Request, body map[string]any) map[string]any {
body["req_key"] = upstreamModelName(request.Candidate)
if body["prompt"] == nil {
body["prompt"] = mediaPromptText(body)
}
return body
},
}
}
func blackforestSpec() providerTaskSpec {
return providerTaskSpec{
Name: "blackforest",
SubmitPath: func(request Request, body map[string]any) string {
return configuredPath(request, "/"+upstreamModelName(request.Candidate), "submitPath", "submit_path")
},
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string { return upstreamTaskID },
Auth: "x-key",
TaskIDPaths: []string{"polling_url"},
StatusPaths: []string{"status"},
SuccessStatuses: []string{"ready"},
FailureStatuses: []string{"error", "task not found"},
}
}
func hunyuanImageSpec() providerTaskSpec {
return providerTaskSpec{
Name: "tencent-hunyuan-image",
SubmitPath: func(request Request, _ map[string]any) string {
return configuredPath(request, "?Action=SubmitHunyuanImageJob&Version=2023-09-01", "submitPath", "submit_path")
},
PollPath: func(request Request, _ string, _ map[string]any) string {
return configuredPath(request, "?Action=QueryHunyuanImageJob&Version=2023-09-01&JobId=${taskId}", "pollPath", "poll_path")
},
Auth: "bearer",
TaskIDPaths: []string{"Response.JobId"},
StatusPaths: []string{"Response.Status"},
SuccessStatuses: []string{"done"},
FailureStatuses: []string{"fail"},
DefaultSubmitBody: func(request Request, body map[string]any) map[string]any {
body["Prompt"] = mediaPromptText(body)
body["Model"] = upstreamModelName(request.Candidate)
return body
},
}
}
func hunyuanVideoSpec() providerTaskSpec {
return providerTaskSpec{
Name: "tencent-hunyuan-video",
SubmitPath: func(request Request, _ map[string]any) string {
return configuredPath(request, "?Action=SubmitTextToVideoJob&Version=2024-01-01", "submitPath", "submit_path")
},
PollPath: func(request Request, _ string, _ map[string]any) string {
return configuredPath(request, "?Action=QueryVideoJob&Version=2024-01-01&JobId=${taskId}", "pollPath", "poll_path")
},
Auth: "bearer",
TaskIDPaths: []string{"Response.JobId"},
StatusPaths: []string{"Response.Status"},
SuccessStatuses: []string{"done"},
FailureStatuses: []string{"fail"},
DefaultSubmitBody: func(request Request, body map[string]any) map[string]any {
body["Prompt"] = mediaPromptText(body)
body["Model"] = upstreamModelName(request.Candidate)
return body
},
}
}
func minimaxSpec() providerTaskSpec {
return providerTaskSpec{
Name: "minimax",
SubmitPath: func(Request, map[string]any) string { return "/video_generation" },
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string {
return "/query/video_generation?task_id=" + upstreamTaskID
},
Auth: "bearer",
TaskIDPaths: []string{"task_id"},
StatusPaths: []string{"status"},
SuccessStatuses: []string{"success"},
FailureStatuses: []string{"failed", "expired"},
}
}
func midjourneySpec() providerTaskSpec {
return providerTaskSpec{
Name: "midjourney",
SubmitPath: func(request Request, body map[string]any) string {
return configuredPath(request, "/diffusion", "submitPath", "submit_path")
},
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string { return "/job/" + upstreamTaskID },
Auth: "bearer",
TaskIDPaths: []string{"job_id", "id"},
StatusPaths: []string{"status"},
SuccessStatuses: []string{"success", "completed"},
FailureStatuses: []string{"failed"},
DefaultSubmitBody: func(request Request, body map[string]any) map[string]any {
if body["prompt"] == nil && body["text"] == nil {
body["prompt"] = mediaPromptText(body)
}
return body
},
}
}
func viduSpec() providerTaskSpec {
return providerTaskSpec{
Name: "vidu",
SubmitPath: func(request Request, body map[string]any) string {
if path := configuredPath(request, "", "submitPath", "submit_path"); path != "" {
return path
}
taskType := firstNonEmptyString(body["type"], body["task_type"], "text2video")
if taskType == "multiframe" {
return "/multiframe"
}
return "/" + taskType
},
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string {
return "/tasks/" + upstreamTaskID + "/creations"
},
Auth: "token",
TaskIDPaths: []string{"task_id"},
StatusPaths: []string{"state", "status"},
SuccessStatuses: []string{"success", "succeeded"},
FailureStatuses: []string{"failed"},
}
}
func aliyunBailianSpec() providerTaskSpec {
return providerTaskSpec{
Name: "aliyun-bailian",
SubmitPath: func(Request, map[string]any) string { return "/services/aigc/video-generation/video-synthesis" },
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string { return "/tasks/" + upstreamTaskID },
Auth: "bearer",
TaskIDPaths: []string{"output.task_id"},
StatusPaths: []string{"output.task_status"},
SuccessStatuses: []string{"succeeded", "success"},
FailureStatuses: []string{"failed"},
}
}
func newAPISpec() providerTaskSpec {
return providerTaskSpec{
Name: "newapi",
SubmitPath: func(Request, map[string]any) string { return "/videos/generations" },
PollPath: func(_ Request, upstreamTaskID string, _ map[string]any) string {
return "/videos/generations/" + upstreamTaskID
},
Auth: "bearer",
TaskIDPaths: []string{"task_id"},
StatusPaths: []string{"status"},
SuccessStatuses: []string{"success"},
FailureStatuses: []string{"failure", "failed"},
}
}
func configuredPath(request Request, fallback string, keys ...string) string {
for _, key := range keys {
if value := strings.TrimSpace(stringFromAny(request.Candidate.PlatformConfig[key])); value != "" {
return value
}
}
return fallback
}
+453
View File
@@ -0,0 +1,453 @@
package clients
import (
"bytes"
"context"
"encoding/json"
"fmt"
"net/http"
"strings"
"time"
)
type providerTaskSpec struct {
Name string
SubmitPath func(Request, map[string]any) string
PollPath func(Request, string, map[string]any) string
Auth string
TaskIDPaths []string
StatusPaths []string
SuccessStatuses []string
FailureStatuses []string
ProcessStatuses []string
DefaultSubmitBody func(Request, map[string]any) map[string]any
}
type providerTaskClient struct {
HTTPClient *http.Client
Spec providerTaskSpec
}
func (c providerTaskClient) Run(ctx context.Context, request Request) (Response, error) {
if request.Kind != "images.generations" && request.Kind != "images.edits" && request.Kind != "videos.generations" {
return Response{}, &ClientError{Code: "unsupported_kind", Message: "unsupported " + c.Spec.Name + " request kind", Retryable: false}
}
startedAt := time.Now()
payload := cloneBody(request.Body)
if c.Spec.DefaultSubmitBody != nil {
payload = c.Spec.DefaultSubmitBody(request, payload)
} else {
payload["model"] = upstreamModelName(request.Candidate)
}
upstreamTaskID := strings.TrimSpace(request.RemoteTaskID)
requestID := upstreamTaskID
var submitResult map[string]any
if upstreamTaskID == "" {
result, id, err := c.submit(ctx, request, payload)
if err != nil {
return Response{}, annotateResponseError(err, id, startedAt, time.Now())
}
submitResult = result
requestID = firstNonEmptyString(id, requestIDFromResult(result))
if isProviderTaskFailure(c.Spec, result) {
return Response{}, providerTaskFailure(c.Spec, result, requestID, startedAt)
}
if isProviderTaskSuccess(c.Spec, result) && hasProviderTaskResult(result) {
return Response{
Result: normalizeProviderTaskResult(request, c.Spec, result, ""),
RequestID: requestID,
Progress: providerProgress(request),
ResponseStartedAt: startedAt,
ResponseFinishedAt: time.Now(),
ResponseDurationMS: responseDurationMS(startedAt, time.Now()),
}, nil
}
upstreamTaskID = providerTaskID(c.Spec, result)
if upstreamTaskID == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: c.Spec.Name + " task id is missing", RequestID: requestID, Retryable: false}
}
if request.OnRemoteTaskSubmitted != nil {
if err := request.OnRemoteTaskSubmitted(upstreamTaskID, map[string]any{"payload": payload, "submit": submitResult}); err != nil {
return Response{}, err
}
}
} else if request.RemoteTaskPayload != nil {
if existingPayload, ok := request.RemoteTaskPayload["payload"].(map[string]any); ok {
payload = existingPayload
}
}
interval := providerPollInterval(request)
timeout := providerPollTimeout(request)
deadline := time.NewTimer(timeout)
defer deadline.Stop()
ticker := time.NewTicker(interval)
defer ticker.Stop()
var lastResult map[string]any
for {
pollStarted := time.Now()
result, pollRequestID, err := c.poll(ctx, request, upstreamTaskID, payload)
pollFinished := time.Now()
if err != nil {
return Response{}, annotateResponseError(err, firstNonEmptyString(pollRequestID, requestID, upstreamTaskID), pollStarted, pollFinished)
}
lastResult = result
requestID = firstNonEmptyString(pollRequestID, requestID, requestIDFromResult(result), upstreamTaskID)
if isProviderTaskSuccess(c.Spec, result) {
finishedAt := time.Now()
return Response{
Result: normalizeProviderTaskResult(request, c.Spec, result, upstreamTaskID),
RequestID: requestID,
Progress: append(providerProgress(request), Progress{Phase: "polling", Progress: 0.65, Message: "provider task polled", Payload: map[string]any{"upstreamTaskId": upstreamTaskID}}),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
if isProviderTaskFailure(c.Spec, result) {
return Response{}, providerTaskFailure(c.Spec, result, requestID, startedAt)
}
select {
case <-ctx.Done():
return Response{}, &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
case <-deadline.C:
return Response{}, &ClientError{Code: "timeout", Message: fmt.Sprintf("%s task %s did not finish before timeout; last status: %s", c.Spec.Name, upstreamTaskID, providerTaskStatus(c.Spec, lastResult)), RequestID: requestID, Retryable: true}
case <-ticker.C:
}
}
}
func (c providerTaskClient) submit(ctx context.Context, request Request, payload map[string]any) (map[string]any, string, error) {
path := c.Spec.SubmitPath(request, payload)
return providerPostJSON(ctx, httpClient(request.HTTPClient, c.HTTPClient), providerURL(request.Candidate.BaseURL, path), payload, request.Candidate.Credentials, c.Spec.Auth)
}
func (c providerTaskClient) poll(ctx context.Context, request Request, upstreamTaskID string, payload map[string]any) (map[string]any, string, error) {
path := resolveProviderPathTemplate(c.Spec.PollPath(request, upstreamTaskID, payload), upstreamTaskID)
url := path
if !strings.HasPrefix(path, "http://") && !strings.HasPrefix(path, "https://") {
url = providerURL(request.Candidate.BaseURL, path)
}
if c.Spec.Name == "jimeng" {
body := map[string]any{"task_id": upstreamTaskID, "req_key": upstreamModelName(request.Candidate)}
return providerPostJSON(ctx, httpClient(request.HTTPClient, c.HTTPClient), url, body, request.Candidate.Credentials, c.Spec.Auth)
}
return providerGetJSON(ctx, httpClient(request.HTTPClient, c.HTTPClient), url, request.Candidate.Credentials, c.Spec.Auth)
}
func providerPostJSON(ctx context.Context, client *http.Client, url string, body map[string]any, credentials map[string]any, auth string) (map[string]any, string, error) {
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(raw))
if err != nil {
return nil, "", err
}
req.Header.Set("Content-Type", "application/json")
applyProviderAuth(req, credentials, auth)
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
return result, requestID, err
}
func providerGetJSON(ctx context.Context, client *http.Client, url string, credentials map[string]any, auth string) (map[string]any, string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, "", err
}
applyProviderAuth(req, credentials, auth)
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
return result, requestID, err
}
func applyProviderAuth(req *http.Request, credentials map[string]any, auth string) {
apiKey := credential(credentials, "apiKey", "api_key", "key", "token")
switch auth {
case "token":
if apiKey != "" {
req.Header.Set("Authorization", "Token "+apiKey)
}
case "x-key":
if apiKey != "" {
req.Header.Set("x-key", apiKey)
}
case "none":
default:
if apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
}
}
func providerURL(base string, path string) string {
path = strings.TrimSpace(path)
if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
return path
}
if path == "" {
path = "/"
}
if !strings.HasPrefix(path, "/") && !strings.HasPrefix(path, "?") {
path = "/" + path
}
return joinURL(base, path)
}
func resolveProviderPathTemplate(path string, upstreamTaskID string) string {
replacements := [][2]string{
{"${upstream_task_id}", upstreamTaskID},
{"{{upstream_task_id}}", upstreamTaskID},
{"{upstream_task_id}", upstreamTaskID},
{"${task_id}", upstreamTaskID},
{"{{task_id}}", upstreamTaskID},
{"{task_id}", upstreamTaskID},
{"${taskId}", upstreamTaskID},
{"${taskID}", upstreamTaskID},
{"{{taskId}}", upstreamTaskID},
{"{{taskID}}", upstreamTaskID},
{"{taskId}", upstreamTaskID},
{"{taskID}", upstreamTaskID},
}
for _, replacement := range replacements {
path = strings.ReplaceAll(path, replacement[0], replacement[1])
}
return path
}
func providerTaskID(spec providerTaskSpec, result map[string]any) string {
paths := append([]string{}, spec.TaskIDPaths...)
paths = append(paths, "task_id", "taskId", "id", "job_id", "Response.JobId", "output.task_id", "data.task_id", "polling_url")
for _, path := range paths {
if value := stringFromPathValue(valueAtPath(result, path)); value != "" {
return value
}
}
return ""
}
func providerTaskStatus(spec providerTaskSpec, result map[string]any) string {
if result == nil {
return ""
}
if value, ok := valueAtPath(result, "status").(float64); ok {
if value == 2 {
return "success"
}
if value == 3 {
return "failed"
}
return "process"
}
paths := append([]string{}, spec.StatusPaths...)
paths = append(paths, "status", "state", "task_status", "output.task_status", "Response.Status", "data.status")
for _, path := range paths {
if value := stringFromPathValue(valueAtPath(result, path)); value != "" {
return strings.ToLower(value)
}
}
return ""
}
func stringFromPathValue(value any) string {
if value == nil {
return ""
}
text := strings.TrimSpace(fmt.Sprint(value))
if text == "" || text == "<nil>" {
return ""
}
return text
}
func isProviderTaskSuccess(spec providerTaskSpec, result map[string]any) bool {
return containsStatus(append([]string{"success", "succeeded", "completed", "complete", "done", "ready", "succeed", "succeeded", "suceeded", "done", "done"}, spec.SuccessStatuses...), providerTaskStatus(spec, result))
}
func isProviderTaskFailure(spec providerTaskSpec, result map[string]any) bool {
return containsStatus(append([]string{"failed", "failure", "error", "cancelled", "canceled", "fail", "expired", "task not found"}, spec.FailureStatuses...), providerTaskStatus(spec, result))
}
func containsStatus(values []string, status string) bool {
status = strings.ToLower(strings.TrimSpace(status))
for _, value := range values {
if strings.ToLower(strings.TrimSpace(value)) == status {
return true
}
}
return false
}
func hasProviderTaskResult(result map[string]any) bool {
return result["data"] != nil || valueAtPath(result, "output.image_urls") != nil || valueAtPath(result, "output.video_url") != nil || valueAtPath(result, "Response.ResultVideoUrl") != nil || valueAtPath(result, "Response.ResultImages") != nil || result["urls"] != nil
}
func normalizeProviderTaskResult(request Request, spec providerTaskSpec, result map[string]any, upstreamTaskID string) map[string]any {
out := cloneMapAny(result)
out["status"] = "success"
if upstreamTaskID != "" {
out["upstream_task_id"] = upstreamTaskID
}
if out["created"] == nil {
out["created"] = time.Now().UnixMilli()
}
if out["model"] == nil {
out["model"] = request.Model
}
if _, ok := out["data"].([]any); !ok {
if out["data"] != nil {
out["raw_data"] = out["data"]
}
out["data"] = providerTaskData(request, result)
}
return out
}
func providerTaskData(request Request, result map[string]any) []any {
fileType := "image"
if request.Kind == "videos.generations" || strings.Contains(request.ModelType, "video") {
fileType = "video"
}
urlValues := []any{}
for _, path := range []string{
"urls",
"image_urls",
"data.image_urls",
"data.images",
"output.image_urls",
"output.video_url",
"output.output",
"data.output",
"data.video_url",
"video_url",
"preview_url",
"Response.ResultImages",
"Response.ResultVideoUrl",
} {
appendURLValues(&urlValues, valueAtPath(result, path))
}
data := make([]any, 0, len(urlValues))
for _, raw := range urlValues {
if url := strings.TrimSpace(fmt.Sprint(raw)); url != "" {
data = append(data, map[string]any{"type": fileType, "url": url})
}
}
if len(data) == 0 {
if base64Values := valueAtPath(result, "data.binary_data_base64"); base64Values != nil {
values := []any{}
appendURLValues(&values, base64Values)
for _, raw := range values {
if content := strings.TrimSpace(fmt.Sprint(raw)); content != "" {
data = append(data, map[string]any{"type": fileType, "content": content, "uploaded": false})
}
}
}
}
return data
}
func appendURLValues(out *[]any, value any) {
switch typed := value.(type) {
case nil:
case string:
*out = append(*out, typed)
case []any:
for _, item := range typed {
appendURLValues(out, item)
}
case []string:
for _, item := range typed {
*out = append(*out, item)
}
case map[string]any:
for _, key := range []string{"url", "image_url", "imageUrl", "video_url", "videoUrl", "content", "output"} {
if item := strings.TrimSpace(fmt.Sprint(typed[key])); item != "" && item != "<nil>" {
*out = append(*out, item)
return
}
}
}
}
func providerTaskFailure(spec providerTaskSpec, result map[string]any, requestID string, startedAt time.Time) error {
message := firstNonEmptyString(valueAtPath(result, "message"), valueAtPath(result, "error.message"), valueAtPath(result, "error"), valueAtPath(result, "Response.ErrorMessage"), valueAtPath(result, "comment"), spec.Name+" task failed")
return &ClientError{
Code: firstNonEmptyString(valueAtPath(result, "code"), valueAtPath(result, "error_code"), valueAtPath(result, "Response.ErrorCode"), "provider_failed"),
Message: message,
RequestID: requestID,
ResponseStartedAt: startedAt,
ResponseFinishedAt: time.Now(),
ResponseDurationMS: responseDurationMS(startedAt, time.Now()),
Retryable: false,
}
}
func providerPollInterval(request Request) time.Duration {
return durationFromConfig(request.Candidate.PlatformConfig, 2*time.Second, "pollIntervalMs", "poll_interval_ms")
}
func providerPollTimeout(request Request) time.Duration {
return durationFromConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
}
func durationFromConfig(config map[string]any, fallback time.Duration, keys ...string) time.Duration {
for _, key := range keys {
switch value := config[key].(type) {
case int:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case int64:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case float64:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case string:
if parsed, err := time.ParseDuration(value); err == nil && parsed > 0 {
return parsed
}
}
}
return fallback
}
func valueAtPath(values map[string]any, path string) any {
if values == nil || strings.TrimSpace(path) == "" {
return nil
}
var current any = values
for _, part := range strings.Split(path, ".") {
object, ok := current.(map[string]any)
if !ok {
return nil
}
current = object[part]
}
return current
}
func mediaPromptText(body map[string]any) string {
if prompt := strings.TrimSpace(stringFromAny(body["prompt"])); prompt != "" {
return prompt
}
content, _ := body["content"].([]any)
for _, item := range content {
if part, ok := item.(map[string]any); ok && strings.TrimSpace(stringFromAny(part["type"])) == "text" {
if text := strings.TrimSpace(stringFromAny(part["text"])); text != "" {
return text
}
}
}
return ""
}
@@ -0,0 +1,339 @@
package clients
import (
"context"
"errors"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestProviderTaskClientsSubmitAndPoll(t *testing.T) {
cases := []struct {
name string
client Client
provider string
specType string
submitMatch func(*http.Request) bool
submitResponse string
pollMatch func(*http.Request) bool
pollResponse string
authHeader string
resultURL string
}{
{
name: "jimeng",
client: JimengClient{},
provider: "jimeng",
specType: "jimeng",
submitMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && r.URL.Query().Get("Action") == "CVSubmitTask"
},
submitResponse: `{"code":10000,"data":{"task_id":"remote-1"}}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && r.URL.Query().Get("Action") == "CVSync2AsyncGetResult"
},
pollResponse: `{"code":10000,"data":{"status":"done","video_url":"https://cdn.example/jimeng.mp4"}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/jimeng.mp4",
},
{
name: "blackforest",
client: BlackforestClient{},
provider: "blackforest",
specType: "blackforest",
submitMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && strings.HasPrefix(r.URL.Path, "/provider-model")
},
submitResponse: `{"polling_url":"__SERVER__/poll/remote-1"}`,
pollMatch: func(r *http.Request) bool { return r.Method == http.MethodGet && r.URL.Path == "/poll/remote-1" },
pollResponse: `{"status":"Ready","urls":["https://cdn.example/flux.png"]}`,
authHeader: "test-key",
resultURL: "https://cdn.example/flux.png",
},
{
name: "hunyuan-image",
client: HunyuanImageClient{},
provider: "tencent-hunyuan-image",
specType: "tencent-hunyuan-image",
submitMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && r.URL.Query().Get("Action") == "SubmitHunyuanImageJob"
},
submitResponse: `{"Response":{"JobId":"remote-1"}}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodGet && r.URL.Query().Get("Action") == "QueryHunyuanImageJob"
},
pollResponse: `{"Response":{"Status":"DONE","ResultImages":["https://cdn.example/hunyuan.png"]}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/hunyuan.png",
},
{
name: "hunyuan-video",
client: HunyuanVideoClient{},
provider: "tencent-hunyuan-video",
specType: "tencent-hunyuan-video",
submitMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && r.URL.Query().Get("Action") == "SubmitTextToVideoJob"
},
submitResponse: `{"Response":{"JobId":"remote-1"}}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodGet && r.URL.Query().Get("Action") == "QueryVideoJob"
},
pollResponse: `{"Response":{"Status":"DONE","ResultVideoUrl":"https://cdn.example/hunyuan.mp4"}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/hunyuan.mp4",
},
{
name: "minimax",
client: MinimaxClient{},
provider: "minimax",
specType: "minimax",
submitMatch: func(r *http.Request) bool { return r.Method == http.MethodPost && r.URL.Path == "/video_generation" },
submitResponse: `{"task_id":123}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodGet && r.URL.Path == "/query/video_generation" && r.URL.Query().Get("task_id") == "123"
},
pollResponse: `{"status":"Success","file_id":"file-1","video_url":"https://cdn.example/minimax.mp4"}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/minimax.mp4",
},
{
name: "midjourney",
client: MidjourneyClient{},
provider: "midjourney",
specType: "midjourney",
submitMatch: func(r *http.Request) bool { return r.Method == http.MethodPost && r.URL.Path == "/diffusion" },
submitResponse: `{"job_id":"remote-1"}`,
pollMatch: func(r *http.Request) bool { return r.Method == http.MethodGet && r.URL.Path == "/job/remote-1" },
pollResponse: `{"status":"completed","output":{"image_urls":["https://cdn.example/mj.png"]}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/mj.png",
},
{
name: "vidu",
client: ViduClient{},
provider: "vidu",
specType: "vidu",
submitMatch: func(r *http.Request) bool { return r.Method == http.MethodPost && r.URL.Path == "/text2video" },
submitResponse: `{"task_id":"remote-1"}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodGet && r.URL.Path == "/tasks/remote-1/creations"
},
pollResponse: `{"state":"success","video_url":"https://cdn.example/vidu.mp4"}`,
authHeader: "Token test-key",
resultURL: "https://cdn.example/vidu.mp4",
},
{
name: "aliyun-bailian",
client: AliyunBailianClient{},
provider: "aliyun-bailian",
specType: "aliyun-bailian",
submitMatch: func(r *http.Request) bool {
return r.Method == http.MethodPost && r.URL.Path == "/services/aigc/video-generation/video-synthesis"
},
submitResponse: `{"output":{"task_id":"remote-1"}}`,
pollMatch: func(r *http.Request) bool { return r.Method == http.MethodGet && r.URL.Path == "/tasks/remote-1" },
pollResponse: `{"output":{"task_status":"SUCCEEDED","video_url":"https://cdn.example/aliyun.mp4"}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/aliyun.mp4",
},
{
name: "newapi",
client: NewAPIClient{},
provider: "newapi",
specType: "newapi",
submitMatch: func(r *http.Request) bool { return r.Method == http.MethodPost && r.URL.Path == "/videos/generations" },
submitResponse: `{"task_id":"remote-1"}`,
pollMatch: func(r *http.Request) bool {
return r.Method == http.MethodGet && r.URL.Path == "/videos/generations/remote-1"
},
pollResponse: `{"status":"SUCCESS","data":{"output":"https://cdn.example/newapi.mp4"}}`,
authHeader: "Bearer test-key",
resultURL: "https://cdn.example/newapi.mp4",
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if tc.authHeader == "test-key" {
if r.Header.Get("x-key") != tc.authHeader {
t.Fatalf("unexpected x-key header: %q", r.Header.Get("x-key"))
}
} else if r.Header.Get("Authorization") != tc.authHeader {
t.Fatalf("unexpected auth header: %q", r.Header.Get("Authorization"))
}
w.Header().Set("Content-Type", "application/json")
w.Header().Set("x-request-id", "req-"+tc.name)
switch {
case tc.submitMatch(r):
_, _ = w.Write([]byte(strings.ReplaceAll(tc.submitResponse, "__SERVER__", "http://"+r.Host)))
case tc.pollMatch(r):
_, _ = w.Write([]byte(tc.pollResponse))
default:
t.Fatalf("unexpected request: %s %s", r.Method, r.URL.String())
}
}))
defer server.Close()
var submittedRemoteTaskID string
request := Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "alias-model",
Body: map[string]any{"model": "alias-model", "prompt": "hello"},
Candidate: store.RuntimeModelCandidate{
Provider: tc.provider,
SpecType: tc.specType,
BaseURL: server.URL,
Credentials: map[string]any{"apiKey": "test-key"},
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
ModelName: "alias-model",
ProviderModelName: "provider-model",
ModelType: "video_generate",
ClientID: tc.name + ":test",
},
OnRemoteTaskSubmitted: func(remoteTaskID string, payload map[string]any) error {
submittedRemoteTaskID = remoteTaskID
if payload["payload"] == nil || payload["submit"] == nil {
t.Fatalf("missing remote payload: %#v", payload)
}
return nil
},
}
response, err := tc.client.Run(context.Background(), request)
if err != nil {
t.Fatalf("run failed: %v", err)
}
data, ok := response.Result["data"].([]any)
if !ok || len(data) == 0 {
t.Fatalf("missing data: %#v", response.Result)
}
first, _ := data[0].(map[string]any)
if first["url"] != tc.resultURL {
t.Fatalf("unexpected result url: %#v", response.Result)
}
if response.RequestID != "req-"+tc.name {
t.Fatalf("unexpected request id: %q", response.RequestID)
}
if submittedRemoteTaskID == "" {
t.Fatalf("expected remote task submission")
}
})
}
}
func TestProviderTaskClientFailureAndRetryableErrors(t *testing.T) {
t.Run("poll failure", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.Header().Set("x-request-id", "req-failed")
switch r.URL.Path {
case "/videos/generations":
_, _ = w.Write([]byte(`{"task_id":"remote-1"}`))
case "/videos/generations/remote-1":
_, _ = w.Write([]byte(`{"status":"failed","code":"UPSTREAM_FAILED","message":"provider rejected"}`))
default:
t.Fatalf("unexpected request: %s", r.URL.String())
}
}))
defer server.Close()
_, err := (NewAPIClient{}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "alias-model",
Body: map[string]any{"model": "alias-model", "prompt": "hello"},
Candidate: store.RuntimeModelCandidate{
Provider: "newapi",
SpecType: "newapi",
BaseURL: server.URL,
Credentials: map[string]any{"apiKey": "test-key"},
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
ModelName: "alias-model",
ProviderModelName: "provider-model",
ModelType: "video_generate",
},
})
var clientErr *ClientError
if !errors.As(err, &clientErr) || clientErr.Code != "UPSTREAM_FAILED" || clientErr.Retryable {
t.Fatalf("expected non-retryable upstream failure, got %#v", err)
}
if clientErr.RequestID != "req-failed" {
t.Fatalf("unexpected request id: %q", clientErr.RequestID)
}
})
t.Run("submit rate limit", func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.Header().Set("x-request-id", "req-rate-limit")
w.WriteHeader(http.StatusTooManyRequests)
_, _ = w.Write([]byte(`{"error":{"message":"slow down"}}`))
}))
defer server.Close()
_, err := (NewAPIClient{}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "alias-model",
Body: map[string]any{"model": "alias-model", "prompt": "hello"},
Candidate: store.RuntimeModelCandidate{
Provider: "newapi",
SpecType: "newapi",
BaseURL: server.URL,
Credentials: map[string]any{"apiKey": "test-key"},
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
ModelName: "alias-model",
ProviderModelName: "provider-model",
ModelType: "video_generate",
},
})
var clientErr *ClientError
if !errors.As(err, &clientErr) || !clientErr.Retryable || clientErr.RequestID != "req-rate-limit" {
t.Fatalf("expected retryable rate limit with request id, got %#v", err)
}
})
}
func TestProviderTaskClientResumeSkipsSubmit(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Method == http.MethodPost {
t.Fatal("submit should not run for resumed remote task")
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"status":"success","data":{"output":"https://cdn.example/resume.mp4"}}`))
}))
defer server.Close()
response, err := (NewAPIClient{}).Run(context.Background(), Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "alias-model",
Body: map[string]any{"model": "alias-model", "prompt": "hello"},
RemoteTaskID: "remote-1",
RemoteTaskPayload: map[string]any{"payload": map[string]any{"prompt": "old"}},
Candidate: store.RuntimeModelCandidate{
Provider: "newapi",
SpecType: "newapi",
BaseURL: server.URL,
Credentials: map[string]any{"apiKey": "test-key"},
PlatformConfig: map[string]any{"pollIntervalMs": 1, "pollTimeoutMs": 1000},
ModelName: "alias-model",
ProviderModelName: "provider-model",
ModelType: "video_generate",
},
})
if err != nil {
t.Fatalf("run failed: %v", err)
}
data := response.Result["data"].([]any)
first := data[0].(map[string]any)
if first["url"] != "https://cdn.example/resume.mp4" || response.Result["upstream_task_id"] != "remote-1" {
t.Fatalf("unexpected response: %#v", response.Result)
}
}
+481
View File
@@ -0,0 +1,481 @@
package clients
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"net/http"
"strings"
"time"
scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script"
)
type UniversalClient struct {
HTTPClient *http.Client
ScriptExecutor *scriptengine.Executor
}
func (c UniversalClient) Run(ctx context.Context, request Request) (Response, error) {
executor := c.ScriptExecutor
if executor == nil {
executor = &scriptengine.Executor{}
}
startedAt := time.Now()
modelType := strings.TrimSpace(request.ModelType)
if modelType == "" {
modelType = strings.TrimSpace(request.Candidate.ModelType)
}
payload := cloneBody(request.Body)
upstreamTaskID := strings.TrimSpace(request.RemoteTaskID)
submitRequestID := upstreamTaskID
var submitResult map[string]any
if upstreamTaskID == "" {
var err error
payload, err = c.universalGetParams(ctx, executor, request, modelType)
if err != nil {
return Response{}, err
}
submitResult, submitRequestID, err = c.universalSubmit(ctx, executor, request, modelType, payload)
if err != nil {
return Response{}, annotateResponseError(err, submitRequestID, startedAt, time.Now())
}
if isUniversalSuccess(submitResult) && submitResult["data"] != nil {
return Response{
Result: normalizeUniversalResult(request, submitResult, ""),
RequestID: firstNonEmptyString(submitRequestID, requestIDFromResult(submitResult)),
Progress: providerProgress(request),
ResponseStartedAt: startedAt,
ResponseFinishedAt: time.Now(),
ResponseDurationMS: responseDurationMS(startedAt, time.Now()),
}, nil
}
if isUniversalFailure(submitResult) {
return Response{}, universalFailureError(submitResult, firstNonEmptyString(submitRequestID, requestIDFromResult(submitResult)), startedAt)
}
upstreamTaskID = universalTaskID(submitResult)
if upstreamTaskID == "" {
return Response{}, &ClientError{Code: "invalid_response", Message: "universal task id is missing", RequestID: submitRequestID, Retryable: false}
}
if request.OnRemoteTaskSubmitted != nil {
if err := request.OnRemoteTaskSubmitted(upstreamTaskID, map[string]any{"payload": payload, "submit": submitResult}); err != nil {
return Response{}, err
}
}
} else if request.RemoteTaskPayload != nil {
if existingPayload, ok := request.RemoteTaskPayload["payload"].(map[string]any); ok {
payload = existingPayload
}
}
result, requestID, err := c.universalPollUntilDone(ctx, executor, request, modelType, upstreamTaskID, payload, firstNonEmptyString(submitRequestID, upstreamTaskID), startedAt)
if err != nil {
return Response{}, err
}
finishedAt := time.Now()
return Response{
Result: normalizeUniversalResult(request, result, upstreamTaskID),
RequestID: firstNonEmptyString(requestID, submitRequestID, requestIDFromResult(result), upstreamTaskID),
Progress: universalProgress(request, upstreamTaskID),
ResponseStartedAt: startedAt,
ResponseFinishedAt: finishedAt,
ResponseDurationMS: responseDurationMS(startedAt, finishedAt),
}, nil
}
func (c UniversalClient) universalGetParams(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string) (map[string]any, error) {
if scriptText := universalSceneScript(request.Candidate.PlatformConfig, modelType, "customGetParamsScript", "custom_get_params_script"); scriptText != "" {
scriptContext := universalScriptContext(request, modelType, nil)
out, err := executor.Execute(ctx, scriptengine.Options{
Script: scriptText,
Args: []any{cloneBody(request.Body), scriptContext},
ContextData: scriptContext,
ScriptName: "custom_get_params_script:" + modelType,
PreferredEntryNames: []string{"getGenerateParams", "getParams", "main", "handler"},
Timeout: 30 * time.Second,
HTTPClient: httpClient(request.HTTPClient, c.HTTPClient),
})
if err != nil {
return nil, universalScriptError(err)
}
if params, ok := out.(map[string]any); ok && params != nil {
if params["_originalParams"] == nil {
params["_originalParams"] = cloneBody(request.Body)
}
return params, nil
}
return nil, &ClientError{Code: "invalid_response", Message: "custom get params script must return an object", Retryable: false}
}
body := universalDefaultPayload(request)
body["_originalParams"] = cloneBody(request.Body)
return body, nil
}
func (c UniversalClient) universalSubmit(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string, payload map[string]any) (map[string]any, string, error) {
if scriptText := universalSceneScript(request.Candidate.PlatformConfig, modelType, "customSubmitScript", "custom_submit_script"); scriptText != "" {
scriptContext := universalScriptContext(request, modelType, payload)
out, err := executor.Execute(ctx, scriptengine.Options{
Script: scriptText,
Args: []any{cloneBody(payload), scriptContext},
ContextData: scriptContext,
ScriptName: "custom_submit_script:" + modelType,
PreferredEntryNames: []string{"submitTask", "submitParams", "submit", "main", "handler"},
Timeout: 30 * time.Second,
HTTPClient: httpClient(request.HTTPClient, c.HTTPClient),
})
if err != nil {
return nil, "", universalScriptError(err)
}
result, ok := out.(map[string]any)
if !ok || result == nil {
return nil, "", &ClientError{Code: "invalid_response", Message: "custom submit script must return an object", Retryable: false}
}
return result, requestIDFromResult(result), nil
}
endpoint := universalSubmitEndpoint(request)
result, requestID, err := universalPostJSON(ctx, httpClient(request.HTTPClient, c.HTTPClient), request.Candidate.BaseURL, endpoint, universalStripPrivatePayload(payload), request.Candidate.Credentials)
return result, requestID, err
}
func (c UniversalClient) universalPollUntilDone(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string, upstreamTaskID string, payload map[string]any, requestID string, startedAt time.Time) (map[string]any, string, error) {
interval := universalDurationConfig(request.Candidate.PlatformConfig, 2*time.Second, "pollIntervalMs", "poll_interval_ms")
timeout := universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms", "timeoutMs")
deadline := time.NewTimer(timeout)
defer deadline.Stop()
ticker := time.NewTicker(interval)
defer ticker.Stop()
var lastResult map[string]any
for {
pollStarted := time.Now()
result, pollRequestID, err := c.universalPoll(ctx, executor, request, modelType, upstreamTaskID, payload)
pollFinished := time.Now()
if err != nil {
return nil, "", annotateResponseError(err, firstNonEmptyString(pollRequestID, requestID, upstreamTaskID), pollStarted, pollFinished)
}
lastResult = result
requestID = firstNonEmptyString(pollRequestID, requestID, requestIDFromResult(result), upstreamTaskID)
if isUniversalSuccess(result) {
return result, requestID, nil
}
if isUniversalFailure(result) {
return nil, "", universalFailureError(result, requestID, startedAt)
}
select {
case <-ctx.Done():
return nil, "", &ClientError{Code: "cancelled", Message: ctx.Err().Error(), RequestID: requestID, Retryable: true}
case <-deadline.C:
return nil, "", &ClientError{Code: "timeout", Message: fmt.Sprintf("universal task %s did not finish before timeout; last status: %s", upstreamTaskID, universalStatus(lastResult)), RequestID: requestID, Retryable: true}
case <-ticker.C:
}
}
}
func (c UniversalClient) universalPoll(ctx context.Context, executor *scriptengine.Executor, request Request, modelType string, upstreamTaskID string, payload map[string]any) (map[string]any, string, error) {
if scriptText := universalSceneScript(request.Candidate.PlatformConfig, modelType, "customPollScript", "custom_poll_script"); scriptText != "" {
scriptContext := universalScriptContext(request, modelType, payload)
out, err := executor.Execute(ctx, scriptengine.Options{
Script: scriptText,
Args: []any{upstreamTaskID, scriptContext},
ContextData: scriptContext,
ScriptName: "custom_poll_script:" + modelType,
PreferredEntryNames: []string{"pollTask", "poll", "main", "handler"},
Timeout: 30 * time.Second,
HTTPClient: httpClient(request.HTTPClient, c.HTTPClient),
})
if err != nil {
return nil, "", universalScriptError(err)
}
result, ok := out.(map[string]any)
if !ok || result == nil {
return nil, "", &ClientError{Code: "invalid_response", Message: "custom poll script must return an object", Retryable: false}
}
return result, requestIDFromResult(result), nil
}
pollURL := resolveUniversalTaskURL(request.Candidate.PlatformConfig, upstreamTaskID)
if pollURL == "" {
return nil, "", &ClientError{Code: "missing_configuration", Message: "universal getTaskURL is required", Retryable: false}
}
return universalGetJSON(ctx, httpClient(request.HTTPClient, c.HTTPClient), pollURL, request.Candidate.Credentials)
}
func universalScriptContext(request Request, modelType string, payload map[string]any) map[string]any {
baseURL := strings.TrimRight(strings.TrimSpace(request.Candidate.BaseURL), "/")
getTaskURL := universalConfigString(request.Candidate.PlatformConfig, "getTaskURL", "get_task_url")
context := map[string]any{
"__easyaiScriptContext": true,
"baseURL": baseURL,
"getTaskURL": getTaskURL,
"authValues": cloneMapAny(request.Candidate.Credentials),
"headers": map[string]any{},
"payload": cloneMapAny(payload),
"type": modelType,
"options": map[string]any{
"task_id": request.RemoteTaskID,
"upstream_task_id": request.RemoteTaskID,
"model": request.Model,
"providerModelName": request.Candidate.ProviderModelName,
"platformId": request.Candidate.PlatformID,
"platformModelId": request.Candidate.PlatformModelID,
"canonicalModelKey": request.Candidate.CanonicalModelKey,
"modelType": modelType,
"timeout": universalDurationConfig(request.Candidate.PlatformConfig, 10*time.Minute, "pollTimeoutMs", "poll_timeout_ms").Milliseconds(),
},
"env": cloneMapAny(request.Candidate.PlatformConfig),
"candidate": universalCandidateSnapshot(request),
}
context["createRequestURL"] = func(path string, base ...string) string {
selectedBase := baseURL
if len(base) > 0 && strings.TrimSpace(base[0]) != "" {
selectedBase = strings.TrimRight(strings.TrimSpace(base[0]), "/")
}
if strings.HasPrefix(path, "http://") || strings.HasPrefix(path, "https://") {
return path
}
return selectedBase + "/" + strings.TrimLeft(path, "/")
}
context["creatRequestURL"] = context["createRequestURL"]
context["resolveGetTaskURL"] = func(taskID string) string {
return resolveUniversalTaskURL(request.Candidate.PlatformConfig, taskID)
}
return context
}
func universalCandidateSnapshot(request Request) map[string]any {
return map[string]any{
"modelName": request.Candidate.ModelName,
"modelAlias": request.Candidate.ModelAlias,
"providerModelName": request.Candidate.ProviderModelName,
"provider": request.Candidate.Provider,
"platformId": request.Candidate.PlatformID,
"platformModelId": request.Candidate.PlatformModelID,
"capabilities": cloneMapAny(request.Candidate.Capabilities),
}
}
func universalDefaultPayload(request Request) map[string]any {
body := cloneBody(request.Body)
body["model"] = upstreamModelName(request.Candidate)
if request.Kind == "images.generations" {
if n := firstPresent(body["n"], body["numImages"]); n != nil {
body["numImages"] = n
}
if aspectRatio := strings.TrimSpace(stringFromAny(body["aspect_ratio"])); aspectRatio != "" {
body["aspectRatio"] = aspectRatio
}
}
return body
}
func universalSubmitEndpoint(request Request) string {
if endpoint := universalConfigString(request.Candidate.PlatformConfig, "submitPath", "submit_path"); endpoint != "" {
return endpoint
}
switch request.Kind {
case "images.generations":
return "/images/generations"
case "images.edits":
return "/images/edits"
case "videos.generations":
return "/video/generations"
default:
return "/" + strings.ReplaceAll(request.Kind, ".", "/")
}
}
func universalPostJSON(ctx context.Context, client *http.Client, baseURL string, endpoint string, body map[string]any, credentials map[string]any) (map[string]any, string, error) {
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, providerURL(baseURL, endpoint), bytes.NewReader(raw))
if err != nil {
return nil, "", err
}
req.Header.Set("Content-Type", "application/json")
if apiKey := credential(credentials, "apiKey", "api_key", "key", "token"); apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
return result, requestID, err
}
func universalGetJSON(ctx context.Context, client *http.Client, url string, credentials map[string]any) (map[string]any, string, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, url, nil)
if err != nil {
return nil, "", err
}
if apiKey := credential(credentials, "apiKey", "api_key", "key", "token"); apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
resp, err := client.Do(req)
if err != nil {
return nil, "", &ClientError{Code: "network", Message: err.Error(), Retryable: true}
}
requestID := requestIDFromHTTPResponse(resp)
result, err := decodeHTTPResponse(resp)
return result, requestID, err
}
func normalizeUniversalResult(request Request, result map[string]any, upstreamTaskID string) map[string]any {
out := cloneMapAny(result)
if out["created"] == nil {
out["created"] = time.Now().UnixMilli()
}
if out["task_id"] == nil {
out["task_id"] = upstreamTaskID
}
if out["upstream_task_id"] == nil {
out["upstream_task_id"] = upstreamTaskID
}
if out["model"] == nil {
out["model"] = request.Model
}
if out["status"] == nil {
out["status"] = "success"
}
return out
}
func universalScriptError(err error) error {
var scriptErr *scriptengine.Error
if strings.TrimSpace(err.Error()) == "" {
return &ClientError{Code: "script_error", Message: "script execution failed", Retryable: false}
}
if errors.As(err, &scriptErr) {
return &ClientError{Code: scriptErr.ErrorCode(), Message: scriptErr.Error(), Retryable: scriptErr.ErrorCode() == "script_timeout"}
}
return &ClientError{Code: "script_error", Message: err.Error(), Retryable: false}
}
func universalFailureError(result map[string]any, requestID string, startedAt time.Time) error {
message := firstNonEmptyString(result["message"], result["error"], result["error_message"], "universal task failed")
return &ClientError{
Code: firstNonEmptyString(result["code"], result["error_code"], "provider_failed"),
Message: message,
RequestID: requestID,
ResponseStartedAt: startedAt,
ResponseFinishedAt: time.Now(),
ResponseDurationMS: responseDurationMS(startedAt, time.Now()),
Retryable: false,
}
}
func isUniversalSuccess(result map[string]any) bool {
switch universalStatus(result) {
case "success", "succeeded", "completed", "complete", "done":
return true
default:
return false
}
}
func isUniversalFailure(result map[string]any) bool {
switch universalStatus(result) {
case "failed", "failure", "error", "cancelled", "canceled":
return true
default:
return false
}
}
func universalStatus(result map[string]any) string {
return strings.ToLower(strings.TrimSpace(firstNonEmptyString(result["status"], result["state"], result["task_status"])))
}
func universalTaskID(result map[string]any) string {
return firstNonEmptyString(result["upstream_task_id"], result["task_id"], result["taskId"], result["id"])
}
func universalProgress(request Request, upstreamTaskID string) []Progress {
progress := providerProgress(request)
progress = append(progress, Progress{Phase: "polling", Progress: 0.65, Message: "provider task polled", Payload: map[string]any{"upstreamTaskId": upstreamTaskID}})
return progress
}
func universalStripPrivatePayload(payload map[string]any) map[string]any {
out := cloneMapAny(payload)
for _, key := range []string{"_originalParams", "_resolution", "_duration"} {
delete(out, key)
}
return out
}
func universalSceneScript(config map[string]any, modelType string, keys ...string) string {
for _, key := range keys {
value := config[key]
switch typed := value.(type) {
case string:
if strings.TrimSpace(typed) != "" {
return strings.TrimSpace(typed)
}
case map[string]any:
if script := firstNonEmptyString(typed[modelType], typed["common"]); script != "" {
return script
}
}
}
return ""
}
func universalConfigString(config map[string]any, keys ...string) string {
for _, key := range keys {
if value := strings.TrimSpace(fmt.Sprint(config[key])); value != "" && value != "<nil>" {
return value
}
}
return ""
}
func universalDurationConfig(config map[string]any, fallback time.Duration, keys ...string) time.Duration {
for _, key := range keys {
switch value := config[key].(type) {
case int:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case int64:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case float64:
if value > 0 {
return time.Duration(value) * time.Millisecond
}
case string:
if parsed, err := time.ParseDuration(value); err == nil && parsed > 0 {
return parsed
}
}
}
return fallback
}
func resolveUniversalTaskURL(config map[string]any, upstreamTaskID string) string {
template := universalConfigString(config, "getTaskURL", "get_task_url")
out := strings.TrimSpace(template)
replacements := [][2]string{
{"${upstream_task_id}", upstreamTaskID},
{"{{upstream_task_id}}", upstreamTaskID},
{"{upstream_task_id}", upstreamTaskID},
{"${task_id}", upstreamTaskID},
{"{{task_id}}", upstreamTaskID},
{"{task_id}", upstreamTaskID},
{"${taskId}", upstreamTaskID},
{"${taskID}", upstreamTaskID},
{"{{taskId}}", upstreamTaskID},
{"{{taskID}}", upstreamTaskID},
{"{taskId}", upstreamTaskID},
{"{taskID}", upstreamTaskID},
}
for _, replacement := range replacements {
out = strings.ReplaceAll(out, replacement[0], replacement[1])
}
return out
}
+132
View File
@@ -0,0 +1,132 @@
package clients
import (
"context"
"net/http"
"net/http/httptest"
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestUniversalClientRunsCustomScripts(t *testing.T) {
request := Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "custom-video",
Body: map[string]any{"model": "custom-video", "prompt": "hello"},
Candidate: testUniversalCandidate(map[string]any{
"customGetParamsScript": map[string]any{
"video_generate": `async function getGenerateParams(params, context) {
return { prompt: params.prompt + "-payload", model: context.candidate.providerModelName };
}`,
},
"customSubmitScript": map[string]any{
"video_generate": `async function submitTask(payload) {
return { status: "submitted", task_id: "task-" + payload.prompt };
}`,
},
"customPollScript": map[string]any{
"video_generate": `async function pollTask(taskId) {
return { status: "success", upstream_task_id: taskId, data: [{ url: "https://cdn.example/video.mp4" }] };
}`,
},
"pollIntervalMs": 1,
"pollTimeoutMs": 1000,
}),
}
var submitted string
request.OnRemoteTaskSubmitted = func(remoteTaskID string, payload map[string]any) error {
submitted = remoteTaskID
return nil
}
response, err := (UniversalClient{}).Run(context.Background(), request)
if err != nil {
t.Fatalf("run failed: %v", err)
}
if submitted != "task-hello-payload" {
t.Fatalf("unexpected remote task id: %q", submitted)
}
if response.Result["upstream_task_id"] != "task-hello-payload" {
t.Fatalf("unexpected result: %#v", response.Result)
}
}
func TestUniversalClientDefaultSubmitAndPoll(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.Header.Get("Authorization") != "Bearer test-key" {
t.Fatalf("missing authorization header")
}
w.Header().Set("Content-Type", "application/json")
switch r.URL.Path {
case "/video/generations":
_, _ = w.Write([]byte(`{"status":"submitted","task_id":"remote-1"}`))
case "/tasks/remote-1":
_, _ = w.Write([]byte(`{"status":"success","data":[{"url":"https://cdn.example/default.mp4"}]}`))
default:
t.Fatalf("unexpected path: %s", r.URL.Path)
}
}))
defer server.Close()
request := Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "default-video",
Body: map[string]any{"model": "default-video", "prompt": "hello"},
Candidate: testUniversalCandidate(map[string]any{
"getTaskURL": server.URL + "/tasks/{{task_id}}",
"pollIntervalMs": 1,
"pollTimeoutMs": 1000,
}),
}
request.Candidate.BaseURL = server.URL
response, err := (UniversalClient{}).Run(context.Background(), request)
if err != nil {
t.Fatalf("run failed: %v", err)
}
if response.Result["upstream_task_id"] != "remote-1" {
t.Fatalf("unexpected result: %#v", response.Result)
}
}
func TestUniversalClientResumeSkipsSubmit(t *testing.T) {
request := Request{
Kind: "videos.generations",
ModelType: "video_generate",
Model: "resume-video",
Body: map[string]any{"model": "resume-video", "prompt": "hello"},
RemoteTaskID: "existing-1",
RemoteTaskPayload: map[string]any{"payload": map[string]any{"prompt": "old"}},
Candidate: testUniversalCandidate(map[string]any{
"customSubmitScript": `async function submitTask() { throw new Error("submit should not run"); }`,
"customPollScript": `async function pollTask(taskId) { return { status: "success", upstream_task_id: taskId, data: [{ url: "https://cdn.example/resume.mp4" }] }; }`,
"pollIntervalMs": 1,
"pollTimeoutMs": 1000,
}),
}
response, err := (UniversalClient{}).Run(context.Background(), request)
if err != nil {
t.Fatalf("run failed: %v", err)
}
if response.Result["upstream_task_id"] != "existing-1" {
t.Fatalf("unexpected result: %#v", response.Result)
}
}
func testUniversalCandidate(config map[string]any) store.RuntimeModelCandidate {
return store.RuntimeModelCandidate{
Provider: "universal",
SpecType: "universal",
BaseURL: "https://provider.example",
Credentials: map[string]any{"apiKey": "test-key"},
PlatformConfig: config,
ModelName: "alias-model",
ProviderModelName: "provider-model",
ModelType: "video_generate",
ClientID: "universal:test",
}
}