feat: improve model catalog aggregation
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
|
||||
@@ -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 != "" {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 150(5秒基准)" {
|
||||
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
|
||||
}
|
||||
@@ -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)))
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
|
||||
@@ -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) = '';
|
||||
Reference in New Issue
Block a user