0.11.0: 旗舰版功能补齐(License/登录记录/会话管理/角色管理/门户定时任务/模型配额/输出脱敏/供应链扫描/记忆管理/AI助手/真实概览)
- 新增迁移 000031-000034(登录日志/角色/模型配额/记忆) - 新增包: license/memory/modelquota/assistant,扫描引擎 - 全部功能后端+前端+端到端验证通过(25 包单测)
This commit is contained in:
@@ -0,0 +1,98 @@
|
||||
package modelquota
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"aigateway.local/core/internal/identity"
|
||||
"aigateway.local/core/internal/platform/apiresponse"
|
||||
)
|
||||
|
||||
// HTTPHandler 提供管理端模型配额 CRUD。
|
||||
type HTTPHandler struct {
|
||||
service *Service
|
||||
identity *identity.Service
|
||||
mux *http.ServeMux
|
||||
}
|
||||
|
||||
func NewHTTPHandler(service *Service, identityService *identity.Service) *HTTPHandler {
|
||||
h := &HTTPHandler{service: service, identity: identityService, mux: http.NewServeMux()}
|
||||
h.mux.HandleFunc("GET /api/v1/admin/model-quotas", h.list)
|
||||
h.mux.HandleFunc("POST /api/v1/admin/model-quotas", h.save)
|
||||
h.mux.HandleFunc("PUT /api/v1/admin/model-quotas/{quota_id}", h.save)
|
||||
h.mux.HandleFunc("DELETE /api/v1/admin/model-quotas/{quota_id}", h.delete)
|
||||
return h
|
||||
}
|
||||
|
||||
func (h *HTTPHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) { h.mux.ServeHTTP(w, r) }
|
||||
|
||||
func (h *HTTPHandler) require(w http.ResponseWriter, r *http.Request, permission string) (identity.Account, bool) {
|
||||
account, err := h.identity.Authenticate(r.Context(), identity.KindAdmin, r.Header.Get("Authorization"))
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusUnauthorized, "登录状态无效或已过期")
|
||||
return identity.Account{}, false
|
||||
}
|
||||
if !identity.HasPermission(account, permission) {
|
||||
apiresponse.Error(w, http.StatusForbidden, "缺少模型配额管理权限")
|
||||
return identity.Account{}, false
|
||||
}
|
||||
return account, true
|
||||
}
|
||||
|
||||
func (h *HTTPHandler) list(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPricingRead); !ok {
|
||||
return
|
||||
}
|
||||
items, err := h.service.List(r.Context())
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusServiceUnavailable, "模型配额查询失败")
|
||||
return
|
||||
}
|
||||
apiresponse.OK(w, items)
|
||||
}
|
||||
|
||||
type quotaInput struct {
|
||||
ProviderCode string `json:"provider_code"`
|
||||
ModelPattern string `json:"model_pattern"`
|
||||
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
|
||||
Enabled *bool `json:"enabled"`
|
||||
}
|
||||
|
||||
func (h *HTTPHandler) save(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPricingManage); !ok {
|
||||
return
|
||||
}
|
||||
var input quotaInput
|
||||
decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20))
|
||||
decoder.DisallowUnknownFields()
|
||||
if decoder.Decode(&input) != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, "请求格式无效")
|
||||
return
|
||||
}
|
||||
enabled := true
|
||||
if input.Enabled != nil {
|
||||
enabled = *input.Enabled
|
||||
}
|
||||
item, err := h.service.Save(r.Context(), r.PathValue("quota_id"), input.ProviderCode, input.ModelPattern, input.MonthlyTokenQuota, enabled)
|
||||
if err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
_ = h.service.Reload(r.Context())
|
||||
apiresponse.OK(w, item)
|
||||
}
|
||||
|
||||
func (h *HTTPHandler) delete(w http.ResponseWriter, r *http.Request) {
|
||||
if _, ok := h.require(w, r, identity.PermissionPricingManage); !ok {
|
||||
return
|
||||
}
|
||||
if err := h.service.Delete(r.Context(), r.PathValue("quota_id")); err != nil {
|
||||
apiresponse.Error(w, http.StatusBadRequest, err.Error())
|
||||
return
|
||||
}
|
||||
_ = h.service.Reload(r.Context())
|
||||
apiresponse.OK(w, map[string]bool{"deleted": true})
|
||||
}
|
||||
|
||||
var _ = strings.TrimSpace
|
||||
@@ -0,0 +1,313 @@
|
||||
// Package modelquota 实现模型级 Token 配额:按 Provider+模型模式设置
|
||||
// 企业总配额,所有 API Key 共享同一自然月计数(与 API Key 级配额叠加),
|
||||
// 用于模型级成本管控。
|
||||
package modelquota
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
platformid "aigateway.local/core/internal/platform/id"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
var ErrUnavailable = errors.New("model quota unavailable")
|
||||
|
||||
// Quota 是一条模型配额记录。
|
||||
type Quota struct {
|
||||
ID string `json:"id"`
|
||||
ProviderCode string `json:"provider_code"`
|
||||
ModelPattern string `json:"model_pattern"`
|
||||
MonthlyTokenQuota int64 `json:"monthly_token_quota"`
|
||||
Enabled bool `json:"enabled"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
// Service 持有配额快照(定时刷新)并提供 Redis 原子预留/提交。
|
||||
type Service struct {
|
||||
pool *pgxpool.Pool
|
||||
client *redis.Client
|
||||
logger *slog.Logger
|
||||
snapshot atomic.Pointer[quotaSnapshot]
|
||||
reserve *redis.Script
|
||||
commitScript *redis.Script
|
||||
mu sync.Mutex // 保护 admin 写路径
|
||||
}
|
||||
|
||||
type quotaSnapshot struct {
|
||||
quotas []Quota
|
||||
}
|
||||
|
||||
// NewService 创建服务(pool=PostgreSQL, client=critical Redis)。
|
||||
func NewService(pool *pgxpool.Pool, client *redis.Client, logger *slog.Logger) *Service {
|
||||
return &Service{
|
||||
pool: pool, client: client, logger: logger,
|
||||
reserve: redis.NewScript(modelReserveScript), commitScript: redis.NewScript(modelCommitScript),
|
||||
}
|
||||
}
|
||||
|
||||
// Reload 从数据库刷新配额快照。
|
||||
func (s *Service) Reload(ctx context.Context) error {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas WHERE enabled ORDER BY provider_code,model_pattern`)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Quota{}
|
||||
for rows.Next() {
|
||||
var q Quota
|
||||
if err := rows.Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt); err != nil {
|
||||
return err
|
||||
}
|
||||
items = append(items, q)
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
s.snapshot.Store("aSnapshot{quotas: items})
|
||||
return nil
|
||||
}
|
||||
|
||||
// Run 周期刷新快照。
|
||||
func (s *Service) Run(ctx context.Context, interval time.Duration) {
|
||||
if interval <= 0 {
|
||||
interval = 30 * time.Second
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if err := s.Reload(ctx); err != nil && s.logger != nil {
|
||||
s.logger.Warn("model quota refresh failed; retaining last snapshot", "error", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Lookup 返回匹配 (provider, model) 的最高配额(模型模式最具体优先);
|
||||
// 未配置返回 0(不限制)。
|
||||
func (s *Service) Lookup(providerCode, model string) int64 {
|
||||
if s == nil {
|
||||
return 0
|
||||
}
|
||||
current := s.snapshot.Load()
|
||||
if current == nil {
|
||||
return 0
|
||||
}
|
||||
providerCode = strings.ToLower(providerCode)
|
||||
best := int64(0)
|
||||
bestLen := -1
|
||||
for _, q := range current.quotas {
|
||||
if q.ProviderCode != providerCode {
|
||||
continue
|
||||
}
|
||||
if !matchPattern(q.ModelPattern, model) {
|
||||
continue
|
||||
}
|
||||
// 更具体的模式(更长前缀)优先;同长度取配额更大者(防御性)。
|
||||
if len(q.ModelPattern) > bestLen || (len(q.ModelPattern) == bestLen && q.MonthlyTokenQuota > best) {
|
||||
best = q.MonthlyTokenQuota
|
||||
bestLen = len(q.ModelPattern)
|
||||
}
|
||||
}
|
||||
return best
|
||||
}
|
||||
|
||||
func matchPattern(pattern, model string) bool {
|
||||
pattern = strings.TrimSpace(pattern)
|
||||
if pattern == "" || pattern == "*" {
|
||||
return true
|
||||
}
|
||||
if strings.HasSuffix(pattern, "*") {
|
||||
return strings.HasPrefix(model, strings.TrimSuffix(pattern, "*"))
|
||||
}
|
||||
return pattern == model
|
||||
}
|
||||
|
||||
// Reserve 为 (provider, model) 的月度计数预留 estimate;Allowed=false 表示超限。
|
||||
// 返回 any 以适配 gateway.ModelQuotaController 接口。
|
||||
func (s *Service) Reserve(ctx context.Context, providerCode, model string, estimate int64, now time.Time) (any, error) {
|
||||
if s == nil || s.client == nil {
|
||||
return Reservation{}, ErrUnavailable
|
||||
}
|
||||
if estimate < 0 {
|
||||
estimate = 0
|
||||
}
|
||||
now = now.UTC()
|
||||
key := monthlyKey(providerCode, model, now)
|
||||
reset := time.Date(now.Year(), now.Month()+1, 1, 0, 0, 0, 0, time.UTC)
|
||||
quota := s.Lookup(providerCode, model)
|
||||
if quota <= 0 {
|
||||
return Reservation{Allowed: true}, nil
|
||||
}
|
||||
result, err := s.reserve.Run(ctx, s.client, []string{key}, estimate, quota, int64(reset.Sub(now).Seconds())+86400).Slice()
|
||||
if err != nil || len(result) != 2 {
|
||||
return Reservation{}, fmt.Errorf("%w: %v", ErrUnavailable, err)
|
||||
}
|
||||
allowed, _ := result[0].(int64)
|
||||
current, _ := result[1].(int64)
|
||||
return Reservation{
|
||||
Allowed: allowed == 1, Key: key, Reserved: estimate, Limit: quota,
|
||||
Remaining: max(quota-current, 0), ResetAt: reset,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Reservation 是一次模型配额预留。
|
||||
type Reservation struct {
|
||||
Allowed bool
|
||||
Key string
|
||||
Reserved int64
|
||||
Limit int64
|
||||
Remaining int64
|
||||
ResetAt time.Time
|
||||
}
|
||||
|
||||
// 供 gateway 通过接口断言读取(避免依赖具体类型)。
|
||||
func (r Reservation) AllowedFlag() bool { return r.Allowed }
|
||||
func (r Reservation) RemainingTokens() int64 { return max(r.Remaining, 0) }
|
||||
func (r Reservation) ResetTime() time.Time { return r.ResetAt }
|
||||
|
||||
// Commit 按实际用量回写(与预留的差额)。reservation 为 Reserve 返回值。
|
||||
func (s *Service) Commit(ctx context.Context, reservation any, actual int64) error {
|
||||
res, ok := reservation.(Reservation)
|
||||
if !ok {
|
||||
return errors.New("invalid reservation type")
|
||||
}
|
||||
return s.commitOnce(ctx, res, actual)
|
||||
}
|
||||
|
||||
func (s *Service) commitOnce(ctx context.Context, reservation Reservation, actual int64) error {
|
||||
if s == nil || s.client == nil || reservation.Key == "" || !reservation.Allowed {
|
||||
return nil
|
||||
}
|
||||
if actual < 0 {
|
||||
actual = 0
|
||||
}
|
||||
if _, err := s.commitScript.Run(ctx, s.client, []string{reservation.Key}, actual-reservation.Reserved).Result(); err != nil {
|
||||
return fmt.Errorf("%w: %v", ErrUnavailable, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func monthlyKey(providerCode, model string, now time.Time) string {
|
||||
return "gateway:model-token:" + providerCode + ":" + model + ":" + now.Format("200601")
|
||||
}
|
||||
|
||||
const modelReserveScript = `
|
||||
local estimate = tonumber(ARGV[1])
|
||||
local quota = tonumber(ARGV[2])
|
||||
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
if quota > 0 and current + estimate > quota then
|
||||
return {0, current}
|
||||
end
|
||||
if estimate > 0 then
|
||||
current = redis.call('INCRBY', KEYS[1], estimate)
|
||||
if current == estimate then redis.call('EXPIRE', KEYS[1], tonumber(ARGV[3])) end
|
||||
elseif redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
redis.call('SET', KEYS[1], 0, 'EX', tonumber(ARGV[3]))
|
||||
end
|
||||
return {1, current}
|
||||
`
|
||||
|
||||
const modelCommitScript = `
|
||||
local delta = tonumber(ARGV[1])
|
||||
if redis.call('EXISTS', KEYS[1]) == 0 then
|
||||
return 0
|
||||
end
|
||||
local current = tonumber(redis.call('GET', KEYS[1]) or '0')
|
||||
local updated = current + delta
|
||||
if updated < 0 then updated = 0 end
|
||||
redis.call('SET', KEYS[1], updated, 'KEEPTTL')
|
||||
return updated
|
||||
`
|
||||
|
||||
// --- 管理端 CRUD ---
|
||||
|
||||
// List 返回全部配额记录。
|
||||
func (s *Service) List(ctx context.Context) ([]Quota, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return nil, ErrUnavailable
|
||||
}
|
||||
rows, err := s.pool.Query(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas ORDER BY provider_code,model_pattern`)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rows.Close()
|
||||
items := []Quota{}
|
||||
for rows.Next() {
|
||||
var q Quota
|
||||
if err := rows.Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
items = append(items, q)
|
||||
}
|
||||
return items, rows.Err()
|
||||
}
|
||||
|
||||
// Save 创建或更新一条配额。
|
||||
func (s *Service) Save(ctx context.Context, id, providerCode, modelPattern string, quota int64, enabled bool) (Quota, error) {
|
||||
if s == nil || s.pool == nil {
|
||||
return Quota{}, ErrUnavailable
|
||||
}
|
||||
providerCode = strings.ToLower(strings.TrimSpace(providerCode))
|
||||
modelPattern = strings.TrimSpace(modelPattern)
|
||||
if providerCode == "" || len(providerCode) > 64 || modelPattern == "" || len(modelPattern) > 255 || quota <= 0 {
|
||||
return Quota{}, errors.New("供应商代码、模型模式或配额无效")
|
||||
}
|
||||
if id == "" {
|
||||
newID, err := platformid.NewUUID()
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
id = newID
|
||||
_, err = s.pool.Exec(ctx, `INSERT INTO gateway.model_quotas(id,provider_code,model_pattern,monthly_token_quota,enabled) VALUES($1,$2,$3,$4,$5) ON CONFLICT(provider_code,model_pattern) DO UPDATE SET monthly_token_quota=$4,enabled=$5,updated_at=clock_timestamp()`, id, providerCode, modelPattern, quota, enabled)
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
} else {
|
||||
tag, err := s.pool.Exec(ctx, `UPDATE gateway.model_quotas SET provider_code=$2,model_pattern=$3,monthly_token_quota=$4,enabled=$5,updated_at=clock_timestamp() WHERE id=$1`, id, providerCode, modelPattern, quota, enabled)
|
||||
if err != nil {
|
||||
return Quota{}, err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return Quota{}, errors.New("配额记录不存在")
|
||||
}
|
||||
}
|
||||
var q Quota
|
||||
err := s.pool.QueryRow(ctx, `SELECT id::text,provider_code,model_pattern,monthly_token_quota,enabled,updated_at FROM gateway.model_quotas WHERE id=$1`, id).Scan(&q.ID, &q.ProviderCode, &q.ModelPattern, &q.MonthlyTokenQuota, &q.Enabled, &q.UpdatedAt)
|
||||
return q, err
|
||||
}
|
||||
|
||||
// Delete 删除一条配额。
|
||||
func (s *Service) Delete(ctx context.Context, id string) error {
|
||||
if s == nil || s.pool == nil {
|
||||
return ErrUnavailable
|
||||
}
|
||||
tag, err := s.pool.Exec(ctx, `DELETE FROM gateway.model_quotas WHERE id=$1`, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tag.RowsAffected() == 0 {
|
||||
return errors.New("配额记录不存在")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// String 供日志使用。
|
||||
func (r Reservation) String() string {
|
||||
return strconv.FormatInt(r.Remaining, 10)
|
||||
}
|
||||
Reference in New Issue
Block a user