feat: improve model catalog aggregation

This commit is contained in:
2026-05-11 17:44:57 +08:00
parent ec87816c95
commit 0431cb8157
41 changed files with 4745 additions and 550 deletions
+19 -15
View File
@@ -127,15 +127,16 @@ func TestOpenAIClientChatContract(t *testing.T) {
Model: "openai:gpt-4o-mini",
Body: map[string]any{"model": "openai:gpt-4o-mini", "messages": []any{map[string]any{"role": "user", "content": "ping"}}},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "gpt-4o-mini",
Credentials: map[string]any{"apiKey": "test-key"},
BaseURL: server.URL,
ModelName: "gpt-4o-mini",
ProviderModelName: "openai-compatible-gpt-4o-mini",
Credentials: map[string]any{"apiKey": "test-key"},
},
})
if err != nil {
t.Fatalf("run openai client: %v", err)
}
if gotPath != "/chat/completions" || gotAuth != "Bearer test-key" || gotModel != "gpt-4o-mini" {
if gotPath != "/chat/completions" || gotAuth != "Bearer test-key" || gotModel != "openai-compatible-gpt-4o-mini" {
t.Fatalf("unexpected request path=%s auth=%s model=%s", gotPath, gotAuth, gotModel)
}
if response.Usage.TotalTokens != 5 || response.Result["id"] != "chatcmpl-test" {
@@ -231,16 +232,17 @@ func TestGeminiClientChatContract(t *testing.T) {
"messages": []any{map[string]any{"role": "user", "content": "ping"}},
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "gemini-2.5-flash",
ModelType: "chat",
Credentials: map[string]any{"apiKey": "gemini-key"},
BaseURL: server.URL,
ModelName: "gemini-2.5-flash",
ProviderModelName: "gemini-compatible-2.5-flash",
ModelType: "chat",
Credentials: map[string]any{"apiKey": "gemini-key"},
},
})
if err != nil {
t.Fatalf("run gemini client: %v", err)
}
if gotPath != "/v1beta/models/gemini-2.5-flash:generateContent" || gotKey != "gemini-key" || gotText != "ping" {
if gotPath != "/v1beta/models/gemini-compatible-2.5-flash:generateContent" || gotKey != "gemini-key" || gotText != "ping" {
t.Fatalf("unexpected request path=%s key=%s text=%s", gotPath, gotKey, gotText)
}
if response.Usage.TotalTokens != 10 || extractText(response.Result) != "gemini ok" {
@@ -290,9 +292,10 @@ func TestVolcesClientImageEditUsesGenerationEndpoint(t *testing.T) {
"image": "https://example.com/source.png",
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "doubao-seedream-4-0-250828",
Credentials: map[string]any{"apiKey": "volces-key"},
BaseURL: server.URL,
ModelName: "豆包Seedream-4.0",
ProviderModelName: "doubao-seedream-4-0-250828",
Credentials: map[string]any{"apiKey": "volces-key"},
Capabilities: map[string]any{
"image_edit": map[string]any{"output_multiple_images": true},
},
@@ -366,9 +369,10 @@ func TestVolcesClientVideoSubmitsAndPollsTask(t *testing.T) {
"aspect_ratio": "16:9",
},
Candidate: store.RuntimeModelCandidate{
BaseURL: server.URL,
ModelName: "doubao-seedance-2-0-260128",
Credentials: map[string]any{"apiKey": "volces-key"},
BaseURL: server.URL,
ModelName: "豆包Seedance-2.0",
ProviderModelName: "doubao-seedance-2-0-260128",
Credentials: map[string]any{"apiKey": "volces-key"},
PlatformConfig: map[string]any{
"volcesPollIntervalMs": 100,
"volcesPollTimeoutSeconds": 1,
+1 -1
View File
@@ -22,7 +22,7 @@ func (c GeminiClient) Run(ctx context.Context, request Request) (Response, error
}
body := geminiBody(request)
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, geminiURL(request.Candidate.BaseURL, request.Candidate.ModelName, apiKey), bytes.NewReader(raw))
req, err := http.NewRequestWithContext(ctx, http.MethodPost, geminiURL(request.Candidate.BaseURL, upstreamModelName(request.Candidate), apiKey), bytes.NewReader(raw))
if err != nil {
return Response{}, err
}
+1 -1
View File
@@ -23,7 +23,7 @@ func (c OpenAIClient) Run(ctx context.Context, request Request) (Response, error
return Response{}, &ClientError{Code: "unsupported_kind", Message: "unsupported openai request kind", Retryable: false}
}
body := cloneBody(request.Body)
body["model"] = request.Candidate.ModelName
body["model"] = upstreamModelName(request.Candidate)
stream := request.Stream || boolValue(body, "stream")
raw, _ := json.Marshal(body)
req, err := http.NewRequestWithContext(ctx, http.MethodPost, joinURL(request.Candidate.BaseURL, endpoint), bytes.NewReader(raw))
+8
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"net/http"
"strings"
"time"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
@@ -72,6 +73,13 @@ func IsRetryable(err error) bool {
return errors.As(err, &clientErr) && clientErr.Retryable
}
func upstreamModelName(candidate store.RuntimeModelCandidate) string {
if name := strings.TrimSpace(candidate.ProviderModelName); name != "" {
return name
}
return candidate.ModelName
}
func ErrorCode(err error) string {
var clientErr *ClientError
if errors.As(err, &clientErr) && clientErr.Code != "" {
+3 -3
View File
@@ -177,7 +177,7 @@ func (c VolcesClient) getJSON(ctx context.Context, baseURL string, path string,
func volcesImageBody(request Request) map[string]any {
body := cleanProviderBody(request.Body)
body["model"] = request.Candidate.ModelName
body["model"] = upstreamModelName(request.Candidate)
if _, ok := body["watermark"]; !ok {
body["watermark"] = false
}
@@ -200,7 +200,7 @@ func volcesImageBody(request Request) map[string]any {
func volcesVideoBody(request Request) map[string]any {
body := cleanProviderBody(request.Body)
body["model"] = request.Candidate.ModelName
body["model"] = upstreamModelName(request.Candidate)
content := contentItems(body["content"])
if len(content) == 0 {
content = buildVolcesContentFromBody(body)
@@ -515,7 +515,7 @@ func volcesVideoSuccessResult(request Request, upstreamTaskID string, raw map[st
"id": upstreamTaskID,
"object": "video.generation",
"created": created,
"model": request.Candidate.ModelName,
"model": upstreamModelName(request.Candidate),
"status": "succeeded",
"upstream_task_id": upstreamTaskID,
"data": data,
+80 -14
View File
@@ -264,7 +264,7 @@ func (s *Server) createPlatformModel(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "create platform model failed")
return
}
writeJSON(w, http.StatusCreated, model)
writeJSON(w, http.StatusCreated, s.platformModelResponse(r.Context(), model))
}
func (s *Server) replacePlatformModels(w http.ResponseWriter, r *http.Request) {
@@ -292,7 +292,7 @@ func (s *Server) replacePlatformModels(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "replace platform models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": models})
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
func (s *Server) deletePlatformModel(w http.ResponseWriter, r *http.Request) {
@@ -315,7 +315,7 @@ func (s *Server) listModels(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "list models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": models})
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
@@ -326,7 +326,7 @@ func (s *Server) listPlayableModels(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusInternalServerError, "list playable models failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": models})
writeJSON(w, http.StatusOK, map[string]any{"items": s.platformModelResponses(r.Context(), models)})
}
func (s *Server) listPricingRules(w http.ResponseWriter, r *http.Request) {
@@ -623,22 +623,88 @@ func (s *Server) listTasks(w http.ResponseWriter, r *http.Request) {
writeError(w, http.StatusUnauthorized, "unauthorized")
return
}
limit := 50
if raw := strings.TrimSpace(r.URL.Query().Get("limit")); raw != "" {
parsed, err := strconv.Atoi(raw)
if err != nil || parsed <= 0 {
writeError(w, http.StatusBadRequest, "invalid limit")
return
}
limit = parsed
query := r.URL.Query()
page, err := positiveQueryInt(query.Get("page"), 1)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid page")
return
}
tasks, err := s.store.ListTasks(r.Context(), user, limit)
pageSizeRaw := query.Get("pageSize")
if pageSizeRaw == "" {
pageSizeRaw = query.Get("limit")
}
pageSize, err := positiveQueryInt(pageSizeRaw, 50)
if err != nil {
writeError(w, http.StatusBadRequest, "invalid pageSize")
return
}
createdFrom, err := parseTaskListTime(query.Get("createdFrom"), query.Get("from"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdFrom")
return
}
createdTo, err := parseTaskListTime(query.Get("createdTo"), query.Get("to"))
if err != nil {
writeError(w, http.StatusBadRequest, "invalid createdTo")
return
}
result, err := s.store.ListTasks(r.Context(), user, store.TaskListFilter{
Query: firstNonEmpty(query.Get("q"), query.Get("query")),
ModelType: firstNonEmpty(query.Get("modelType"), query.Get("type")),
CreatedFrom: createdFrom,
CreatedTo: createdTo,
Page: page,
PageSize: pageSize,
})
if err != nil {
s.logger.Error("list tasks failed", "error", err)
writeError(w, http.StatusInternalServerError, "list tasks failed")
return
}
writeJSON(w, http.StatusOK, map[string]any{"items": tasks})
writeJSON(w, http.StatusOK, map[string]any{
"items": result.Items,
"total": result.Total,
"page": result.Page,
"pageSize": result.PageSize,
})
}
func positiveQueryInt(raw string, fallback int) (int, error) {
raw = strings.TrimSpace(raw)
if raw == "" {
return fallback, nil
}
value, err := strconv.Atoi(raw)
if err != nil || value <= 0 {
return 0, fmt.Errorf("invalid positive integer")
}
return value, nil
}
func parseTaskListTime(values ...string) (*time.Time, error) {
raw := strings.TrimSpace(firstNonEmpty(values...))
if raw == "" {
return nil, nil
}
layouts := []string{time.RFC3339Nano, time.RFC3339, "2006-01-02T15:04", "2006-01-02 15:04:05", "2006-01-02"}
var lastErr error
for _, layout := range layouts {
parsed, err := time.ParseInLocation(layout, raw, time.Local)
if err == nil {
return &parsed, nil
}
lastErr = err
}
return nil, lastErr
}
func firstNonEmpty(values ...string) string {
for _, value := range values {
if trimmed := strings.TrimSpace(value); trimmed != "" {
return trimmed
}
}
return ""
}
func boolValue(body map[string]any, key string) bool {
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,221 @@
package httpapi
import (
"testing"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func TestBuildModelCatalogAggregatesSources(t *testing.T) {
models := []store.PlatformModel{
{
ID: "model-a",
PlatformID: "platform-a",
ModelName: "seedance",
ModelAlias: "Seedance-2.0",
ModelType: store.StringList{"image_generate"},
DisplayName: "Seedance Source A",
BillingConfig: map[string]any{
"image": map[string]any{"basePrice": float64(10), "dynamicWeight": map[string]any{"1K": float64(1), "2K": float64(2)}},
},
RateLimitPolicy: map[string]any{
"platformLimits": map[string]any{
"max_request_per_minute": 60,
"max_token_per_minute": 1000,
"max_concurrent_requests": 2,
},
},
PricingMode: "inherit_discount",
Enabled: true,
},
{
ID: "model-b",
PlatformID: "platform-b",
ModelName: "seedance",
ModelAlias: "Seedance-2.0",
ModelType: store.StringList{"image_generate"},
DisplayName: "Seedance Source B",
BillingConfig: map[string]any{
"image": map[string]any{"basePrice": float64(10), "dynamicWeight": map[string]any{"1K": float64(1), "2K": float64(2)}},
},
RateLimitPolicy: map[string]any{
"rpm": 40,
"tpm": 2000,
"concurrent": 3,
},
DiscountFactor: 0.8,
PricingMode: "custom",
Enabled: true,
},
}
platforms := []store.Platform{
{ID: "platform-a", Provider: "volces", Name: "火山引擎", Status: "enabled", Priority: 20, DefaultDiscountFactor: 1},
{ID: "platform-b", Provider: "gemini", Name: "Gemini", Status: "enabled", Priority: 10, DefaultDiscountFactor: 1},
}
providers := []store.CatalogProvider{
{ProviderKey: "volces", DisplayName: "火山引擎", IconPath: "volces.png"},
{ProviderKey: "gemini", DisplayName: "Google Gemini", IconPath: "gemini.png"},
}
accessRules := []store.AccessRule{
{SubjectType: "user_group", SubjectID: "group-vip", ResourceType: "platform", ResourceID: "platform-b", Effect: "allow", Status: "active"},
{SubjectType: "user_group", SubjectID: "group-blocked", ResourceType: "platform", ResourceID: "platform-a", Effect: "deny", Status: "active"},
}
userGroups := []store.UserGroup{
{ID: "group-vip", GroupKey: "vip", Name: "VIP 用户组"},
{ID: "group-blocked", GroupKey: "blocked", Name: "Blocked 用户组"},
}
baseModels := []store.BaseModel{
{ID: "", Metadata: map[string]any{"description": "高质量图像生成模型"}},
}
response := buildModelCatalog(models, platforms, providers, nil, accessRules, userGroups, baseModels)
if response.Summary.ModelCount != 1 || response.Summary.SourceCount != 2 {
t.Fatalf("unexpected summary: %+v", response.Summary)
}
item := response.Items[0]
if item.SourceCount != 2 {
t.Fatalf("expected merged source count, got %d", item.SourceCount)
}
if item.Source.Label != "2 个源" {
t.Fatalf("expected source label to only show count, got %q", item.Source.Label)
}
if item.RateLimits.RPM == nil || *item.RateLimits.RPM != 100 {
t.Fatalf("expected summed rpm 100, got %+v", item.RateLimits.RPM)
}
if item.RateLimits.TPM == nil || *item.RateLimits.TPM != 3000 {
t.Fatalf("expected summed tpm 3000, got %+v", item.RateLimits.TPM)
}
if item.RateLimits.Concurrent == nil || *item.RateLimits.Concurrent != 5 {
t.Fatalf("expected summed concurrency 5, got %+v", item.RateLimits.Concurrent)
}
if item.Permission.Label != "用户组 VIP 用户组;拒绝 Blocked 用户组" {
t.Fatalf("expected permission label from access rules, got %q", item.Permission.Label)
}
if len(item.Permission.AllowGroups) != 1 || item.Permission.AllowGroups[0] != "VIP 用户组" {
t.Fatalf("expected allow permission groups, got %+v", item.Permission.AllowGroups)
}
if len(item.Permission.DenyGroups) != 1 || item.Permission.DenyGroups[0] != "Blocked 用户组" {
t.Fatalf("expected deny permission groups, got %+v", item.Permission.DenyGroups)
}
if item.Discount.Label != "80% - 无折扣" {
t.Fatalf("expected friendly discount label, got %q", item.Discount.Label)
}
if len(item.ProviderKeys) != 2 {
t.Fatalf("expected both providers on merged item, got %+v", item.ProviderKeys)
}
if !hasFilterCount(response.Filters.Providers, "volces", 1) || !hasFilterCount(response.Filters.Providers, "gemini", 1) {
t.Fatalf("expected provider filters to count merged model for each provider: %+v", response.Filters.Providers)
}
if !hasFilterCount(response.Filters.Capabilities, "image", 1) {
t.Fatalf("expected image capability filter: %+v", response.Filters.Capabilities)
}
if got := item.Pricing.Lines[0]; got != "图像:1K 10 / 2K 20" {
t.Fatalf("unexpected pricing line %q", got)
}
}
func TestBuildModelCatalogUsesBaseModelProviderForProviderFilters(t *testing.T) {
models := []store.PlatformModel{
{
ID: "glm-volces",
PlatformID: "platform-volces",
BaseModelID: "base-glm",
ModelName: "glm-4.7",
ModelAlias: "GLM-4.7",
ModelType: store.StringList{"text_generate"},
DisplayName: "GLM-4.7",
Enabled: true,
},
{
ID: "glm-zhipu",
PlatformID: "platform-zhipu",
BaseModelID: "base-glm",
ModelName: "glm-4.7",
ModelAlias: "GLM-4.7",
ModelType: store.StringList{"text_generate"},
DisplayName: "GLM-4.7",
Enabled: true,
},
}
platforms := []store.Platform{
{ID: "platform-volces", Provider: "volces-openai", Name: "火山引擎(OpenAI兼容)", Status: "enabled"},
{ID: "platform-zhipu", Provider: "zhipu-openai", Name: "智谱官方", Status: "enabled"},
}
providers := []store.CatalogProvider{
{ProviderKey: "volces-openai", DisplayName: "火山引擎(OpenAI兼容)", IconPath: "volces.png"},
{ProviderKey: "zhipu-openai", DisplayName: "智谱AI", IconPath: "zhipu.png"},
}
baseModels := []store.BaseModel{
{ID: "base-glm", ProviderKey: "zhipu-openai", ProviderModelName: "glm-4.7", ModelAlias: "GLM-4.7"},
}
response := buildModelCatalog(models, platforms, providers, nil, nil, nil, baseModels)
if response.Summary.ModelCount != 1 || response.Summary.SourceCount != 2 {
t.Fatalf("unexpected summary: %+v", response.Summary)
}
item := response.Items[0]
if len(item.ProviderKeys) != 1 || item.ProviderKeys[0] != "zhipu-openai" {
t.Fatalf("expected model provider zhipu-openai only, got %+v", item.ProviderKeys)
}
if len(item.Providers) != 1 || item.Providers[0].Name != "智谱AI" || item.Providers[0].SourceCount != 2 {
t.Fatalf("expected provider summary to aggregate both sources under model provider, got %+v", item.Providers)
}
if !hasFilterCount(response.Filters.Providers, "zhipu-openai", 1) {
t.Fatalf("expected zhipu provider filter count 1, got %+v", response.Filters.Providers)
}
if hasFilterCount(response.Filters.Providers, "volces-openai", 1) {
t.Fatalf("did not expect platform provider in model provider filters: %+v", response.Filters.Providers)
}
}
func TestBillingConfigLinesShowsTextInputAndOutputPricing(t *testing.T) {
lines := billingConfigLines(map[string]any{
"text_total": map[string]any{
"basePrice": 0.01,
"formulaConfig": map[string]any{
"inputTokenPrice": 0.01,
"outputTokenPrice": 0.03,
},
},
})
if len(lines) != 2 {
t.Fatalf("expected input and output pricing lines, got %+v", lines)
}
if lines[0] != "输入 0.01/k tokens" {
t.Fatalf("unexpected input pricing line %q", lines[0])
}
if lines[1] != "输出 0.03/k tokens" {
t.Fatalf("unexpected output pricing line %q", lines[1])
}
}
func TestBillingConfigLinesShowsVideoFiveSecondBasis(t *testing.T) {
lines := billingConfigLines(map[string]any{
"video": map[string]any{
"basePrice": float64(75),
"dynamicWeight": map[string]any{"480p": float64(1), "720p": float64(2)},
},
})
if len(lines) != 1 {
t.Fatalf("expected one video pricing line, got %+v", lines)
}
if lines[0] != "视频:480p 75 / 720p 1505秒基准)" {
t.Fatalf("unexpected video pricing line %q", lines[0])
}
flatLines := billingConfigLines(map[string]any{"videoBase": float64(100)})
if len(flatLines) != 1 || flatLines[0] != "视频:100 / 5秒基准" {
t.Fatalf("unexpected flat video pricing line %+v", flatLines)
}
}
func hasFilterCount(options []ModelCatalogFilterOption, value string, count int) bool {
for _, option := range options {
if option.Value == value && option.Count == count {
return true
}
}
return false
}
@@ -0,0 +1,48 @@
package httpapi
import (
"context"
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)
func (s *Server) platformModelResponse(ctx context.Context, model store.PlatformModel) store.PlatformModel {
model = s.withEffectiveResponseBillingConfig(ctx, model)
return store.FilterPlatformModelBillingConfig(model)
}
func (s *Server) platformModelResponses(ctx context.Context, models []store.PlatformModel) []store.PlatformModel {
items := make([]store.PlatformModel, len(models))
for i, model := range models {
items[i] = s.platformModelResponse(ctx, model)
}
return items
}
func (s *Server) withEffectiveResponseBillingConfig(ctx context.Context, model store.PlatformModel) store.PlatformModel {
config := model.BillingConfig
if model.PricingRuleSetID != "" {
if ruleSetConfig, err := s.store.PricingRuleSetBillingConfig(ctx, model.PricingRuleSetID); err == nil && len(ruleSetConfig) > 0 {
config = ruleSetConfig
}
}
if len(model.BillingConfigOverride) > 0 {
config = mergeResponseBillingConfig(config, model.BillingConfigOverride)
}
model.BillingConfig = config
return model
}
func mergeResponseBillingConfig(base map[string]any, override map[string]any) map[string]any {
if len(base) == 0 && len(override) == 0 {
return nil
}
out := make(map[string]any, len(base)+len(override))
for key, value := range base {
out[key] = value
}
for key, value := range override {
out[key] = value
}
return out
}
+1
View File
@@ -92,6 +92,7 @@ func NewServer(cfg config.Config, db *store.Store, logger *slog.Logger) http.Han
mux.Handle("POST /api/admin/platform-models", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.createPlatformModel)))
mux.Handle("DELETE /api/admin/platform-models/{modelID}", server.requireAdmin(auth.PermissionManager, http.HandlerFunc(server.deletePlatformModel)))
mux.Handle("GET /api/admin/models", server.requireAdmin(auth.PermissionPower, http.HandlerFunc(server.listModels)))
mux.Handle("GET /api/v1/model-catalog", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listModelCatalog)))
mux.Handle("GET /api/v1/platforms", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listPlayablePlatforms)))
mux.Handle("GET /api/v1/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listPlayableModels)))
mux.Handle("GET /api/v1/playground/models", server.auth.Require(auth.PermissionBasic, http.HandlerFunc(server.listPlayableModels)))
+13 -7
View File
@@ -17,7 +17,7 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
p.retry_policy, p.rate_limit_policy,
COALESCE(p.dynamic_priority, p.priority) AS effective_priority,
m.id::text, COALESCE(m.base_model_id::text, ''), COALESCE(b.canonical_model_key, ''),
COALESCE(b.provider_model_name, ''), m.model_name, COALESCE(m.model_alias, ''),
COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), m.model_name, COALESCE(m.model_alias, ''),
$2 AS requested_model_type, m.display_name, m.capabilities, m.capability_override,
COALESCE(b.base_billing_config, '{}'::jsonb), m.billing_config, m.billing_config_override,
m.pricing_mode, COALESCE(m.discount_factor, 0)::float8, COALESCE(m.pricing_rule_set_id::text, ''),
@@ -32,17 +32,22 @@ LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provi
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
LEFT JOIN runtime_client_states s
ON s.client_id = p.platform_key || ':' || $2 || ':' || m.model_name
ON s.client_id = p.platform_key || ':' || $2 || ':' || COALESCE(NULLIF(m.provider_model_name, ''), m.model_name)
WHERE p.status = 'enabled'
AND p.deleted_at IS NULL
AND m.enabled = true
AND m.model_type @> jsonb_build_array($2)
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
AND (
m.model_name = $1
OR m.model_alias = $1
OR b.canonical_model_key = $1
OR b.provider_model_name = $1
(COALESCE(m.model_alias, '') <> '' AND m.model_alias = $1)
OR (
COALESCE(m.model_alias, '') = ''
AND (
m.model_name = $1
OR b.canonical_model_key = $1
OR b.provider_model_name = $1
)
)
)
ORDER BY effective_priority ASC,
COALESCE(s.limiter_ratio, 0) ASC,
@@ -137,7 +142,8 @@ ORDER BY effective_priority ASC,
item.RuntimeRateLimitPolicy = decodeObject(runtimeRateLimitPolicy)
item.AutoDisablePolicy = decodeObject(autoDisablePolicy)
item.DegradePolicy = decodeObject(degradePolicy)
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, item.ModelName)
upstreamModelName := firstNonEmpty(item.ProviderModelName, item.ModelName)
item.ClientID = fmt.Sprintf("%s:%s:%s", item.PlatformKey, item.ModelType, upstreamModelName)
item.QueueKey = item.ClientID
items = append(items, item)
}
@@ -0,0 +1,147 @@
package store
import "strings"
func FilterPlatformModelBillingConfigs(models []PlatformModel) []PlatformModel {
filtered := make([]PlatformModel, len(models))
for i, model := range models {
filtered[i] = FilterPlatformModelBillingConfig(model)
}
return filtered
}
func FilterPlatformModelBillingConfig(model PlatformModel) PlatformModel {
model.BillingConfig = filterBillingConfigByModelTypes(model.BillingConfig, model.ModelType)
model.BillingConfigOverride = filterBillingConfigByModelTypes(model.BillingConfigOverride, model.ModelType)
return model
}
func filterBillingConfigByModelTypes(config map[string]any, modelTypes []string) map[string]any {
if len(config) == 0 {
return nil
}
resources := billingResourcesForModelTypes(modelTypes)
if len(resources) == 0 {
return cloneBillingConfig(config)
}
filtered := map[string]any{}
for key, value := range config {
if billingConfigKeyAllowed(key, resources) {
filtered[key] = value
}
}
if resources["image_edit"] && !hasAnyKey(filtered, "image_edit", "editBase") && !hasAnyKey(config, "image_edit", "editBase") {
copyBillingConfigKey(filtered, config, "image")
copyBillingConfigKey(filtered, config, "imageBase")
}
if len(filtered) == 0 {
return nil
}
return filtered
}
func billingResourcesForModelTypes(modelTypes []string) map[string]bool {
resources := map[string]bool{}
for _, modelType := range modelTypes {
switch normalizeBillingType(modelType) {
case "chat", "text", "responses", "text_generate", "text_embedding", "embedding",
"image_analysis", "video_understanding", "audio_understanding", "omni", "tools_call":
resources["text"] = true
case "image", "images.generations", "image_generate":
resources["image"] = true
case "images.edits", "image_edit":
resources["image_edit"] = true
case "video", "videos.generations", "video_generate", "image_to_video", "text_to_video",
"video_edit", "omni_video", "video_reference", "video_first_last_frame":
resources["video"] = true
case "audio", "audio_generate", "text_to_speech", "speech":
resources["audio"] = true
case "music", "music_generate":
resources["music"] = true
case "digital_human", "digital_human_generate":
resources["digital_human"] = true
case "model", "model_3d", "text_to_model", "image_to_model", "multiview_to_model", "mesh_edit":
resources["model"] = true
default:
inferBillingResources(modelType, resources)
}
}
return resources
}
func inferBillingResources(modelType string, resources map[string]bool) {
normalized := normalizeBillingType(modelType)
switch {
case strings.Contains(normalized, "digital_human"):
resources["digital_human"] = true
case strings.Contains(normalized, "video"):
resources["video"] = true
case strings.Contains(normalized, "image"):
resources["image"] = true
case strings.Contains(normalized, "audio") || strings.Contains(normalized, "speech"):
resources["audio"] = true
case strings.Contains(normalized, "music"):
resources["music"] = true
case strings.Contains(normalized, "model") || strings.Contains(normalized, "mesh"):
resources["model"] = true
case strings.Contains(normalized, "text") || strings.Contains(normalized, "token"):
resources["text"] = true
}
}
func billingConfigKeyAllowed(key string, resources map[string]bool) bool {
switch normalizeBillingConfigKey(key) {
case "text", "texttotal", "text_total", "textinputper1k", "textoutputper1k", "textinput", "textoutput", "text_input", "text_output", "inputtokenprice", "outputtokenprice":
return resources["text"]
case "image", "imagebase":
return resources["image"]
case "image_edit", "imageedit", "editbase":
return resources["image_edit"]
case "video", "videobase":
return resources["video"]
case "audio", "audiobase":
return resources["audio"]
case "music", "musicbase":
return resources["music"]
case "digital_human", "digitalhuman", "digitalhumanbase":
return resources["digital_human"]
case "model", "modelbase":
return resources["model"]
default:
return false
}
}
func normalizeBillingType(value string) string {
return strings.ToLower(strings.TrimSpace(value))
}
func normalizeBillingConfigKey(value string) string {
return strings.ReplaceAll(strings.ReplaceAll(normalizeBillingType(value), "-", "_"), " ", "")
}
func cloneBillingConfig(config map[string]any) map[string]any {
clone := make(map[string]any, len(config))
for key, value := range config {
clone[key] = value
}
return clone
}
func hasAnyKey(config map[string]any, keys ...string) bool {
for _, key := range keys {
if _, ok := config[key]; ok {
return true
}
}
return false
}
func copyBillingConfigKey(target map[string]any, source map[string]any, key string) {
if value, ok := source[key]; ok {
target[key] = value
}
}
@@ -0,0 +1,99 @@
package store
import "testing"
func TestFilterPlatformModelBillingConfigKeepsOnlyImagePricing(t *testing.T) {
model := PlatformModel{
ModelType: StringList{"image_generate"},
BillingConfig: map[string]any{
"text": map[string]any{"basePrice": 0.01},
"image": map[string]any{"basePrice": 10},
"image_edit": map[string]any{"basePrice": 12},
"video": map[string]any{"basePrice": 100},
},
BillingConfigOverride: map[string]any{
"image": map[string]any{"basePrice": 8},
"video": map[string]any{"basePrice": 80},
},
}
filtered := FilterPlatformModelBillingConfig(model)
assertHasKeys(t, filtered.BillingConfig, "image")
assertMissingKeys(t, filtered.BillingConfig, "text", "image_edit", "video")
assertHasKeys(t, filtered.BillingConfigOverride, "image")
assertMissingKeys(t, filtered.BillingConfigOverride, "video")
}
func TestFilterPlatformModelBillingConfigKeepsImageEditOrImageFallback(t *testing.T) {
model := PlatformModel{
ModelType: StringList{"image_edit"},
BillingConfig: map[string]any{
"image": map[string]any{"basePrice": 10},
"video": map[string]any{"basePrice": 100},
},
}
filtered := FilterPlatformModelBillingConfig(model)
assertHasKeys(t, filtered.BillingConfig, "image")
assertMissingKeys(t, filtered.BillingConfig, "video")
}
func TestFilterPlatformModelBillingConfigKeepsVideoPricing(t *testing.T) {
model := PlatformModel{
ModelType: StringList{"video_generate", "image_to_video"},
BillingConfig: map[string]any{
"image": map[string]any{"basePrice": 10},
"video": map[string]any{"basePrice": 100},
"videoBase": 100,
},
}
filtered := FilterPlatformModelBillingConfig(model)
assertHasKeys(t, filtered.BillingConfig, "video", "videoBase")
assertMissingKeys(t, filtered.BillingConfig, "image")
}
func TestFilterPlatformModelBillingConfigKeepsTextFlatPricing(t *testing.T) {
model := PlatformModel{
ModelType: StringList{"text_generate"},
BillingConfig: map[string]any{
"textInputPer1k": 0.01,
"textOutputPer1k": 0.02,
"text_total": map[string]any{
"basePrice": 0.01,
"formulaConfig": map[string]any{"inputTokenPrice": 0.01, "outputTokenPrice": 0.02},
"dynamicWeight": map[string]any{"cached": 0.5},
"dimensionSchema": map[string]any{"unit": "1k_tokens"},
},
"inputTokenPrice": 0.01,
"outputTokenPrice": 0.02,
"image": map[string]any{"basePrice": 10},
},
}
filtered := FilterPlatformModelBillingConfig(model)
assertHasKeys(t, filtered.BillingConfig, "textInputPer1k", "textOutputPer1k", "text_total", "inputTokenPrice", "outputTokenPrice")
assertMissingKeys(t, filtered.BillingConfig, "image")
}
func assertHasKeys(t *testing.T, value map[string]any, keys ...string) {
t.Helper()
for _, key := range keys {
if _, ok := value[key]; !ok {
t.Fatalf("expected key %q in %#v", key, value)
}
}
}
func assertMissingKeys(t *testing.T, value map[string]any, keys ...string) {
t.Helper()
for _, key := range keys {
if _, ok := value[key]; ok {
t.Fatalf("expected key %q to be filtered from %#v", key, value)
}
}
}
+14 -6
View File
@@ -79,6 +79,8 @@ WHERE resource_type = 'platform_model'
}
func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier, input CreatePlatformModelInput) (PlatformModel, error) {
input.ModelName = strings.TrimSpace(input.ModelName)
input.ProviderModelName = strings.TrimSpace(input.ProviderModelName)
base, err := s.lookupBaseModel(ctx, q, input.BaseModelID, input.CanonicalModelKey, input.ModelName)
if err != nil && !IsNotFound(err) {
return PlatformModel{}, err
@@ -93,6 +95,9 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
if input.ModelName == "" {
input.ModelName = base.ProviderModelName
}
if input.ProviderModelName == "" {
input.ProviderModelName = input.ModelName
}
if input.DisplayName == "" {
input.DisplayName = firstNonEmpty(base.DisplayName, input.ModelName)
}
@@ -153,19 +158,20 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
var modelTypeBytes []byte
err = q.QueryRow(ctx, `
INSERT INTO platform_models (
platform_id, base_model_id, model_name, model_alias, model_type, display_name,
platform_id, base_model_id, model_name, provider_model_name, model_alias, model_type, display_name,
capability_override, capabilities, pricing_mode, discount_factor,
pricing_rule_set_id, billing_config_override, billing_config, permission_config, retry_policy, rate_limit_policy,
runtime_policy_set_id, runtime_policy_override, enabled
)
VALUES (
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5::jsonb, $6,
$7::jsonb, $8::jsonb, $9, $10::numeric,
NULLIF($11, '')::uuid, $12::jsonb, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb,
NULLIF($17, '')::uuid, $18::jsonb, true
$1::uuid, $2::uuid, $3, NULLIF($4, ''), NULLIF($5, ''), $6::jsonb, $7,
$8::jsonb, $9::jsonb, $10, $11::numeric,
NULLIF($12, '')::uuid, $13::jsonb, $14::jsonb, $15::jsonb, $16::jsonb, $17::jsonb,
NULLIF($18, '')::uuid, $19::jsonb, true
)
ON CONFLICT (platform_id, model_name) DO UPDATE
SET base_model_id = EXCLUDED.base_model_id,
provider_model_name = EXCLUDED.provider_model_name,
model_alias = EXCLUDED.model_alias,
display_name = EXCLUDED.display_name,
capability_override = EXCLUDED.capability_override,
@@ -183,7 +189,7 @@ SET base_model_id = EXCLUDED.base_model_id,
enabled = true,
updated_at = now()
RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_name,
COALESCE(model_alias, ''), model_type, display_name, capability_override,
COALESCE(NULLIF(provider_model_name, ''), model_name), COALESCE(model_alias, ''), model_type, display_name, capability_override,
capabilities, pricing_mode, COALESCE(discount_factor, 0)::float8,
COALESCE(pricing_rule_set_id::text, ''), billing_config_override, billing_config,
permission_config, retry_policy, rate_limit_policy, COALESCE(runtime_policy_set_id::text, ''), runtime_policy_override,
@@ -191,6 +197,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
input.PlatformID,
baseID,
input.ModelName,
input.ProviderModelName,
input.ModelAlias,
string(modelTypeJSON),
input.DisplayName,
@@ -211,6 +218,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
&model.PlatformID,
&model.BaseModelID,
&model.ModelName,
&model.ProviderModelName,
&model.ModelAlias,
&modelTypeBytes,
&model.DisplayName,
+3 -1
View File
@@ -133,6 +133,7 @@ type PlatformModel struct {
Provider string `json:"provider,omitempty"`
PlatformName string `json:"platformName,omitempty"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
@@ -691,7 +692,7 @@ func (s *Store) listModels(ctx context.Context, platformID string) ([]PlatformMo
rows, err := s.pool.Query(ctx, `
SELECT m.id::text, m.platform_id::text, COALESCE(m.base_model_id::text, ''), p.provider, p.name,
m.model_name, COALESCE(m.model_alias, ''), m.model_type, m.display_name,
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''), m.model_type, m.display_name,
m.capability_override, m.capabilities, m.pricing_mode, COALESCE(m.discount_factor, 0)::float8,
COALESCE(m.pricing_rule_set_id::text, ''), m.billing_config_override, m.billing_config,
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, ''), m.runtime_policy_override,
@@ -724,6 +725,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
&model.Provider,
&model.PlatformName,
&model.ModelName,
&model.ProviderModelName,
&model.ModelAlias,
&modelTypeBytes,
&model.DisplayName,
+58 -2
View File
@@ -3,6 +3,7 @@ package store
import (
"context"
"encoding/json"
"strconv"
"strings"
"github.com/jackc/pgx/v5"
@@ -182,7 +183,7 @@ func (s *Store) PricingRuleSetBillingConfig(ctx context.Context, id string) (map
return nil, nil
}
rows, err := s.pool.Query(ctx, `
SELECT resource_type, base_price::float8, dynamic_weight
SELECT resource_type, base_price::float8, dynamic_weight, formula_config
FROM model_pricing_rules
WHERE rule_set_id = $1::uuid
AND status = 'active'
@@ -197,15 +198,31 @@ ORDER BY priority ASC, resource_type ASC`, id)
var resourceType string
var basePrice float64
var dynamicWeightBytes []byte
if err := rows.Scan(&resourceType, &basePrice, &dynamicWeightBytes); err != nil {
var formulaConfigBytes []byte
if err := rows.Scan(&resourceType, &basePrice, &dynamicWeightBytes, &formulaConfigBytes); err != nil {
return nil, err
}
dynamicWeight := decodeObject(dynamicWeightBytes)
formulaConfig := decodeObject(formulaConfigBytes)
switch resourceType {
case "text_input":
config["textInputPer1k"] = basePrice
case "text_output":
config["textOutputPer1k"] = basePrice
case "text_total":
inputPrice := basePrice
if value, ok := pricingRuleNumberFromKeys(formulaConfig, "inputTokenPrice", "input_token_price", "textInputPer1k", "text_input"); ok {
inputPrice = value
}
config["textInputPer1k"] = inputPrice
if outputPrice, ok := pricingRuleNumberFromKeys(formulaConfig, "outputTokenPrice", "output_token_price", "textOutputPer1k", "text_output"); ok {
config["textOutputPer1k"] = outputPrice
}
resourceConfig := pricingResourceConfig(basePrice, dynamicWeight)
if len(formulaConfig) > 0 {
resourceConfig["formulaConfig"] = formulaConfig
}
config["text_total"] = resourceConfig
case "image":
config["imageBase"] = basePrice
config["image"] = pricingResourceConfig(basePrice, dynamicWeight)
@@ -225,6 +242,45 @@ ORDER BY priority ASC, resource_type ASC`, id)
return config, nil
}
func pricingRuleNumberFromKeys(config map[string]any, keys ...string) (float64, bool) {
if len(config) == 0 {
return 0, false
}
for _, key := range keys {
if value, ok := pricingRuleNumberValue(config[key]); ok {
return value, true
}
}
return 0, false
}
func pricingRuleNumberValue(value any) (float64, bool) {
switch typed := value.(type) {
case float64:
return typed, true
case float32:
return float64(typed), true
case int:
return float64(typed), true
case int64:
return float64(typed), true
case int32:
return float64(typed), true
case json.Number:
number, err := typed.Float64()
return number, err == nil
case string:
trimmed := strings.TrimSpace(typed)
if trimmed == "" {
return 0, false
}
number, err := strconv.ParseFloat(trimmed, 64)
return number, err == nil
default:
return 0, false
}
}
func pricingResourceConfig(basePrice float64, dynamicWeight map[string]any) map[string]any {
config := map[string]any{"basePrice": basePrice}
if len(dynamicWeight) > 0 {
+1
View File
@@ -15,6 +15,7 @@ type CreatePlatformModelInput struct {
BaseModelID string `json:"baseModelId"`
CanonicalModelKey string `json:"canonicalModelKey"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName"`
ModelAlias string `json:"modelAlias"`
ModelType StringList `json:"modelType"`
DisplayName string `json:"displayName"`
+96 -13
View File
@@ -11,13 +11,35 @@ import (
"github.com/jackc/pgx/v5/pgconn"
)
func (s *Store) ListTasks(ctx context.Context, user *auth.User, limit int) ([]GatewayTask, error) {
if limit <= 0 {
limit = 50
type TaskListFilter struct {
Query string
ModelType string
CreatedFrom *time.Time
CreatedTo *time.Time
Page int
PageSize int
}
type TaskListResult struct {
Items []GatewayTask
Total int
Page int
PageSize int
}
func (s *Store) ListTasks(ctx context.Context, user *auth.User, filter TaskListFilter) (TaskListResult, error) {
page := filter.Page
if page <= 0 {
page = 1
}
if limit > 100 {
limit = 100
pageSize := filter.PageSize
if pageSize <= 0 {
pageSize = 50
}
if pageSize > 100 {
pageSize = 100
}
offset := (page - 1) * pageSize
gatewayUserID := localGatewayUserID(user)
apiKeyID := ""
userID := ""
@@ -26,11 +48,22 @@ func (s *Store) ListTasks(ctx context.Context, user *auth.User, limit int) ([]Ga
userID = strings.TrimSpace(user.ID)
}
if gatewayUserID == "" && userID == "" {
return nil, ErrLocalUserRequired
return TaskListResult{}, ErrLocalUserRequired
}
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
queryPattern := ""
if query := strings.TrimSpace(filter.Query); query != "" {
queryPattern = "%" + query + "%"
}
args := []any{
gatewayUserID,
userID,
apiKeyID,
queryPattern,
strings.TrimSpace(filter.ModelType),
nullableTaskListTime(filter.CreatedFrom),
nullableTaskListTime(filter.CreatedTo),
}
whereSQL := `
WHERE (
(
NULLIF($1, '')::uuid IS NOT NULL
@@ -46,10 +79,45 @@ WHERE (
NULLIF($3, '') IS NULL
OR api_key_id = $3
)
AND (
NULLIF($4, '') IS NULL
OR id::text ILIKE $4
OR COALESCE(request_id, '') ILIKE $4
OR kind ILIKE $4
OR model ILIKE $4
OR COALESCE(requested_model, '') ILIKE $4
OR COALESCE(resolved_model, '') ILIKE $4
OR COALESCE(api_key_id, '') ILIKE $4
OR COALESCE(api_key_name, '') ILIKE $4
OR COALESCE(api_key_prefix, '') ILIKE $4
OR status ILIKE $4
OR COALESCE(model_type, '') ILIKE $4
)
AND (
NULLIF($5, '') IS NULL
OR model_type = $5
)
AND (
$6::timestamptz IS NULL
OR created_at >= $6::timestamptz
)
AND (
$7::timestamptz IS NULL
OR created_at <= $7::timestamptz
)`
var total int
if err := s.pool.QueryRow(ctx, `SELECT count(*) FROM gateway_tasks `+whereSQL, args...).Scan(&total); err != nil {
return TaskListResult{}, err
}
queryArgs := append(args, pageSize, offset)
rows, err := s.pool.Query(ctx, `
SELECT `+gatewayTaskColumns+`
FROM gateway_tasks
`+whereSQL+`
ORDER BY created_at DESC
LIMIT $4`, gatewayUserID, userID, apiKeyID, limit)
LIMIT $8 OFFSET $9`, queryArgs...)
if err != nil {
return nil, err
return TaskListResult{}, err
}
defer rows.Close()
@@ -57,11 +125,26 @@ LIMIT $4`, gatewayUserID, userID, apiKeyID, limit)
for rows.Next() {
task, err := scanGatewayTask(rows)
if err != nil {
return nil, err
return TaskListResult{}, err
}
items = append(items, task)
}
return items, rows.Err()
if err := rows.Err(); err != nil {
return TaskListResult{}, err
}
return TaskListResult{
Items: items,
Total: total,
Page: page,
PageSize: pageSize,
}, nil
}
func nullableTaskListTime(value *time.Time) any {
if value == nil {
return nil
}
return *value
}
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
@@ -0,0 +1,6 @@
ALTER TABLE IF EXISTS platform_models
ADD COLUMN IF NOT EXISTS provider_model_name text;
UPDATE platform_models
SET provider_model_name = model_name
WHERE provider_model_name IS NULL OR btrim(provider_model_name) = '';