BCAI/easyai-ai-gateway
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add priority demotion controls

Browse Source
This commit is contained in:
wangbo
2026-05-12 18:43:20 +08:00
parent 3917b84b5d
commit 98abd247d6
22 changed files with 917 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/api/internal/store/rate_limit_status.go
+159 -20
View File
@@ -2,8 +2,11 @@ package store
import (
"context"
"database/sql"
"sort"
"strconv"
"strings"
"time"
)
type RateLimitMetricStatus struct {
@@ -17,30 +20,53 @@ type RateLimitMetricStatus struct {
}
type ModelRateLimitStatus struct {
PlatformModelID string `json:"platformModelId"`
PlatformID string `json:"platformId"`
PlatformName string `json:"platformName"`
Provider string `json:"provider"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
PlatformCooldownUntil string `json:"platformCooldownUntil,omitempty"`
ModelCooldownUntil string `json:"modelCooldownUntil,omitempty"`
Concurrent RateLimitMetricStatus `json:"concurrent"`
QueuedTasks float64 `json:"queuedTasks"`
RPM RateLimitMetricStatus `json:"rpm"`
TPM RateLimitMetricStatus `json:"tpm"`
LoadRatio float64 `json:"loadRatio"`
PlatformModelID string `json:"platformModelId"`
PlatformID string `json:"platformId"`
PlatformName string `json:"platformName"`
Provider string `json:"provider"`
PlatformPriority int `json:"platformPriority"`
PlatformDynamicPriority *int `json:"platformDynamicPriority,omitempty"`
PlatformEffectivePriority int `json:"platformEffectivePriority"`
ModelName string `json:"modelName"`
ProviderModelName string `json:"providerModelName,omitempty"`
ModelAlias string `json:"modelAlias,omitempty"`
DisplayName string `json:"displayName"`
ModelType []string `json:"modelType"`
Enabled bool `json:"enabled"`
RateLimitPolicy map[string]any `json:"rateLimitPolicy,omitempty"`
PlatformCooldownUntil string `json:"platformCooldownUntil,omitempty"`
ModelCooldownUntil string `json:"modelCooldownUntil,omitempty"`
Concurrent RateLimitMetricStatus `json:"concurrent"`
QueuedTasks float64 `json:"queuedTasks"`
RPM RateLimitMetricStatus `json:"rpm"`
TPM RateLimitMetricStatus `json:"tpm"`
LoadRatio float64 `json:"loadRatio"`
RecentPriorityDemotions []PriorityDemotionRecord `json:"recentPriorityDemotions,omitempty"`
}
type PriorityDemotionRecord struct {
ID string `json:"id"`
TaskID string `json:"taskId"`
PlatformID string `json:"platformId"`
PlatformModelID string `json:"platformModelId,omitempty"`
Reason string `json:"reason,omitempty"`
ErrorCode string `json:"errorCode,omitempty"`
ErrorMessage string `json:"errorMessage,omitempty"`
Category string `json:"category,omitempty"`
StatusCode int `json:"statusCode,omitempty"`
PolicySource string `json:"policySource,omitempty"`
Policy string `json:"policy,omitempty"`
PolicyRule string `json:"policyRule,omitempty"`
MatchedValue string `json:"matchedValue,omitempty"`
DynamicPriority int `json:"dynamicPriority,omitempty"`
CreatedAt time.Time `json:"createdAt"`
}
func (s *Store) ListModelRateLimitStatuses(ctx context.Context) ([]ModelRateLimitStatus, error) {
rows, err := s.pool.Query(ctx, `
SELECT m.id::text, m.platform_id::text, p.name, p.provider,
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
SELECT m.id::text, m.platform_id::text, p.name, p.provider,
p.priority, p.dynamic_priority, COALESCE(p.dynamic_priority, p.priority),
m.model_name, COALESCE(NULLIF(m.provider_model_name, ''), m.model_name), COALESCE(m.model_alias, ''),
m.model_type, m.display_name, m.enabled,
p.rate_limit_policy, COALESCE(rp.rate_limit_policy, '{}'::jsonb), COALESCE(NULLIF(m.runtime_policy_override, '{}'::jsonb), b.runtime_policy_override, '{}'::jsonb), m.rate_limit_policy,
COALESCE(to_char(p.cooldown_until AT TIME ZONE 'UTC', 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'), ''),
@@ -119,6 +145,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
var runtimePolicyBytes []byte
var runtimeOverrideBytes []byte
var modelPolicyBytes []byte
var platformDynamicPriority sql.NullInt64
var platformCooldownUntil string
var modelCooldownUntil string
var concurrentCurrent float64
@@ -134,6 +161,9 @@ ORDER BY p.priority ASC, m.model_name ASC`)
&item.PlatformID,
&item.PlatformName,
&item.Provider,
&item.PlatformPriority,
&platformDynamicPriority,
&item.PlatformEffectivePriority,
&item.ModelName,
&item.ProviderModelName,
&item.ModelAlias,
@@ -157,6 +187,7 @@ ORDER BY p.priority ASC, m.model_name ASC`)
); err != nil {
return nil, err
}
item.PlatformDynamicPriority = intPointerFromNull(platformDynamicPriority)
item.ModelType = decodeStringArray(modelTypeBytes)
policy := effectiveModelRateLimitPolicy(
decodeObject(platformPolicyBytes),
@@ -177,6 +208,13 @@ ORDER BY p.priority ASC, m.model_name ASC`)
if err := rows.Err(); err != nil {
return nil, err
}
demotions, err := s.listRecentPriorityDemotionsByPlatform(ctx, items, 10)
if err != nil {
return nil, err
}
for index := range items {
items[index].RecentPriorityDemotions = demotions[items[index].PlatformID]
}
sort.SliceStable(items, func(i, j int) bool {
if items[i].LoadRatio == items[j].LoadRatio {
return strings.ToLower(items[i].DisplayName) < strings.ToLower(items[j].DisplayName)
@@ -186,6 +224,91 @@ ORDER BY p.priority ASC, m.model_name ASC`)
return items, nil
}
func (s *Store) listRecentPriorityDemotionsByPlatform(ctx context.Context, statuses []ModelRateLimitStatus, limit int) (map[string][]PriorityDemotionRecord, error) {
out := map[string][]PriorityDemotionRecord{}
if limit <= 0 || len(statuses) == 0 {
return out, nil
}
seen := map[string]bool{}
platformIDs := make([]string, 0, len(statuses))
for _, status := range statuses {
platformID := strings.TrimSpace(status.PlatformID)
if platformID == "" || seen[platformID] {
continue
}
seen[platformID] = true
platformIDs = append(platformIDs, platformID)
}
if len(platformIDs) == 0 {
return out, nil
}
rows, err := s.pool.Query(ctx, `
SELECT id::text, task_id::text, COALESCE(message, ''), payload, created_at
FROM (
SELECT e.*,
row_number() OVER (
PARTITION BY e.payload->>'platformId'
ORDER BY e.created_at DESC, e.seq DESC
) AS demotion_rank
FROM gateway_task_events e
WHERE e.event_type = 'task.policy.priority_demoted'
AND e.payload->>'platformId' = ANY($1::text[])
) ranked
WHERE demotion_rank <= $2
ORDER BY payload->>'platformId' ASC, created_at DESC`, platformIDs, limit)
if err != nil {
return nil, err
}
defer rows.Close()
for rows.Next() {
var id string
var taskID string
var message string
var payloadBytes []byte
var createdAt time.Time
if err := rows.Scan(&id, &taskID, &message, &payloadBytes, &createdAt); err != nil {
return nil, err
}
record := priorityDemotionRecordFromEventPayload(id, taskID, message, decodeObject(payloadBytes), createdAt)
if record.PlatformID == "" {
continue
}
out[record.PlatformID] = append(out[record.PlatformID], record)
}
return out, rows.Err()
}
func priorityDemotionRecordFromEventPayload(id string, taskID string, message string, payload map[string]any, createdAt time.Time) PriorityDemotionRecord {
errorMessage := stringValue(payload["errorMessage"])
if errorMessage == "" {
errorMessage = stringValue(payload["message"])
}
if errorMessage == "" {
errorMessage = strings.TrimSpace(message)
}
errorCode := stringValue(payload["errorCode"])
if errorCode == "" {
errorCode = stringValue(payload["code"])
}
return PriorityDemotionRecord{
ID: id,
TaskID: taskID,
PlatformID: stringValue(payload["platformId"]),
PlatformModelID: stringValue(payload["platformModelId"]),
Reason: stringValue(payload["reason"]),
ErrorCode: errorCode,
ErrorMessage: errorMessage,
Category: stringValue(payload["category"]),
StatusCode: intValue(payload["statusCode"]),
PolicySource: stringValue(payload["policySource"]),
Policy: stringValue(payload["policy"]),
PolicyRule: stringValue(payload["policyRule"]),
MatchedValue: stringValue(payload["matchedValue"]),
DynamicPriority: intValue(payload["dynamicPriority"]),
CreatedAt: createdAt,
}
}
func effectiveModelRateLimitPolicy(platformPolicy map[string]any, runtimePolicy map[string]any, runtimeOverride map[string]any, modelPolicy map[string]any) map[string]any {
policy := platformPolicy
if hasRateLimitRules(runtimePolicy) {
@@ -279,3 +402,19 @@ func floatValue(value any) float64 {
return 0
}
}
func intValue(value any) int {
switch typed := value.(type) {
case int:
return typed
case int64:
return int(typed)
case float64:
return int(typed)
case string:
parsed, _ := strconv.Atoi(strings.TrimSpace(typed))
return parsed
default:
return 0
}
}