AI Gateway Go 0.10.0 源码快照 + 旗舰版需求规划报告
M0-M7 已完成:核心网关(身份/RBAC/TOTP/OIDC/SAML/Provider/配额/路由/内容策略/审计/定价)+ 资源市场(MCP/Skills/数字员工)。 含 22 个 PostgreSQL 迁移、管理端/门户端前端源码、OpenAPI 契约、部署 compose。 Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,42 @@
|
||||
package apikey
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
type UsageStore struct{ redis *redis.Client }
|
||||
|
||||
func NewUsageStore(client *redis.Client) *UsageStore { return &UsageStore{redis: client} }
|
||||
|
||||
func MonthlyTokenUsageKey(apiKeyID string, now time.Time) string {
|
||||
return fmt.Sprintf("gateway:usage:api-key:%s:tokens:%s", apiKeyID, now.UTC().Format("200601"))
|
||||
}
|
||||
|
||||
func (s *UsageStore) MonthlyTokens(ctx context.Context, apiKeyIDs []string, now time.Time) (map[string]int64, error) {
|
||||
usage := make(map[string]int64, len(apiKeyIDs))
|
||||
if len(apiKeyIDs) == 0 || s == nil || s.redis == nil {
|
||||
return usage, nil
|
||||
}
|
||||
pipe := s.redis.Pipeline()
|
||||
commands := make(map[string]*redis.StringCmd, len(apiKeyIDs))
|
||||
for _, id := range apiKeyIDs {
|
||||
commands[id] = pipe.Get(ctx, MonthlyTokenUsageKey(id, now))
|
||||
}
|
||||
_, err := pipe.Exec(ctx)
|
||||
if err != nil && err != redis.Nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, err)
|
||||
}
|
||||
for id, command := range commands {
|
||||
value, commandErr := command.Int64()
|
||||
if commandErr == nil {
|
||||
usage[id] = max(value, 0)
|
||||
} else if commandErr != redis.Nil {
|
||||
return nil, fmt.Errorf("%w: %v", ErrStore, commandErr)
|
||||
}
|
||||
}
|
||||
return usage, nil
|
||||
}
|
||||
Reference in New Issue
Block a user