Kris
552aef767c
新增QQ Bot适配器完整代码栈，包含： 1. 基础适配器入口与工具类封装 2. 会话管理、重试队列与流量控制 3. 命令系统与内置指令（ping/help/status等） 4. 富媒体消息处理与格式转换 5. 引用存储与审批管理 6. 凭证备份与会话持久化 7. 健康检查与交互回调系统
88 lines
2.9 KiB
Python
88 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class TokenBucket:
|
|
rate: float
|
|
burst: int
|
|
|
|
_tokens: float = field(default=0, init=False)
|
|
_last_refill: float = field(default_factory=time.monotonic, init=False)
|
|
_lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
|
|
|
|
def __post_init__(self):
|
|
self._tokens = float(self.burst)
|
|
|
|
def _refill(self) -> None:
|
|
now = time.monotonic()
|
|
elapsed = now - self._last_refill
|
|
self._tokens = min(self._tokens + elapsed * self.rate, float(self.burst))
|
|
self._last_refill = now
|
|
|
|
async def acquire(self, tokens: float = 1.0) -> None:
|
|
while True:
|
|
async with self._lock:
|
|
self._refill()
|
|
if self._tokens >= tokens:
|
|
self._tokens -= tokens
|
|
return
|
|
|
|
wait = (tokens - self._tokens) / self.rate if self.rate > 0 else 0.1
|
|
await asyncio.sleep(max(wait, 0.01))
|
|
|
|
def try_acquire(self, tokens: float = 1.0) -> bool:
|
|
self._refill()
|
|
if self._tokens >= tokens:
|
|
self._tokens -= tokens
|
|
return True
|
|
return False
|
|
|
|
|
|
@dataclass
|
|
class RouteRateLimiter:
|
|
defaults: dict[str, tuple[float, int]] = field(
|
|
default_factory=lambda: {
|
|
"send_message": (5.0, 10),
|
|
"send_media": (1.0, 3),
|
|
"upload_media": (0.5, 2),
|
|
"default": (10.0, 20),
|
|
}
|
|
)
|
|
|
|
_buckets: dict[str, TokenBucket] = field(default_factory=dict, init=False)
|
|
_lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False)
|
|
|
|
async def acquire(self, route: str = "default", tokens: float = 1.0) -> None:
|
|
bucket = await self._get_or_create_bucket(route)
|
|
await bucket.acquire(tokens)
|
|
|
|
def try_acquire(self, route: str = "default", tokens: float = 1.0) -> bool:
|
|
bucket = self._get_or_create_bucket_sync(route)
|
|
return bucket.try_acquire(tokens)
|
|
|
|
async def _get_or_create_bucket(self, route: str) -> TokenBucket:
|
|
async with self._lock:
|
|
if route not in self._buckets:
|
|
rate, burst = self.defaults.get(route, self.defaults["default"])
|
|
self._buckets[route] = TokenBucket(rate=rate, burst=burst)
|
|
return self._buckets[route]
|
|
|
|
def _get_or_create_bucket_sync(self, route: str) -> TokenBucket:
|
|
if route not in self._buckets:
|
|
rate, burst = self.defaults.get(route, self.defaults["default"])
|
|
self._buckets[route] = TokenBucket(rate=rate, burst=burst)
|
|
return self._buckets[route]
|
|
|
|
def get_stats(self) -> dict[str, dict]:
|
|
return {
|
|
route: {"tokens": bucket._tokens, "rate": bucket.rate, "burst": bucket.burst}
|
|
for route, bucket in self._buckets.items()
|
|
}
|