90 lines
2.7 KiB
Go
90 lines
2.7 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"strings"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/auth"
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
|
|
)
|
|
|
|
func (s *Service) rateLimitReservations(ctx context.Context, user *auth.User, candidate store.RuntimeModelCandidate, body map[string]any) []store.RateLimitReservation {
|
|
out := make([]store.RateLimitReservation, 0)
|
|
out = append(out, reservationsFromPolicy("platform_model", candidate.PlatformModelID, effectiveRateLimitPolicy(candidate), body)...)
|
|
if group, err := s.store.ResolveUserGroupPolicy(ctx, user); err == nil && group.ID != "" {
|
|
out = append(out, reservationsFromPolicy("user_group", group.ID, group.RateLimitPolicy, body)...)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func effectiveRateLimitPolicy(candidate store.RuntimeModelCandidate) map[string]any {
|
|
if hasRules(candidate.ModelRateLimitPolicy) {
|
|
return candidate.ModelRateLimitPolicy
|
|
}
|
|
if hasRules(candidate.PlatformRateLimitPolicy) {
|
|
return candidate.PlatformRateLimitPolicy
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func reservationsFromPolicy(scopeType string, scopeKey string, policy map[string]any, body map[string]any) []store.RateLimitReservation {
|
|
if scopeKey == "" || !hasRules(policy) {
|
|
return nil
|
|
}
|
|
rules, _ := policy["rules"].([]any)
|
|
out := make([]store.RateLimitReservation, 0, len(rules))
|
|
estimatedTokens := estimateRequestTokens(body)
|
|
for _, rawRule := range rules {
|
|
rule, _ := rawRule.(map[string]any)
|
|
metric := strings.TrimSpace(stringFromMap(rule, "metric"))
|
|
limit := floatFromAny(rule["limit"])
|
|
amount := 1.0
|
|
if strings.HasPrefix(metric, "tpm") {
|
|
amount = float64(estimatedTokens)
|
|
}
|
|
out = append(out, store.RateLimitReservation{
|
|
ScopeType: scopeType,
|
|
ScopeKey: scopeKey,
|
|
Metric: metric,
|
|
Limit: limit,
|
|
Amount: amount,
|
|
WindowSeconds: int(floatFromAny(rule["windowSeconds"])),
|
|
LeaseTTLSeconds: int(floatFromAny(rule["leaseTtlSeconds"])),
|
|
})
|
|
}
|
|
return out
|
|
}
|
|
|
|
func hasRules(policy map[string]any) bool {
|
|
rules, _ := policy["rules"].([]any)
|
|
return len(rules) > 0
|
|
}
|
|
|
|
func estimateRequestTokens(body map[string]any) int {
|
|
text := ""
|
|
if prompt := stringFromMap(body, "prompt"); prompt != "" {
|
|
text += prompt
|
|
}
|
|
if input := stringFromMap(body, "input"); input != "" {
|
|
text += input
|
|
}
|
|
if messages, ok := body["messages"].([]any); ok {
|
|
for _, raw := range messages {
|
|
message, _ := raw.(map[string]any)
|
|
switch content := message["content"].(type) {
|
|
case string:
|
|
text += content
|
|
case []any:
|
|
for _, rawPart := range content {
|
|
part, _ := rawPart.(map[string]any)
|
|
text += stringFromMap(part, "text")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
if text == "" {
|
|
return 1
|
|
}
|
|
return len([]rune(text))/4 + 1
|
|
}
|