fix gateway loopback validation chains
This commit is contained in:
@@ -59,7 +59,7 @@ func (s *Store) ListBaseModels(ctx context.Context) ([]BaseModel, error) {
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT `+baseModelColumns+`
|
||||
FROM base_model_catalog
|
||||
ORDER BY provider_key ASC, model_type ASC, canonical_model_key ASC`)
|
||||
ORDER BY provider_key ASC, canonical_model_key ASC`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -84,7 +84,7 @@ func (s *Store) CreateBaseModel(ctx context.Context, input BaseModelInput) (Base
|
||||
runtimePolicyOverride, _ := json.Marshal(emptyObjectIfNil(input.RuntimePolicyOverride))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
|
||||
modelType := primaryString(input.ModelType, "text_generate")
|
||||
modelType, _ := json.Marshal(input.ModelType)
|
||||
|
||||
return scanBaseModel(s.pool.QueryRow(ctx, `
|
||||
INSERT INTO base_model_catalog (
|
||||
@@ -94,7 +94,7 @@ INSERT INTO base_model_catalog (
|
||||
)
|
||||
VALUES (
|
||||
(SELECT id FROM model_catalog_providers WHERE provider_key = $1 OR provider_code = $1 LIMIT 1),
|
||||
$1, $2, $3, $4, $5, $6, $7, $8,
|
||||
$1, $2, $3, $4::jsonb, $5, $6, $7, $8,
|
||||
COALESCE(NULLIF($9, '')::uuid, (SELECT id FROM model_pricing_rule_sets WHERE rule_set_key = 'default-multimodal-v1' LIMIT 1)),
|
||||
COALESCE(NULLIF($10, '')::uuid, (SELECT id FROM model_runtime_policy_sets WHERE policy_key = 'default-runtime-v1' LIMIT 1)),
|
||||
$11, $12, NULLIF($13, ''), NULLIF($14::jsonb, '{}'::jsonb), $15, $16
|
||||
@@ -103,7 +103,7 @@ RETURNING `+baseModelColumns,
|
||||
input.ProviderKey,
|
||||
input.CanonicalModelKey,
|
||||
input.ProviderModelName,
|
||||
modelType,
|
||||
string(modelType),
|
||||
input.ModelAlias,
|
||||
capabilities,
|
||||
billingConfig,
|
||||
@@ -127,7 +127,7 @@ func (s *Store) UpdateBaseModel(ctx context.Context, id string, input BaseModelI
|
||||
runtimePolicyOverride, _ := json.Marshal(emptyObjectIfNil(input.RuntimePolicyOverride))
|
||||
metadata, _ := json.Marshal(emptyObjectIfNil(input.Metadata))
|
||||
defaultSnapshot, _ := json.Marshal(emptyObjectIfNil(input.DefaultSnapshot))
|
||||
modelType := primaryString(input.ModelType, "text_generate")
|
||||
modelType, _ := json.Marshal(input.ModelType)
|
||||
|
||||
return scanBaseModel(s.pool.QueryRow(ctx, `
|
||||
UPDATE base_model_catalog
|
||||
@@ -135,7 +135,7 @@ SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = $
|
||||
provider_key = $2,
|
||||
canonical_model_key = $3,
|
||||
provider_model_name = $4,
|
||||
model_type = $5,
|
||||
model_type = $5::jsonb,
|
||||
display_name = $6,
|
||||
capabilities = $7,
|
||||
base_billing_config = $8,
|
||||
@@ -156,7 +156,7 @@ RETURNING `+baseModelColumns,
|
||||
input.ProviderKey,
|
||||
input.CanonicalModelKey,
|
||||
input.ProviderModelName,
|
||||
modelType,
|
||||
string(modelType),
|
||||
input.ModelAlias,
|
||||
capabilities,
|
||||
billingConfig,
|
||||
@@ -195,7 +195,7 @@ SET provider_id = (SELECT id FROM model_catalog_providers WHERE provider_key = C
|
||||
provider_key = COALESCE($2::text, provider_key),
|
||||
canonical_model_key = COALESCE($3::text, canonical_model_key),
|
||||
provider_model_name = COALESCE($4::text, provider_model_name),
|
||||
model_type = COALESCE($5::text, model_type),
|
||||
model_type = COALESCE($5::jsonb, model_type),
|
||||
display_name = COALESCE($6::text, display_name),
|
||||
capabilities = COALESCE($7::jsonb, capabilities),
|
||||
base_billing_config = COALESCE($8::jsonb, base_billing_config),
|
||||
@@ -214,7 +214,7 @@ RETURNING `+baseModelColumns,
|
||||
stringFromSnapshot(snapshot, "providerKey"),
|
||||
stringFromSnapshot(snapshot, "canonicalModelKey"),
|
||||
stringFromSnapshot(snapshot, "providerModelName"),
|
||||
stringFromSnapshot(snapshot, "modelType"),
|
||||
jsonStringListFromSnapshot(snapshot, "modelType"),
|
||||
stringFromSnapshot(snapshot, "modelAlias", "displayName"),
|
||||
jsonFromSnapshot(snapshot, "capabilities"),
|
||||
jsonFromSnapshot(snapshot, "baseBillingConfig"),
|
||||
@@ -242,9 +242,10 @@ SET provider_id = (
|
||||
canonical_model_key = COALESCE(NULLIF(default_snapshot->>'canonicalModelKey', ''), canonical_model_key),
|
||||
provider_model_name = COALESCE(NULLIF(default_snapshot->>'providerModelName', ''), provider_model_name),
|
||||
model_type = COALESCE(NULLIF(CASE
|
||||
WHEN jsonb_typeof(default_snapshot->'modelType') = 'array' THEN default_snapshot->'modelType'->>0
|
||||
ELSE default_snapshot->>'modelType'
|
||||
END, ''), model_type),
|
||||
WHEN jsonb_typeof(default_snapshot->'modelType') = 'array' THEN default_snapshot->'modelType'
|
||||
WHEN COALESCE(default_snapshot->>'modelType', '') <> '' THEN jsonb_build_array(default_snapshot->>'modelType')
|
||||
ELSE NULL
|
||||
END, '[]'::jsonb), model_type),
|
||||
display_name = COALESCE(NULLIF(COALESCE(default_snapshot->>'modelAlias', default_snapshot->>'displayName'), ''), display_name),
|
||||
capabilities = COALESCE(default_snapshot->'capabilities', capabilities),
|
||||
base_billing_config = COALESCE(default_snapshot->'baseBillingConfig', base_billing_config),
|
||||
@@ -292,7 +293,7 @@ func scanBaseModelRows(rows pgx.Rows) ([]BaseModel, error) {
|
||||
|
||||
func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
|
||||
var item BaseModel
|
||||
var modelType string
|
||||
var modelType []byte
|
||||
var modelAlias string
|
||||
var capabilities []byte
|
||||
var billingConfig []byte
|
||||
@@ -330,7 +331,7 @@ func scanBaseModel(scanner baseModelScanner) (BaseModel, error) {
|
||||
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
|
||||
item.Metadata = decodeObject(metadata)
|
||||
item.DefaultSnapshot = decodeObject(defaultSnapshot)
|
||||
item.ModelType = baseModelTypes(item.Capabilities, item.Metadata, modelType)
|
||||
item.ModelType = decodeStringArray(modelType)
|
||||
item.ModelAlias = modelAlias
|
||||
item.DisplayName = modelAlias
|
||||
return item, nil
|
||||
@@ -420,14 +421,22 @@ func jsonFromSnapshot(snapshot map[string]any, key string) any {
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func baseModelTypes(capabilities map[string]any, metadata map[string]any, fallback string) StringList {
|
||||
values := make([]string, 0)
|
||||
values = append(values, stringListFromAny(capabilities["originalTypes"])...)
|
||||
values = append(values, stringListFromAny(metadata["originalTypes"])...)
|
||||
if fallback != "" {
|
||||
values = append(values, fallback)
|
||||
func jsonStringListFromSnapshot(snapshot map[string]any, key string) any {
|
||||
values := stringListFromAny(snapshot[key])
|
||||
if len(values) == 0 {
|
||||
if value, ok := snapshot[key].(string); ok {
|
||||
values = []string{value}
|
||||
}
|
||||
}
|
||||
return uniqueStringList(values)
|
||||
normalized := uniqueStringList(values)
|
||||
if len(normalized) == 0 {
|
||||
return nil
|
||||
}
|
||||
raw, err := json.Marshal(normalized)
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
return string(raw)
|
||||
}
|
||||
|
||||
func stringListFromAny(value any) []string {
|
||||
@@ -451,16 +460,39 @@ func uniqueStringList(values []string) StringList {
|
||||
out := make([]string, 0, len(values))
|
||||
seen := map[string]bool{}
|
||||
for _, value := range values {
|
||||
value = strings.TrimSpace(value)
|
||||
if value == "" || seen[value] {
|
||||
continue
|
||||
for _, normalized := range modelTypeAliases(value) {
|
||||
normalized = strings.TrimSpace(normalized)
|
||||
if normalized == "" || seen[normalized] {
|
||||
continue
|
||||
}
|
||||
seen[normalized] = true
|
||||
out = append(out, normalized)
|
||||
}
|
||||
seen[value] = true
|
||||
out = append(out, value)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func normalizeModelTypeList(values []string) StringList {
|
||||
return uniqueStringList(values)
|
||||
}
|
||||
|
||||
func modelTypeAliases(value string) []string {
|
||||
switch strings.TrimSpace(value) {
|
||||
case "chat", "text", "responses":
|
||||
return []string{"text_generate"}
|
||||
case "image":
|
||||
return []string{"image_generate", "image_edit"}
|
||||
case "images.generations":
|
||||
return []string{"image_generate"}
|
||||
case "images.edits":
|
||||
return []string{"image_edit"}
|
||||
case "video", "videos.generations":
|
||||
return []string{"video_generate"}
|
||||
default:
|
||||
return []string{value}
|
||||
}
|
||||
}
|
||||
|
||||
func primaryString(values []string, fallback string) string {
|
||||
for _, value := range values {
|
||||
if value = strings.TrimSpace(value); value != "" {
|
||||
|
||||
@@ -18,27 +18,25 @@ SELECT p.id::text, p.platform_key, p.name, p.provider,
|
||||
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, ''),
|
||||
m.model_type, m.display_name, m.capabilities, m.capability_override,
|
||||
$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, ''),
|
||||
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
|
||||
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb)
|
||||
COALESCE(b.pricing_rule_set_id::text, ''),
|
||||
m.permission_config, m.retry_policy, m.rate_limit_policy, COALESCE(m.runtime_policy_set_id::text, COALESCE(b.runtime_policy_set_id::text, '')),
|
||||
COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb),
|
||||
COALESCE(rp.retry_policy, '{}'::jsonb), COALESCE(rp.rate_limit_policy, '{}'::jsonb),
|
||||
COALESCE(rp.auto_disable_policy, '{}'::jsonb), COALESCE(rp.degrade_policy, '{}'::jsonb)
|
||||
FROM platform_models m
|
||||
JOIN integration_platforms p ON p.id = m.platform_id
|
||||
LEFT JOIN model_catalog_providers cp ON cp.provider_key = p.provider OR cp.provider_code = p.provider
|
||||
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 || ':' || m.model_type || ':' || m.model_name
|
||||
ON s.client_id = p.platform_key || ':' || $2 || ':' || m.model_name
|
||||
WHERE p.status = 'enabled'
|
||||
AND p.deleted_at IS NULL
|
||||
AND m.enabled = true
|
||||
AND (
|
||||
m.model_type = $2
|
||||
OR ($2 = 'text_generate' AND m.model_type IN ('chat', 'responses', 'text'))
|
||||
OR ($2 = 'image_generate' AND m.model_type IN ('image', 'images.generations'))
|
||||
OR ($2 = 'image_edit' AND m.model_type IN ('images.edits'))
|
||||
OR ($2 = 'video_generate' AND m.model_type IN ('video', 'videos.generations', 'video_generate', 'text_to_video', 'image_to_video', 'omni_video', 'video_edit', 'video_reference', 'video_first_last_frame'))
|
||||
)
|
||||
AND m.model_type @> jsonb_build_array($2)
|
||||
AND (p.cooldown_until IS NULL OR p.cooldown_until <= now())
|
||||
AND (
|
||||
m.model_name = $1
|
||||
@@ -73,6 +71,10 @@ ORDER BY effective_priority ASC,
|
||||
var modelRetryPolicy []byte
|
||||
var modelRateLimitPolicy []byte
|
||||
var runtimePolicyOverride []byte
|
||||
var runtimeRetryPolicy []byte
|
||||
var runtimeRateLimitPolicy []byte
|
||||
var autoDisablePolicy []byte
|
||||
var degradePolicy []byte
|
||||
if err := rows.Scan(
|
||||
&item.PlatformID,
|
||||
&item.PlatformKey,
|
||||
@@ -105,11 +107,16 @@ ORDER BY effective_priority ASC,
|
||||
&item.PricingMode,
|
||||
&item.DiscountFactor,
|
||||
&item.ModelPricingRuleSetID,
|
||||
&item.BasePricingRuleSetID,
|
||||
&permissionConfig,
|
||||
&modelRetryPolicy,
|
||||
&modelRateLimitPolicy,
|
||||
&item.RuntimePolicySetID,
|
||||
&runtimePolicyOverride,
|
||||
&runtimeRetryPolicy,
|
||||
&runtimeRateLimitPolicy,
|
||||
&autoDisablePolicy,
|
||||
°radePolicy,
|
||||
); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -126,6 +133,10 @@ ORDER BY effective_priority ASC,
|
||||
item.ModelRetryPolicy = decodeObject(modelRetryPolicy)
|
||||
item.ModelRateLimitPolicy = decodeObject(modelRateLimitPolicy)
|
||||
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
|
||||
item.RuntimeRetryPolicy = decodeObject(runtimeRetryPolicy)
|
||||
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)
|
||||
item.QueueKey = item.ClientID
|
||||
items = append(items, item)
|
||||
|
||||
@@ -17,7 +17,7 @@ type modelCatalogSnapshot struct {
|
||||
ProviderKey string
|
||||
CanonicalModelKey string
|
||||
ProviderModelName string
|
||||
ModelType string
|
||||
ModelType StringList
|
||||
DisplayName string
|
||||
Capabilities map[string]any
|
||||
BaseBillingConfig map[string]any
|
||||
@@ -83,9 +83,13 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
|
||||
if err != nil && !IsNotFound(err) {
|
||||
return PlatformModel{}, err
|
||||
}
|
||||
if input.ModelType == "" {
|
||||
if len(input.ModelType) == 0 {
|
||||
input.ModelType = base.ModelType
|
||||
}
|
||||
input.ModelType = normalizeModelTypeList(input.ModelType)
|
||||
if len(input.ModelType) == 0 {
|
||||
input.ModelType = StringList{"text_generate"}
|
||||
}
|
||||
if input.ModelName == "" {
|
||||
input.ModelName = base.ProviderModelName
|
||||
}
|
||||
@@ -104,11 +108,12 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
|
||||
if len(billingConfig) == 0 {
|
||||
billingConfig = mergeObjects(base.BaseBillingConfig, input.BillingConfigOverride)
|
||||
}
|
||||
explicitRuntimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
|
||||
rateLimitPolicy := input.RateLimitPolicy
|
||||
if len(rateLimitPolicy) == 0 {
|
||||
if len(rateLimitPolicy) == 0 && explicitRuntimePolicySetID == "" {
|
||||
rateLimitPolicy = base.DefaultRateLimitPolicy
|
||||
}
|
||||
runtimePolicySetID := strings.TrimSpace(input.RuntimePolicySetID)
|
||||
runtimePolicySetID := explicitRuntimePolicySetID
|
||||
if runtimePolicySetID == "" {
|
||||
runtimePolicySetID = base.RuntimePolicySetID
|
||||
}
|
||||
@@ -119,6 +124,7 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
|
||||
|
||||
capabilityOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.CapabilityOverride))
|
||||
capabilitiesJSON, _ := json.Marshal(emptyObjectIfNil(capabilities))
|
||||
modelTypeJSON, _ := json.Marshal(input.ModelType)
|
||||
billingOverrideJSON, _ := json.Marshal(emptyObjectIfNil(input.BillingConfigOverride))
|
||||
billingJSON, _ := json.Marshal(emptyObjectIfNil(billingConfig))
|
||||
permissionJSON, _ := json.Marshal(emptyObjectIfNil(input.PermissionConfig))
|
||||
@@ -144,6 +150,7 @@ func (s *Store) createPlatformModel(ctx context.Context, q platformModelQuerier,
|
||||
var retryPolicyBytes []byte
|
||||
var rateLimitPolicyBytes []byte
|
||||
var runtimePolicyOverrideBytes []byte
|
||||
var modelTypeBytes []byte
|
||||
err = q.QueryRow(ctx, `
|
||||
INSERT INTO platform_models (
|
||||
platform_id, base_model_id, model_name, model_alias, model_type, display_name,
|
||||
@@ -152,12 +159,12 @@ INSERT INTO platform_models (
|
||||
runtime_policy_set_id, runtime_policy_override, enabled
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, $2::uuid, $3, NULLIF($4, ''), $5, $6,
|
||||
$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
|
||||
)
|
||||
ON CONFLICT (platform_id, model_name, model_type) DO UPDATE
|
||||
ON CONFLICT (platform_id, model_name) DO UPDATE
|
||||
SET base_model_id = EXCLUDED.base_model_id,
|
||||
model_alias = EXCLUDED.model_alias,
|
||||
display_name = EXCLUDED.display_name,
|
||||
@@ -185,7 +192,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
|
||||
baseID,
|
||||
input.ModelName,
|
||||
input.ModelAlias,
|
||||
input.ModelType,
|
||||
string(modelTypeJSON),
|
||||
input.DisplayName,
|
||||
string(capabilityOverrideJSON),
|
||||
string(capabilitiesJSON),
|
||||
@@ -205,7 +212,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
|
||||
&model.BaseModelID,
|
||||
&model.ModelName,
|
||||
&model.ModelAlias,
|
||||
&model.ModelType,
|
||||
&modelTypeBytes,
|
||||
&model.DisplayName,
|
||||
&capabilityOverrideBytes,
|
||||
&capabilitiesBytes,
|
||||
@@ -228,6 +235,7 @@ RETURNING id::text, platform_id::text, COALESCE(base_model_id::text, ''), model_
|
||||
}
|
||||
model.CapabilityOverride = decodeObject(capabilityOverrideBytes)
|
||||
model.Capabilities = decodeObject(capabilitiesBytes)
|
||||
model.ModelType = decodeStringArray(modelTypeBytes)
|
||||
model.BillingConfigOverride = decodeObject(billingOverrideBytes)
|
||||
model.BillingConfig = decodeObject(billingBytes)
|
||||
model.PermissionConfig = decodeObject(permissionBytes)
|
||||
@@ -265,6 +273,7 @@ func (s *Store) lookupBaseModel(ctx context.Context, q platformModelQuerier, id
|
||||
var billingConfig []byte
|
||||
var rateLimitPolicy []byte
|
||||
var runtimePolicyOverride []byte
|
||||
var modelTypeBytes []byte
|
||||
err := q.QueryRow(ctx, `
|
||||
SELECT id::text, provider_key, canonical_model_key, provider_model_name, model_type, display_name,
|
||||
capabilities, base_billing_config, default_rate_limit_policy,
|
||||
@@ -279,7 +288,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
|
||||
&item.ProviderKey,
|
||||
&item.CanonicalModelKey,
|
||||
&item.ProviderModelName,
|
||||
&item.ModelType,
|
||||
&modelTypeBytes,
|
||||
&item.DisplayName,
|
||||
&capabilities,
|
||||
&billingConfig,
|
||||
@@ -297,6 +306,7 @@ LIMIT 1`, strings.TrimSpace(id), strings.TrimSpace(canonicalKey), strings.TrimSp
|
||||
item.BaseBillingConfig = decodeObject(billingConfig)
|
||||
item.DefaultRateLimitPolicy = decodeObject(rateLimitPolicy)
|
||||
item.RuntimePolicyOverride = decodeObject(runtimePolicyOverride)
|
||||
item.ModelType = normalizeModelTypeList(decodeStringArray(modelTypeBytes))
|
||||
return item, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -134,7 +134,7 @@ type PlatformModel struct {
|
||||
PlatformName string `json:"platformName,omitempty"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias,omitempty"`
|
||||
ModelType string `json:"modelType"`
|
||||
ModelType StringList `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
CapabilityOverride map[string]any `json:"capabilityOverride,omitempty"`
|
||||
Capabilities map[string]any `json:"capabilities,omitempty"`
|
||||
@@ -390,6 +390,8 @@ type GatewayTask struct {
|
||||
ResponseDurationMS int64 `json:"responseDurationMs"`
|
||||
FinishedAt string `json:"finishedAt,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
ErrorCode string `json:"errorCode,omitempty"`
|
||||
ErrorMessage string `json:"errorMessage,omitempty"`
|
||||
CreatedAt time.Time `json:"createdAt"`
|
||||
UpdatedAt time.Time `json:"updatedAt"`
|
||||
}
|
||||
@@ -404,6 +406,7 @@ request, status, COALESCE(result, '{}'::jsonb), COALESCE(billings, '[]'::jsonb),
|
||||
COALESCE(usage, '{}'::jsonb), COALESCE(metrics, '{}'::jsonb), COALESCE(billing_summary, '{}'::jsonb),
|
||||
COALESCE(final_charge_amount, 0)::float8, COALESCE(response_started_at::text, ''),
|
||||
COALESCE(response_finished_at::text, ''), COALESCE(response_duration_ms, 0), COALESCE(error, ''),
|
||||
COALESCE(error_code, ''), COALESCE(error_message, ''),
|
||||
created_at, updated_at, COALESCE(finished_at::text, '')`
|
||||
|
||||
type TaskEvent struct {
|
||||
@@ -713,6 +716,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
|
||||
var retryPolicy []byte
|
||||
var rateLimitPolicy []byte
|
||||
var runtimePolicyOverride []byte
|
||||
var modelTypeBytes []byte
|
||||
if err := rows.Scan(
|
||||
&model.ID,
|
||||
&model.PlatformID,
|
||||
@@ -721,7 +725,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
|
||||
&model.PlatformName,
|
||||
&model.ModelName,
|
||||
&model.ModelAlias,
|
||||
&model.ModelType,
|
||||
&modelTypeBytes,
|
||||
&model.DisplayName,
|
||||
&capabilityOverride,
|
||||
&capabilities,
|
||||
@@ -743,6 +747,7 @@ ORDER BY m.model_type ASC, m.model_name ASC`, args...)
|
||||
}
|
||||
model.CapabilityOverride = decodeObject(capabilityOverride)
|
||||
model.Capabilities = decodeObject(capabilities)
|
||||
model.ModelType = decodeStringArray(modelTypeBytes)
|
||||
model.BillingConfigOverride = decodeObject(billingConfigOverride)
|
||||
model.BillingConfig = decodeObject(billingConfig)
|
||||
model.PermissionConfig = decodeObject(permissionConfig)
|
||||
@@ -1231,7 +1236,7 @@ func (s *Store) VerifyLocalAPIKey(ctx context.Context, secret string) (*auth.Use
|
||||
return nil, auth.ErrUnauthorized
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT k.id::text, k.key_hash, k.key_prefix, k.name, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''),
|
||||
SELECT k.id::text, k.key_hash, k.key_prefix, k.name, k.scopes, COALESCE(k.user_group_id::text, u.default_user_group_id::text, ''),
|
||||
u.id::text, u.username, u.roles, COALESCE(u.gateway_tenant_id::text, ''),
|
||||
COALESCE(u.tenant_id, ''), COALESCE(u.tenant_key, '')
|
||||
FROM gateway_api_keys k
|
||||
@@ -1252,6 +1257,7 @@ WHERE k.key_prefix = $1
|
||||
var hash string
|
||||
var keyPrefix string
|
||||
var keyName string
|
||||
var scopesBytes []byte
|
||||
var userGroupID string
|
||||
var gatewayUserID string
|
||||
var username string
|
||||
@@ -1259,7 +1265,7 @@ WHERE k.key_prefix = $1
|
||||
var gatewayTenantID string
|
||||
var tenantID string
|
||||
var tenantKey string
|
||||
if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
|
||||
if err := rows.Scan(&apiKeyID, &hash, &keyPrefix, &keyName, &scopesBytes, &userGroupID, &gatewayUserID, &username, &rolesBytes, &gatewayTenantID, &tenantID, &tenantKey); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if bcrypt.CompareHashAndPassword([]byte(hash), []byte(secret)) != nil {
|
||||
@@ -1281,6 +1287,7 @@ WHERE k.key_prefix = $1
|
||||
APIKeyID: apiKeyID,
|
||||
APIKeyName: keyName,
|
||||
APIKeyPrefix: keyPrefix,
|
||||
APIKeyScopes: decodeStringArray(scopesBytes),
|
||||
}, nil
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
@@ -1661,6 +1668,8 @@ func scanGatewayTask(scanner taskScanner) (GatewayTask, error) {
|
||||
&task.ResponseFinishedAt,
|
||||
&task.ResponseDurationMS,
|
||||
&task.Error,
|
||||
&task.ErrorCode,
|
||||
&task.ErrorMessage,
|
||||
&task.CreatedAt,
|
||||
&task.UpdatedAt,
|
||||
&task.FinishedAt,
|
||||
|
||||
@@ -176,6 +176,63 @@ func (s *Store) DeletePricingRuleSet(ctx context.Context, id string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) PricingRuleSetBillingConfig(ctx context.Context, id string) (map[string]any, error) {
|
||||
id = strings.TrimSpace(id)
|
||||
if id == "" {
|
||||
return nil, nil
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT resource_type, base_price::float8, dynamic_weight
|
||||
FROM model_pricing_rules
|
||||
WHERE rule_set_id = $1::uuid
|
||||
AND status = 'active'
|
||||
ORDER BY priority ASC, resource_type ASC`, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
config := map[string]any{}
|
||||
for rows.Next() {
|
||||
var resourceType string
|
||||
var basePrice float64
|
||||
var dynamicWeightBytes []byte
|
||||
if err := rows.Scan(&resourceType, &basePrice, &dynamicWeightBytes); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
dynamicWeight := decodeObject(dynamicWeightBytes)
|
||||
switch resourceType {
|
||||
case "text_input":
|
||||
config["textInputPer1k"] = basePrice
|
||||
case "text_output":
|
||||
config["textOutputPer1k"] = basePrice
|
||||
case "image":
|
||||
config["imageBase"] = basePrice
|
||||
config["image"] = pricingResourceConfig(basePrice, dynamicWeight)
|
||||
case "image_edit":
|
||||
config["editBase"] = basePrice
|
||||
config["image_edit"] = pricingResourceConfig(basePrice, dynamicWeight)
|
||||
case "video":
|
||||
config["videoBase"] = basePrice
|
||||
config["video"] = pricingResourceConfig(basePrice, dynamicWeight)
|
||||
default:
|
||||
config[resourceType] = pricingResourceConfig(basePrice, dynamicWeight)
|
||||
}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return config, nil
|
||||
}
|
||||
|
||||
func pricingResourceConfig(basePrice float64, dynamicWeight map[string]any) map[string]any {
|
||||
config := map[string]any{"basePrice": basePrice}
|
||||
if len(dynamicWeight) > 0 {
|
||||
config["dynamicWeight"] = dynamicWeight
|
||||
}
|
||||
return config
|
||||
}
|
||||
|
||||
func insertPricingRules(ctx context.Context, tx pgx.Tx, ruleSetID string, defaultCurrency string, rules []PricingRuleInput) error {
|
||||
for index, rule := range rules {
|
||||
rule = normalizePricingRule(rule, index, defaultCurrency)
|
||||
|
||||
@@ -107,6 +107,33 @@ func (s *Store) DeleteRuntimePolicySet(ctx context.Context, id string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Store) DisableCandidatePlatform(ctx context.Context, platformID string) error {
|
||||
if strings.TrimSpace(platformID) == "" {
|
||||
return nil
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE integration_platforms
|
||||
SET status = 'disabled',
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, platformID)
|
||||
return err
|
||||
}
|
||||
|
||||
func (s *Store) CooldownCandidatePlatform(ctx context.Context, platformID string, cooldownSeconds int) error {
|
||||
if strings.TrimSpace(platformID) == "" {
|
||||
return nil
|
||||
}
|
||||
if cooldownSeconds <= 0 {
|
||||
cooldownSeconds = 300
|
||||
}
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
UPDATE integration_platforms
|
||||
SET cooldown_until = now() + ($2::int * interval '1 second'),
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, platformID, cooldownSeconds)
|
||||
return err
|
||||
}
|
||||
|
||||
func scanRuntimePolicySet(scanner runtimePolicyScanner) (RuntimePolicySet, error) {
|
||||
var item RuntimePolicySet
|
||||
var rateLimitPolicy []byte
|
||||
|
||||
@@ -16,7 +16,7 @@ type CreatePlatformModelInput struct {
|
||||
CanonicalModelKey string `json:"canonicalModelKey"`
|
||||
ModelName string `json:"modelName"`
|
||||
ModelAlias string `json:"modelAlias"`
|
||||
ModelType string `json:"modelType"`
|
||||
ModelType StringList `json:"modelType"`
|
||||
DisplayName string `json:"displayName"`
|
||||
CapabilityOverride map[string]any `json:"capabilityOverride"`
|
||||
Capabilities map[string]any `json:"capabilities"`
|
||||
@@ -64,12 +64,17 @@ type RuntimeModelCandidate struct {
|
||||
PermissionConfig map[string]any
|
||||
PricingMode string
|
||||
DiscountFactor float64
|
||||
BasePricingRuleSetID string
|
||||
PlatformPricingRuleSetID string
|
||||
ModelPricingRuleSetID string
|
||||
ModelRetryPolicy map[string]any
|
||||
ModelRateLimitPolicy map[string]any
|
||||
RuntimePolicySetID string
|
||||
RuntimePolicyOverride map[string]any
|
||||
RuntimeRetryPolicy map[string]any
|
||||
RuntimeRateLimitPolicy map[string]any
|
||||
AutoDisablePolicy map[string]any
|
||||
DegradePolicy map[string]any
|
||||
ClientID string
|
||||
QueueKey string
|
||||
}
|
||||
|
||||
@@ -3,9 +3,67 @@ package store
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
||||
"github.com/jackc/pgx/v5"
|
||||
"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
|
||||
}
|
||||
if limit > 100 {
|
||||
limit = 100
|
||||
}
|
||||
gatewayUserID := localGatewayUserID(user)
|
||||
apiKeyID := ""
|
||||
userID := ""
|
||||
if user != nil {
|
||||
apiKeyID = strings.TrimSpace(user.APIKeyID)
|
||||
userID = strings.TrimSpace(user.ID)
|
||||
}
|
||||
if gatewayUserID == "" && userID == "" {
|
||||
return nil, ErrLocalUserRequired
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `
|
||||
SELECT `+gatewayTaskColumns+`
|
||||
FROM gateway_tasks
|
||||
WHERE (
|
||||
(
|
||||
NULLIF($1, '')::uuid IS NOT NULL
|
||||
AND gateway_user_id = NULLIF($1, '')::uuid
|
||||
)
|
||||
OR (
|
||||
NULLIF($1, '')::uuid IS NULL
|
||||
AND NULLIF($2, '') IS NOT NULL
|
||||
AND user_id = $2
|
||||
)
|
||||
)
|
||||
AND (
|
||||
NULLIF($3, '') IS NULL
|
||||
OR api_key_id = $3
|
||||
)
|
||||
ORDER BY created_at DESC
|
||||
LIMIT $4`, gatewayUserID, userID, apiKeyID, limit)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
items := make([]GatewayTask, 0)
|
||||
for rows.Next() {
|
||||
task, err := scanGatewayTask(rows)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, task)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
func (s *Store) MarkTaskRunning(ctx context.Context, taskID string, modelType string, normalizedRequest map[string]any) error {
|
||||
normalizedJSON, _ := json.Marshal(emptyObjectIfNil(normalizedRequest))
|
||||
_, err := s.pool.Exec(ctx, `
|
||||
@@ -140,6 +198,103 @@ WHERE id = $1::uuid`,
|
||||
return s.GetTask(ctx, input.TaskID)
|
||||
}
|
||||
|
||||
func (s *Store) SettleTaskBilling(ctx context.Context, task GatewayTask) error {
|
||||
if task.FinalChargeAmount <= 0 || strings.TrimSpace(task.GatewayUserID) == "" {
|
||||
return nil
|
||||
}
|
||||
currency := strings.TrimSpace(taskBillingString(task.BillingSummary["currency"]))
|
||||
if currency == "" || currency == "mixed" {
|
||||
currency = "resource"
|
||||
}
|
||||
metadata, _ := json.Marshal(map[string]any{
|
||||
"taskId": task.ID,
|
||||
"kind": task.Kind,
|
||||
"model": task.Model,
|
||||
"resolvedModel": task.ResolvedModel,
|
||||
"billings": task.Billings,
|
||||
"billingSummary": task.BillingSummary,
|
||||
})
|
||||
return pgx.BeginFunc(ctx, s.pool, func(tx pgx.Tx) error {
|
||||
if _, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_accounts (
|
||||
gateway_tenant_id, gateway_user_id, tenant_id, tenant_key, user_id, currency
|
||||
)
|
||||
VALUES (NULLIF($1, '')::uuid, $2::uuid, NULLIF($3, ''), NULLIF($4, ''), NULLIF($5, ''), $6)
|
||||
ON CONFLICT (gateway_user_id, currency) DO NOTHING`,
|
||||
task.GatewayTenantID, task.GatewayUserID, task.TenantID, task.TenantKey, task.UserID, currency); err != nil {
|
||||
return err
|
||||
}
|
||||
var exists bool
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT EXISTS (
|
||||
SELECT 1
|
||||
FROM gateway_wallet_transactions t
|
||||
JOIN gateway_wallet_accounts a ON a.id = t.account_id
|
||||
WHERE a.gateway_user_id = $1::uuid
|
||||
AND a.currency = $2
|
||||
AND t.idempotency_key = $3
|
||||
)`, task.GatewayUserID, currency, billingIdempotencyKey(task.ID)).Scan(&exists); err != nil {
|
||||
return err
|
||||
}
|
||||
if exists {
|
||||
return nil
|
||||
}
|
||||
var accountID string
|
||||
var balanceBefore float64
|
||||
var gatewayTenantID string
|
||||
if err := tx.QueryRow(ctx, `
|
||||
SELECT id::text, balance::float8, COALESCE(gateway_tenant_id::text, '')
|
||||
FROM gateway_wallet_accounts
|
||||
WHERE gateway_user_id = $1::uuid
|
||||
AND currency = $2
|
||||
FOR UPDATE`, task.GatewayUserID, currency).Scan(&accountID, &balanceBefore, &gatewayTenantID); err != nil {
|
||||
return err
|
||||
}
|
||||
amount := roundMoney(task.FinalChargeAmount)
|
||||
balanceAfter := roundMoney(balanceBefore - amount)
|
||||
if _, err := tx.Exec(ctx, `
|
||||
UPDATE gateway_wallet_accounts
|
||||
SET balance = $2,
|
||||
total_spent = total_spent + $3,
|
||||
updated_at = now()
|
||||
WHERE id = $1::uuid`, accountID, balanceAfter, amount); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err := tx.Exec(ctx, `
|
||||
INSERT INTO gateway_wallet_transactions (
|
||||
account_id, gateway_tenant_id, gateway_user_id, direction, transaction_type,
|
||||
amount, balance_before, balance_after, idempotency_key, reference_type, reference_id, metadata
|
||||
)
|
||||
VALUES (
|
||||
$1::uuid, NULLIF($2, '')::uuid, $3::uuid, 'debit', 'task_billing',
|
||||
$4, $5, $6, $7, 'gateway_task', $8, $9::jsonb
|
||||
)`,
|
||||
accountID, firstNonEmpty(gatewayTenantID, task.GatewayTenantID), task.GatewayUserID, amount, roundMoney(balanceBefore), balanceAfter, billingIdempotencyKey(task.ID), task.ID, string(metadata))
|
||||
if pgErr, ok := err.(*pgconn.PgError); ok && pgErr.Code == "23505" {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
})
|
||||
}
|
||||
|
||||
func billingIdempotencyKey(taskID string) string {
|
||||
return "task:" + taskID + ":billing"
|
||||
}
|
||||
|
||||
func roundMoney(value float64) float64 {
|
||||
if value < 0 {
|
||||
return -roundMoney(-value)
|
||||
}
|
||||
return float64(int64(value*1000000+0.5)) / 1000000
|
||||
}
|
||||
|
||||
func taskBillingString(value any) string {
|
||||
if text, ok := value.(string); ok {
|
||||
return text
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (s *Store) FinishTaskFailure(ctx context.Context, input FinishTaskFailureInput) (GatewayTask, error) {
|
||||
metricsJSON, _ := json.Marshal(emptyObjectIfNil(input.Metrics))
|
||||
if _, err := s.pool.Exec(ctx, `
|
||||
|
||||
Reference in New Issue
Block a user