feat(api): migrate media clients and universal scripts

2026-05-20 21:29:07 +08:00
parent 11a2c13e4a
commit af9b281d34
14 changed files with 2648 additions and 18 deletions
@@ -0,0 +1,196 @@
+package runner
+
+import (
+	"context"
+	"fmt"
+	"reflect"
+	"strings"
+
+	scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+func (s *Service) preprocessRequestWithScripts(ctx context.Context, kind string, body map[string]any, candidate store.RuntimeModelCandidate) parameterPreprocessResult {
+	if platformConfigBool(candidate.PlatformConfig, "skipParamNormalization", "skip_param_normalization") {
+		modelType := strings.TrimSpace(candidate.ModelType)
+		if modelType == "" {
+			modelType = modelTypeFromKind(kind, body)
+		}
+		input := cloneMap(body)
+		return parameterPreprocessResult{
+			Body: cloneMap(body),
+			Log: parameterPreprocessingLog{
+				ModelType: modelType,
+				Input:     input,
+				Output:    cloneMap(body),
+				Changed:   false,
+				Changes:   []parameterPreprocessChange{},
+				Model:     preprocessingModelSnapshot(candidate),
+			},
+		}
+	}
+
+	result := preprocessRequestWithLog(kind, body, candidate)
+	if result.Err != nil {
+		return result
+	}
+	scriptText := platformConfigString(candidate.PlatformConfig, "customPreprocessScript", "custom_preprocess_script")
+	if strings.TrimSpace(scriptText) == "" || s.scriptExecutor == nil {
+		return result
+	}
+
+	before := cloneMap(result.Body)
+	scriptContext := s.scriptContext(candidate, result.Log.ModelType, nil, map[string]any{
+		"modelCapability": effectiveModelCapability(candidate),
+		"platformModel":   result.Log.Model,
+		"platform":        candidate.PlatformConfig,
+	})
+	out, err := s.scriptExecutor.Execute(ctx, scriptengine.Options{
+		Script:              scriptText,
+		Args:                []any{cloneMap(result.Body), result.Log.ModelType, scriptContext},
+		ContextData:         scriptContext,
+		ScriptName:          "custom_preprocess_script:" + result.Log.ModelType,
+		PreferredEntryNames: []string{"preprocessParams", "preprocess", "main", "handler"},
+		Timeout:             scriptengine.PreprocessTimeout,
+	})
+	if err != nil {
+		result.Log.recordScriptChange("CustomPreprocessScript", "error", "$", before, result.Body, err.Error())
+		result.Log.Output = cloneMap(result.Body)
+		result.Log.Changed = len(result.Log.Changes) > 0
+		result.Err = err
+		return result
+	}
+	rewritten, ok := out.(map[string]any)
+	if !ok || rewritten == nil {
+		result.Log.Output = cloneMap(result.Body)
+		result.Log.Changed = len(result.Log.Changes) > 0
+		return result
+	}
+	merged := cloneMap(result.Body)
+	for key, value := range rewritten {
+		merged[key] = value
+	}
+	if !mapsEqual(before, merged) {
+		result.Log.recordScriptChange("CustomPreprocessScript", "rewrite", "$", before, merged, "platform custom preprocess script returned parameter updates")
+	}
+	result.Body = merged
+	result.Log.Output = cloneMap(merged)
+	result.Log.Changed = len(result.Log.Changes) > 0
+	return result
+}
+
+func (s *Service) scriptContext(candidate store.RuntimeModelCandidate, modelType string, payload map[string]any, extra map[string]any) map[string]any {
+	getTaskURL := platformConfigString(candidate.PlatformConfig, "getTaskURL", "get_task_url")
+	baseURL := strings.TrimRight(strings.TrimSpace(candidate.BaseURL), "/")
+	env := cloneMap(candidate.PlatformConfig)
+	context := map[string]any{
+		"__easyaiScriptContext": true,
+		"baseURL":               baseURL,
+		"getTaskURL":            getTaskURL,
+		"authValues":            cloneMap(candidate.Credentials),
+		"headers":               map[string]any{},
+		"payload":               cloneMap(payload),
+		"type":                  modelType,
+		"options": map[string]any{
+			"model":              candidate.ModelName,
+			"providerModelName":  candidate.ProviderModelName,
+			"platformId":         candidate.PlatformID,
+			"platformModelId":    candidate.PlatformModelID,
+			"canonicalModelKey":  candidate.CanonicalModelKey,
+			"sourceProviderCode": candidate.Provider,
+		},
+		"env":       env,
+		"candidate": preprocessingModelSnapshot(candidate),
+	}
+	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 resolveTaskURLTemplate(getTaskURL, taskID, "")
+	}
+	for key, value := range extra {
+		context[key] = value
+	}
+	return context
+}
+
+func preprocessingModelSnapshot(candidate store.RuntimeModelCandidate) map[string]any {
+	return map[string]any{
+		"modelName":         candidate.ModelName,
+		"modelAlias":        candidate.ModelAlias,
+		"providerModelName": candidate.ProviderModelName,
+		"provider":          candidate.Provider,
+		"platformId":        candidate.PlatformID,
+		"platformModelId":   candidate.PlatformModelID,
+		"capabilities":      cloneMap(candidate.Capabilities),
+	}
+}
+
+func (log *parameterPreprocessingLog) recordScriptChange(processor string, action string, path string, before any, after any, reason string) {
+	if log == nil {
+		return
+	}
+	log.Changes = append(log.Changes, parameterPreprocessChange{
+		Processor: processor,
+		Action:    action,
+		Path:      path,
+		Before:    cloneAny(before),
+		After:     cloneAny(after),
+		Reason:    reason,
+	})
+}
+
+func platformConfigString(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 platformConfigBool(config map[string]any, keys ...string) bool {
+	for _, key := range keys {
+		switch value := config[key].(type) {
+		case bool:
+			return value
+		case string:
+			return strings.EqualFold(strings.TrimSpace(value), "true")
+		}
+	}
+	return false
+}
+
+func resolveTaskURLTemplate(template string, upstreamTaskID string, taskID string) string {
+	out := strings.TrimSpace(template)
+	replacements := [][2]string{
+		{"${upstream_task_id}", upstreamTaskID},
+		{"{{upstream_task_id}}", upstreamTaskID},
+		{"{upstream_task_id}", upstreamTaskID},
+		{"${task_id}", taskID},
+		{"{{task_id}}", taskID},
+		{"{task_id}", taskID},
+		{"${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
+}
+
+func mapsEqual(left map[string]any, right map[string]any) bool {
+	return reflect.DeepEqual(left, right)
+}
@@ -0,0 +1,64 @@
+package runner
+
+import (
+	"context"
+	"testing"
+
+	scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+func TestPreprocessRequestWithCustomScript(t *testing.T) {
+	service := &Service{scriptExecutor: &scriptengine.Executor{}}
+	candidate := store.RuntimeModelCandidate{
+		Provider:  "universal",
+		ModelName: "image-model",
+		ModelType: "image_generate",
+		Capabilities: map[string]any{
+			"image_generate": map[string]any{"max_output_images": 4},
+		},
+		PlatformConfig: map[string]any{
+			"customPreprocessScript": `(params, type, context) => {
+				return { prompt: params.prompt + "-" + type, n: 2, provider: context.candidate.provider };
+			}`,
+		},
+	}
+
+	result := service.preprocessRequestWithScripts(context.Background(), "images.generations", map[string]any{"prompt": "hello", "n": 8}, candidate)
+	if result.Err != nil {
+		t.Fatalf("unexpected preprocess error: %v", result.Err)
+	}
+	if result.Body["prompt"] != "hello-image_generate" || result.Body["n"].(float64) != 2 {
+		t.Fatalf("unexpected body: %#v", result.Body)
+	}
+	if !result.Log.Changed || len(result.Log.Changes) == 0 {
+		t.Fatalf("expected script change in log: %#v", result.Log)
+	}
+}
+
+func TestPreprocessRequestSkipParamNormalizationSkipsCustomScript(t *testing.T) {
+	service := &Service{scriptExecutor: &scriptengine.Executor{}}
+	candidate := store.RuntimeModelCandidate{
+		ModelName: "image-model",
+		ModelType: "image_generate",
+		Provider:  "universal",
+		Capabilities: map[string]any{
+			"image_generate": map[string]any{"max_output_images": 1},
+		},
+		PlatformConfig: map[string]any{
+			"skipParamNormalization": true,
+			"customPreprocessScript": `(params) => ({ prompt: "changed", n: 1 })`,
+		},
+	}
+
+	result := service.preprocessRequestWithScripts(context.Background(), "images.generations", map[string]any{"prompt": "hello", "n": 9}, candidate)
+	if result.Err != nil {
+		t.Fatalf("unexpected preprocess error: %v", result.Err)
+	}
+	if result.Body["prompt"] != "hello" || result.Body["n"].(int) != 9 {
+		t.Fatalf("skip should keep raw body, got %#v", result.Body)
+	}
+	if result.Log.Changed || len(result.Log.Changes) != 0 {
+		t.Fatalf("skip should not record changes: %#v", result.Log)
+	}
+}
@@ -12,18 +12,20 @@ import (
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/clients"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
+	scriptengine "github.com/easyai/easyai-ai-gateway/apps/api/internal/script"
 	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
 	"github.com/jackc/pgx/v5"
 	"github.com/riverqueue/river"
 )

 type Service struct {
-	cfg         config.Config
-	store       *store.Store
-	logger      *slog.Logger
-	clients     map[string]clients.Client
-	httpClients *httpClientCache
-	riverClient *river.Client[pgx.Tx]
+	cfg            config.Config
+	store          *store.Store
+	logger         *slog.Logger
+	clients        map[string]clients.Client
+	scriptExecutor *scriptengine.Executor
+	httpClients    *httpClientCache
+	riverClient    *river.Client[pgx.Tx]
 }

 type Result struct {
@@ -47,17 +49,29 @@ func (e *TaskQueuedError) Is(target error) bool {

 func New(cfg config.Config, db *store.Store, logger *slog.Logger) *Service {
 	httpClients := newHTTPClientCache()
+	scriptExecutor := &scriptengine.Executor{Logger: logger}
 	return &Service{
-		cfg:    cfg,
-		store:  db,
-		logger: logger,
+		cfg:            cfg,
+		store:          db,
+		logger:         logger,
+		scriptExecutor: scriptExecutor,
 		clients: map[string]clients.Client{
-			"openai":     clients.OpenAIClient{HTTPClient: httpClients.none},
-			"gemini":     clients.GeminiClient{HTTPClient: httpClients.none},
-			"volces":     clients.VolcesClient{HTTPClient: httpClients.none},
-			"keling":     clients.KelingClient{HTTPClient: httpClients.none},
-			"kling":      clients.KelingClient{HTTPClient: httpClients.none},
-			"simulation": clients.SimulationClient{},
+			"openai":                clients.OpenAIClient{HTTPClient: httpClients.none},
+			"aliyun-bailian":        clients.AliyunBailianClient{HTTPClient: httpClients.none},
+			"blackforest":           clients.BlackforestClient{HTTPClient: httpClients.none},
+			"gemini":                clients.GeminiClient{HTTPClient: httpClients.none},
+			"jimeng":                clients.JimengClient{HTTPClient: httpClients.none},
+			"midjourney":            clients.MidjourneyClient{HTTPClient: httpClients.none},
+			"minimax":               clients.MinimaxClient{HTTPClient: httpClients.none},
+			"newapi":                clients.NewAPIClient{HTTPClient: httpClients.none},
+			"tencent-hunyuan-image": clients.HunyuanImageClient{HTTPClient: httpClients.none},
+			"tencent-hunyuan-video": clients.HunyuanVideoClient{HTTPClient: httpClients.none},
+			"vidu":                  clients.ViduClient{HTTPClient: httpClients.none},
+			"volces":                clients.VolcesClient{HTTPClient: httpClients.none},
+			"keling":                clients.KelingClient{HTTPClient: httpClients.none},
+			"kling":                 clients.KelingClient{HTTPClient: httpClients.none},
+			"universal":             clients.UniversalClient{HTTPClient: httpClients.none, ScriptExecutor: scriptExecutor},
+			"simulation":            clients.SimulationClient{},
 		},
 		httpClients: httpClients,
 	}
@@ -147,7 +161,7 @@ func (s *Service) execute(ctx context.Context, task store.GatewayTask, user *aut
 	attemptNo := task.AttemptCount
 	var firstPreprocessing parameterPreprocessingLog
 	if len(candidates) > 0 {
-		preprocessing := preprocessRequestWithLog(task.Kind, body, candidates[0])
+		preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidates[0])
 		firstCandidateBody = preprocessing.Body
 		firstPreprocessing = preprocessing.Log
 		normalizedModelType = candidates[0].ModelType
@@ -225,7 +239,7 @@ candidatesLoop:
 		var candidateErr error
 		for clientAttempt := 1; clientAttempt <= clientAttempts; clientAttempt++ {
 			nextAttemptNo := attemptNo + 1
-			preprocessing := preprocessRequestWithLog(task.Kind, body, candidate)
+			preprocessing := s.preprocessRequestWithScripts(ctx, task.Kind, body, candidate)
 			preprocessingLog := preprocessing.Log
 			lastPreprocessing = &preprocessingLog
 			if preprocessing.Err != nil {
@@ -1090,8 +1104,13 @@ func parameterPreprocessClientError(err error) *clients.ClientError {
 	if err == nil {
 		return nil
 	}
+	code := "invalid_parameter"
+	var coded interface{ ErrorCode() string }
+	if errors.As(err, &coded) && strings.TrimSpace(coded.ErrorCode()) != "" {
+		code = coded.ErrorCode()
+	}
 	return &clients.ClientError{
-		Code:       "invalid_parameter",
+		Code:       code,
 		Message:    err.Error(),
 		StatusCode: 400,
 		Retryable:  false,