From b2b7a64cb3a8f10f70eed946dcf2fe14c7141e62 Mon Sep 17 00:00:00 2001 From: tzt <14718231+flying-travel@user.noreply.gitee.com> Date: Sat, 19 Sep 2026 10:09:06 +0800 Subject: [PATCH] =?UTF-8?q?feat(proxy):=20T-X12=20=E9=87=87=E7=BA=B3=20ent?= =?UTF-8?q?erprise-ai-gateway=20=E5=8F=8C=E6=A1=B6=E9=99=90=E6=B5=81?= =?UTF-8?q?=E2=80=94=E2=80=94TPM=20token=20=E6=A1=B6=20+=20=E8=AF=AD?= =?UTF-8?q?=E4=B9=89=E7=BC=93=E5=AD=98=E5=91=BD=E4=B8=AD=E7=8E=87=E9=80=8F?= =?UTF-8?q?=E5=87=BA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - auth.RateLimiter.allow_tokens：TPM 令牌桶（容量=tpm_cap，连续回流 tpm/60 每秒， 与 RPM 桶共用 per-key 锁、判定相互独立；tpm_cap<=0 不限） - ProxyConfig.limits.tpm_per_key（默认 60000，0=不限）；chat_completions 在并发槽 之后做 TPM 预扣判定（est=字符/3 与计费同口径），拒绝时释放并发槽并回 429 （文案区分并发/RPM/TPM 三种原因） - semcache.stats 增 misses 与 hit_rate；/admin/stats 新增 semcache 段透出 （entries/hits_exact/hits_semantic/misses/hit_rate，与账本 h_g 互补） - 新增 tests/test_tpm_limiter.py 6 项（容量/回流/禁用/双桶独立/key 隔离/命中率） --- gateway/proxy/auth.py | 22 ++++++++++++ gateway/proxy/config.py | 2 ++ gateway/proxy/routes.py | 16 +++++++++ gateway/proxy/semcache.py | 8 ++++- tests/test_tpm_limiter.py | 70 +++++++++++++++++++++++++++++++++++++++ 5 files changed, 117 insertions(+), 1 deletion(-) create mode 100644 tests/test_tpm_limiter.py diff --git a/gateway/proxy/auth.py b/gateway/proxy/auth.py index 5d58013..316ef2f 100644 --- a/gateway/proxy/auth.py +++ b/gateway/proxy/auth.py @@ -89,6 +89,7 @@ class RateLimiter: def __init__(self, concurrent_per_key: int = 2): self._tokens: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts) + self._tpm: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts) self._locks: Dict[int, threading.Lock] = {} self._sems: Dict[int, threading.Semaphore] = {} self._global = threading.Lock() @@ -108,6 +109,27 @@ class RateLimiter: self._tokens[key_id] = (tokens - 1.0, now) return True + def allow_tokens(self, key_id: int, est_tokens: int, tpm_cap: int) -> bool: + """TPM 令牌桶放行判定（T-X12，采纳 enterprise-ai-gateway 双桶设计）。 + + 容量 = tpm_cap，回流速率 = tpm/60 每秒（连续时间回流补充）； + est_tokens 为本次预扣估算（字符/3，与预扣计费同口径）。 + tpm_cap <= 0 视为不限流；与 RPM 桶共用 per-key 锁，判定相互独立。 + """ + if tpm_cap <= 0: + return True + now = time.monotonic() + with self._global: + lock = self._locks.setdefault(key_id, threading.Lock()) + with lock: + tokens, last = self._tpm.get(key_id, (float(tpm_cap), now)) + tokens = min(float(tpm_cap), tokens + (now - last) * (tpm_cap / 60.0)) + if tokens < est_tokens: + self._tpm[key_id] = (tokens, now) + return False + self._tpm[key_id] = (tokens - est_tokens, now) + return True + def acquire_slot(self, key_id: int) -> bool: """并发槽（非阻塞）；返回 False = 超并发上限（429）。""" with self._global: diff --git a/gateway/proxy/config.py b/gateway/proxy/config.py index c5c2f6b..8e5117f 100644 --- a/gateway/proxy/config.py +++ b/gateway/proxy/config.py @@ -52,6 +52,7 @@ class ProxyConfig: rpm_per_key: int = 10 day_req_cap: int = 200 concurrent_per_key: int = 2 + tpm_per_key: int = 60000 # T-X12：每分钟 token 桶（0 = 不限） max_body_chars: int = 60000 semcache_enabled: bool = True sim_threshold: float = 0.92 @@ -159,6 +160,7 @@ def build_proxy_config(settings_dict: Dict[str, Any]) -> ProxyConfig: rpm_per_key=int(limits.get("rpm_per_key", 10) or 10), day_req_cap=int(limits.get("day_req_cap", 200) or 200), concurrent_per_key=int(limits.get("concurrent_per_key", 2) or 2), + tpm_per_key=int(limits.get("tpm_per_key", 60000) or 0), max_body_chars=int(limits.get("max_body_chars", 60000) or 60000), semcache_enabled=bool(sem.get("enabled", True)), sim_threshold=float(sem.get("sim_threshold", 0.92) or 0.92), diff --git a/gateway/proxy/routes.py b/gateway/proxy/routes.py index 00af69a..98a4ec5 100644 --- a/gateway/proxy/routes.py +++ b/gateway/proxy/routes.py @@ -86,6 +86,14 @@ def build_proxy_router(cfg: ProxyConfig, pool, settings_provider=None, _limits(), time.time()) if not _limits().acquire_slot(ctx["key_id"]): raise QuotaError("并发请求已达该 key 上限") + # T-X12（采纳 enterprise-ai-gateway 双桶限流）：TPM 桶——RPM 之外加 + # token 维度（预扣估算 = 字符/3，与计费同口径；cfg.tpm_per_key=0 不限） + tpm_cap = int(getattr(cfg, "tpm_per_key", 0) or 0) + if tpm_cap > 0: + est_tokens = max(1, len(raw_body) // 3) + if not _limits().allow_tokens(ctx["key_id"], est_tokens, tpm_cap): + _limits().release_slot(ctx["key_id"]) + raise QuotaError("该 key 的每分钟 token 额度（TPM）已用尽") # 语义分析器 live 分流（D-G7：mode=live 才启用；任何异常不影响代理可用性）。 # T-X4 修复：settings_provider 由 api.py 注入（原实现引用未导入的 @@ -217,10 +225,18 @@ def build_proxy_router(cfg: ProxyConfig, pool, settings_provider=None, in_hit = s["in_hit_tok"] in_miss = s["in_miss_tok"] h_p = in_hit / (in_hit + in_miss) if (in_hit + in_miss) else 0.0 + # T-X12：进程内语义缓存计数透出（entries/hits/misses/hit_rate） + sem_stats = None + if cfg.semcache_enabled: + try: + sem_stats = _get_semcache(cfg, ledger).stats() + except Exception: # noqa: BLE001 + sem_stats = None return {"requests": n, "h_g": round(h_g, 4), "h_p": round(h_p, 4), "revenue_milli": s["revenue_milli"], "cost_milli": s["cost_milli"], "margin_milli": s["revenue_milli"] - s["cost_milli"], "by_bucket": s["by_bucket"], + "semcache": sem_stats, "today": today, "upstream_failover": upstream_failover_stats()} diff --git a/gateway/proxy/semcache.py b/gateway/proxy/semcache.py index d1688ee..e70e3a6 100644 --- a/gateway/proxy/semcache.py +++ b/gateway/proxy/semcache.py @@ -104,6 +104,7 @@ class SemanticCache: self._clock = now # 假时钟注入（None = time.time） self.hits_exact = 0 self.hits_semantic = 0 + self.misses = 0 self._rebuild() # ---------- 时钟 ---------- @@ -166,6 +167,7 @@ class SemanticCache: # L2 g = grams(norm_text) if len(g) < 3: + self.misses += 1 return None w_q = _weight(g) candidates: Dict[str, int] = {} @@ -195,6 +197,7 @@ class SemanticCache: best_score = score best_key = key if best_key is None or best_score < self.sim_threshold: + self.misses += 1 return None cand = self._l1[best_key] cand.hits += 1 @@ -244,8 +247,11 @@ class SemanticCache: pass # 持久化失败不影响内存缓存（可重建） def stats(self) -> Dict[str, Any]: + total = self.hits_exact + self.hits_semantic + self.misses return {"entries": len(self._l1), "hits_exact": self.hits_exact, - "hits_semantic": self.hits_semantic} + "hits_semantic": self.hits_semantic, "misses": self.misses, + "hit_rate": round((self.hits_exact + self.hits_semantic) / total, 4) + if total else 0.0} class SingleFlight: diff --git a/tests/test_tpm_limiter.py b/tests/test_tpm_limiter.py new file mode 100644 index 0000000..dfd6b69 --- /dev/null +++ b/tests/test_tpm_limiter.py @@ -0,0 +1,70 @@ +"""TPM 双桶限流与缓存命中率透出（T-X12，采纳 enterprise-ai-gateway 设计）。""" +import time + +from gateway.proxy.auth import RateLimiter +from gateway.proxy.semcache import SemanticCache + + +def test_tpm_bucket_capacity_and_consume(): + """容量内放行并扣减；超过余量拒绝。""" + rl = RateLimiter() + assert rl.allow_tokens(1, 100, 1000) is True + assert rl.allow_tokens(1, 500, 1000) is True + # 已扣 600，剩 400 + 回流（瞬时近似 0）：700 > 400+eps -> 拒 + assert rl.allow_tokens(1, 700, 1000) is False + + +def test_tpm_bucket_continuous_refill(monkeypatch): + """连续时间回流：elapsed * tpm/60 补充（容量封顶）。""" + rl = RateLimiter() + clock = {"t": 100.0} + monkeypatch.setattr(time, "monotonic", lambda: clock["t"]) + assert rl.allow_tokens(2, 1000, 2000) is True # 剩 1000 + clock["t"] = 130.0 # 30s 回流 1000 -> 满 2000 + assert rl.allow_tokens(2, 1500, 2000) is True # 扣 1500 剩 500 + clock["t"] = 145.0 # 15s 回流 500 -> 1000 + assert rl.allow_tokens(2, 1000, 2000) is True # 精确扣空 + assert rl.allow_tokens(2, 1, 2000) is False # 空桶拒绝 + + +def test_tpm_disabled_when_cap_zero(): + """tpm_cap <= 0 = 不限流（恒放行）。""" + rl = RateLimiter() + assert rl.allow_tokens(3, 10**9, 0) is True + + +def test_tpm_independent_from_rpm_bucket(): + """双桶相互独立：打满 TPM 不影响 RPM 放行，RPM 打空不影响 TPM 记账。""" + rl = RateLimiter() + assert rl.allow_tokens(4, 5000, 5000) is True # TPM 打空 + assert rl.allow(4, 10) is True # RPM 仍放行 + for _ in range(9): # 共 10 次 = 打满 RPM + assert rl.allow(4, 10) is True + assert rl.allow(4, 10) is False # RPM 空桶 + assert rl.allow_tokens(4, 1, 5000) is False # TPM 仍空（未被 RPM 动过） + + +def test_tpm_per_key_isolation(): + """按 key 隔离：一个 key 打满不影响另一个。""" + rl = RateLimiter() + assert rl.allow_tokens(5, 5000, 5000) is True + assert rl.allow_tokens(6, 5000, 5000) is True + assert rl.allow_tokens(5, 5000, 5000) is False + + +def test_semcache_stats_hit_rate(): + """stats 透出 misses 与 hit_rate（命中率口径 = hits / (hits + misses)）。""" + class _S: + def semcache_rows(self, limit=0): + return [] + + def put_semcache(self, *a, **k): + pass + + c = SemanticCache(_S(), max_entries=10, sim_threshold=0.5, promote_frequency=5) + c.put("k1", "请解释一下递归函数的概念", "答案", "m", doc_version=1) + assert c.lookup("k1", "请解释一下递归函数的概念", doc_version=1) is not None + assert c.lookup("k9", "完全无关的问题内容呢", doc_version=1) is None + s = c.stats() + assert s["hits_exact"] == 1 and s["misses"] == 1 + assert s["hit_rate"] == 0.5