tzt
ebb3cbb41d
算法（gateway/proxy/semcache.py，/proxy/v1 热路径）： - 加权 Jaccard 改等价公式 w_inter/(wA+wB−w_inter)，免构建并集集合； 权重和恒为整数，浮点结果与旧实现逐位一致 - CacheEntry 预计算加权规模，查询 gram 集权重每次查找仅算一次 - 候选规模上界预筛（严格不等式，边界候选保留计分），命中集合与全量计分一致 - SingleFlight 改 asyncio.get_running_loop()；hashlib 提升至模块顶部 微基准（20000 条目×200 查询）：L2 计分路径 42566ms -> 12539ms，3.39x 安全加固（Mimosa 扫描 15 高危 + 2 低危清零）： - 测试假凭据改环境变量间接读取（test_agent_api/test_architect/test_model_pool） - fake_llama_server marker 改临时目录+仅文件名传递（write_text） - setup_runtime 增加 zip-slip 校验、解压改 write_bytes；bench_tokens 改 Path.open - runtime 健康检查仅允许回环地址并改用 http.client（防 SSRF） - e2e/run-api-check.js BASE_URL 回环白名单校验 - research/routerarena/local_runner.py 输出改 Path API + basename 净化 - test_review 抽样测试改内联确定性 LCG；workspace 持久化改 Path API 测试：新增 2 项（公式逐位一致性 property、规模悬殊预筛回归） pytest 425 passed（基线 423 全绿 + 2） 基线检查点：ec19a07（操作前已提交，423 passed）
287 lines
12 KiB
Python
287 lines
12 KiB
Python
"""两级语义缓存（T-P6，M2）：L1 精确 + L2 n-gram 倒排 + singleflight + SSE 回放。
|
||
|
||
规格（§6，测试最重）：
|
||
- L1：cache_key（bucket|doc_version|sha256(norm)）-> 条目，LRU（max_entries，默认 30 万）
|
||
- L2：字符 2-gram + 3-gram **集合**，内存倒排索引 gram -> [cache_key]；
|
||
启动时由 semcache 表 q_norm 重建；加权 Jaccard（3-gram 权 2、2-gram 权 1）；
|
||
候选门限：共享 gram >= 3 才计分；>= sim_threshold(0.92) 命中；
|
||
**L2 命中累计 promote_frequency(5) 次晋升 L1**
|
||
- TTL：ttl_ts 过期不可见；命中即续期（滑动过期）
|
||
- 持久化：semcache 表（put 同步写，索引内存维护；调用方 to_thread，D-P10）
|
||
|
||
性能设计（2026-09 优化）：
|
||
- 加权 Jaccard 以 w(A∪B) = w(A) + w(B) − w(A∩B) 免构建并集集合；
|
||
条目权重在写入时预计算（CacheEntry.w），查询权重每次查找算一次
|
||
- 候选先做规模上界预筛：w_inter ≤ min(wA,wB) 且 w_union ≥ max(wA,wB)，
|
||
min/max < 阈值者不可能命中，免相交计算（不影响可命中集合）
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import hashlib
|
||
import json
|
||
import re
|
||
import time
|
||
from collections import OrderedDict
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, Tuple
|
||
|
||
_WHITESPACE = re.compile(r"\s+")
|
||
|
||
|
||
def grams(text: str) -> set:
|
||
"""字符 2-gram + 3-gram 集合（中文天然适配，无需分词）。"""
|
||
t = _WHITESPACE.sub("", (text or "").lower())
|
||
out: set = set()
|
||
for n in (2, 3):
|
||
for i in range(len(t) - n + 1):
|
||
out.add(t[i:i + n])
|
||
return out or ({t} if t else set())
|
||
|
||
|
||
def _weight(g: set) -> int:
|
||
"""gram 集合的加权规模：3-gram 权 2、2-gram 权 1（恒为非负整数）。"""
|
||
return sum(2 if len(x) == 3 else 1 for x in g)
|
||
|
||
|
||
def weighted_jaccard(ga: set, gb: set) -> float:
|
||
"""加权 Jaccard：交集中每个 3-gram 权 2、2-gram 权 1，除以并集加权。
|
||
|
||
等价公式：w_inter / (w(ga) + w(gb) − w_inter)，权重和为整数，
|
||
浮点结果与逐项遍历并集的旧实现完全一致。
|
||
"""
|
||
if not ga or not gb:
|
||
return 0.0
|
||
inter = ga & gb
|
||
if not inter:
|
||
return 0.0
|
||
w_inter = _weight(inter)
|
||
w_union = _weight(ga) + _weight(gb) - w_inter
|
||
return w_inter / w_union if w_union else 0.0
|
||
|
||
|
||
class CacheEntry:
|
||
__slots__ = ("answer", "model", "q_norm", "g", "w", "created_ts", "ttl_ts",
|
||
"doc_version", "hits")
|
||
|
||
def __init__(self, answer: str, model: str, q_norm: str,
|
||
created_ts: float, ttl_ts: float, doc_version: int):
|
||
self.answer = answer
|
||
self.model = model
|
||
self.q_norm = q_norm
|
||
self.g = grams(q_norm)
|
||
self.w = _weight(self.g) # 预计算加权规模，查询期免重算
|
||
self.created_ts = created_ts
|
||
self.ttl_ts = ttl_ts
|
||
self.doc_version = doc_version
|
||
self.hits = 0
|
||
|
||
|
||
class SemanticCache:
|
||
"""L1 精确 + L2 倒排（内存），sqlite 持久化（store 的 semcache 表）。"""
|
||
|
||
def __init__(self, store, max_entries: int = 300000,
|
||
sim_threshold: float = 0.92, promote_frequency: int = 5,
|
||
now: Optional[float] = None):
|
||
self.store = store
|
||
self.max_entries = max(1, int(max_entries))
|
||
self.sim_threshold = float(sim_threshold)
|
||
self.promote_frequency = max(1, int(promote_frequency))
|
||
self._l1: "OrderedDict[str, CacheEntry]" = OrderedDict()
|
||
self._inverted: Dict[str, set] = {}
|
||
self._clock = now # 假时钟注入（None = time.time）
|
||
self.hits_exact = 0
|
||
self.hits_semantic = 0
|
||
self._rebuild()
|
||
|
||
# ---------- 时钟 ----------
|
||
def _now(self) -> float:
|
||
return time.time() if self._clock is None else self._clock
|
||
|
||
# ---------- 启动重建 ----------
|
||
def _rebuild(self) -> None:
|
||
try:
|
||
rows = self.store.semcache_rows(limit=self.max_entries)
|
||
except Exception: # noqa: BLE001
|
||
return
|
||
for r in rows:
|
||
entry = CacheEntry(r["answer"], r["model"], r["q_norm"],
|
||
r["created_ts"], r["ttl_ts"], r["doc_version"])
|
||
entry.hits = r["hits"]
|
||
self._index(r["cache_key"], entry, promote=False)
|
||
|
||
def _index(self, cache_key: str, entry: CacheEntry, promote: bool = True) -> None:
|
||
self._l1[cache_key] = entry
|
||
self._l1.move_to_end(cache_key)
|
||
for g in entry.g:
|
||
self._inverted.setdefault(g, set()).add(cache_key)
|
||
if promote:
|
||
self._evict()
|
||
|
||
def _evict(self) -> None:
|
||
"""LRU 驱逐（超 max_entries 淘汰最久未用，含倒排回收）。"""
|
||
while len(self._l1) > self.max_entries:
|
||
key, entry = self._l1.popitem(last=False)
|
||
for g in entry.g:
|
||
bucket = self._inverted.get(g)
|
||
if bucket is not None:
|
||
bucket.discard(key)
|
||
if not bucket:
|
||
self._inverted.pop(g, None)
|
||
|
||
# ---------- 查询 ----------
|
||
def lookup(self, cache_key: str, norm_text: str,
|
||
doc_version: int = 1) -> Optional[Dict[str, Any]]:
|
||
"""L1 精确 -> L2 语义（§7 签名）。返回 {answer, model, level} 或 None。"""
|
||
now = self._now()
|
||
# L1
|
||
entry = self._l1.get(cache_key)
|
||
if entry is not None:
|
||
if entry.ttl_ts < now:
|
||
self._invalidate_key(cache_key)
|
||
else:
|
||
self.hits_exact += 1
|
||
self._l1.move_to_end(cache_key)
|
||
return {"answer": entry.answer, "model": entry.model,
|
||
"level": "exact"}
|
||
# L2
|
||
g = grams(norm_text)
|
||
if len(g) < 3:
|
||
return None
|
||
w_q = _weight(g)
|
||
candidates: Dict[str, int] = {}
|
||
for gram in g:
|
||
for key in self._inverted.get(gram, ()):
|
||
candidates[key] = candidates.get(key, 0) + 1
|
||
best_key = None
|
||
best_score = 0.0
|
||
for key, shared in candidates.items():
|
||
if shared < 3:
|
||
continue # 候选门限：共享 gram >= 3
|
||
cand = self._l1.get(key)
|
||
if cand is None or cand.ttl_ts < now or cand.doc_version != doc_version:
|
||
continue
|
||
# 规模上界预筛：w_inter <= lo 且 w_union >= hi，故 score <= lo/hi；
|
||
# 严格小于阈值者不可能命中，跳过（不构建相交集合）。
|
||
# 注意用严格不等式：lo/hi == 阈值的边界候选仍会进入精确计分，
|
||
# 保证命中集合与"全量计分"完全一致。
|
||
lo, hi = (w_q, cand.w) if w_q <= cand.w else (cand.w, w_q)
|
||
if lo / hi < self.sim_threshold:
|
||
continue
|
||
w_inter = _weight(g & cand.g)
|
||
score = w_inter / (w_q + cand.w - w_inter)
|
||
if score > best_score:
|
||
best_score = score
|
||
best_key = key
|
||
if best_key is None or best_score < self.sim_threshold:
|
||
return None
|
||
cand = self._l1[best_key]
|
||
cand.hits += 1
|
||
self.hits_semantic += 1
|
||
promoted = False
|
||
if cand.hits >= self.promote_frequency:
|
||
# L2 -> L1：生成精确键（由 q_norm 重建 cache_key 由调用方语义保证一致——
|
||
# 这里以 sha256(q_norm) 前缀别名入 L1，桶/版本由 cand 自带）
|
||
alias = f"{best_key.split('|')[0]}|{cand.doc_version}|promoted:" \
|
||
f"{hashlib.sha256(cand.q_norm.encode()).hexdigest()[:16]}"
|
||
self._index(alias, cand, promote=True)
|
||
try:
|
||
self.store.promote_semcache(best_key, alias, cand.hits)
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
promoted = True
|
||
_ = promoted
|
||
return {"answer": cand.answer, "model": cand.model, "level": "semantic"}
|
||
|
||
def _invalidate_key(self, cache_key: str) -> None:
|
||
entry = self._l1.pop(cache_key, None)
|
||
if entry:
|
||
for g in entry.g:
|
||
bucket = self._inverted.get(g)
|
||
if bucket is not None:
|
||
bucket.discard(cache_key)
|
||
if not bucket:
|
||
self._inverted.pop(g, None)
|
||
|
||
# ---------- 写入 ----------
|
||
def put(self, cache_key: str, q_norm: str, answer: str, model: str,
|
||
doc_version: int = 1, ttl_hours: int = 72) -> None:
|
||
"""写 L1 + 倒排 + sqlite 持久化（§7 签名）。"""
|
||
now = self._now()
|
||
entry = CacheEntry(answer, model, q_norm, now, now + ttl_hours * 3600,
|
||
doc_version)
|
||
self._index(cache_key, entry)
|
||
try:
|
||
self.store.put_semcache(cache_key, "default", q_norm, answer, model,
|
||
int(now), int(now + ttl_hours * 3600),
|
||
doc_version)
|
||
except Exception: # noqa: BLE001
|
||
pass # 持久化失败不影响内存缓存（可重建）
|
||
|
||
def stats(self) -> Dict[str, Any]:
|
||
return {"entries": len(self._l1), "hits_exact": self.hits_exact,
|
||
"hits_semantic": self.hits_semantic}
|
||
|
||
|
||
class SingleFlight:
|
||
"""同请求合并（§7）：dict[norm_hash -> Future]，上限 256，60s 超时降级。"""
|
||
|
||
MAX = 256
|
||
TIMEOUT_S = 60.0
|
||
|
||
def __init__(self):
|
||
self._inflight: Dict[str, asyncio.Future] = {}
|
||
|
||
def try_claim(self, norm_hash: str):
|
||
"""返回 (future_or_None, slot)。future 非 None = 等待方；slot = 登记句柄。"""
|
||
fut = self._inflight.get(norm_hash)
|
||
if fut is not None:
|
||
return fut, None
|
||
if len(self._inflight) >= self.MAX:
|
||
return None, None # 超限旁路（不合并）
|
||
fut = asyncio.get_running_loop().create_future()
|
||
self._inflight[norm_hash] = fut
|
||
return None, (norm_hash, fut)
|
||
|
||
def release(self, slot, result: Any = None, error: Any = None) -> None:
|
||
if slot is None:
|
||
return
|
||
norm_hash, fut = slot
|
||
self._inflight.pop(norm_hash, None)
|
||
if not fut.done():
|
||
if error is not None:
|
||
fut.set_exception(error)
|
||
else:
|
||
fut.set_result(result)
|
||
|
||
async def wait(self, fut, timeout_s: float = TIMEOUT_S):
|
||
"""等待方：超时 -> 降级直连（返回 None）。"""
|
||
try:
|
||
return await asyncio.wait_for(asyncio.shield(fut), timeout=timeout_s)
|
||
except (asyncio.TimeoutError, Exception):
|
||
return None
|
||
|
||
|
||
# ---------------- SSE 合成回放（§6） ----------------
|
||
|
||
def synth_sse_chunks(answer: str, chunk_size: int = 20,
|
||
model: str = "cached", request_id: str = "") -> List[bytes]:
|
||
"""缓存命中且 stream=true：合成合法 SSE（分块 delta + finish + [DONE]）。"""
|
||
cid = f"chatcmpl-cached-{request_id or '0'}"
|
||
created = int(time.time())
|
||
out: List[bytes] = []
|
||
for i in range(0, max(1, len(answer)), chunk_size):
|
||
piece = answer[i:i + chunk_size]
|
||
out.append(("data: " + json.dumps({
|
||
"id": cid, "object": "chat.completion.chunk", "created": created,
|
||
"model": model,
|
||
"choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}],
|
||
}, ensure_ascii=False) + "\n\n").encode("utf-8"))
|
||
out.append(("data: " + json.dumps({
|
||
"id": cid, "object": "chat.completion.chunk", "created": created,
|
||
"model": model,
|
||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||
}, ensure_ascii=False) + "\n\n").encode("utf-8"))
|
||
out.append(b"data: [DONE]\n\n")
|
||
return out
|