Kris
0b19470c70
本次提交新增了渠道消息处理的完整核心模块，包含以下核心功能： 1. 新增会话围栏类，实现会话并发控制与过期清理 2. 新增媒体清理器，实现过期媒体文件自动清理 3. 新增熔断器组件，实现服务降级与故障隔离 4. 新增消息处理器，完成渠道消息的完整流转处理 5. 新增限流器组件，实现渠道级和账户级流量控制 6. 新增链路追踪模块，集成Langfuse实现调用链路监控 7. 新增指标统计模块，实现消息处理全链路指标采集 8. 新增统一消息模型，封装全渠道消息格式 9. 新增块回复流水线，实现流式回复的合并与去重 10. 新增本地媒体存储模块，实现媒体文件的本地管理 11. 新增回复分发器，实现回复内容的有序发送与延迟处理 12. 完善__init__.py导出所有核心模块与工具类
218 lines
6.7 KiB
Python
218 lines
6.7 KiB
Python
import asyncio
|
||
import logging
|
||
import time
|
||
from collections.abc import AsyncIterator, Callable
|
||
from dataclasses import dataclass
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
_global_semaphore = asyncio.Semaphore(50)
|
||
|
||
|
||
def get_global_semaphore() -> asyncio.Semaphore:
|
||
return _global_semaphore
|
||
|
||
DEFAULT_RATE = 10.0
|
||
DEFAULT_CAPACITY = 20
|
||
DEFAULT_PER_ACCOUNT_CONCURRENCY = 5
|
||
DEFAULT_QUEUE_DEPTH = 20
|
||
DEFAULT_MAX_LIMITERS = 10_000
|
||
|
||
|
||
@dataclass(slots=True)
|
||
class AcquireResult:
|
||
allowed: bool
|
||
retry_after_sec: float = 0.0
|
||
|
||
def __bool__(self) -> bool:
|
||
return self.allowed
|
||
|
||
|
||
class TokenBucket:
|
||
"""令牌桶算法 — 控制单账户消息速率。
|
||
|
||
参数：
|
||
- rate: 每秒补充的令牌数（默认 10）
|
||
- capacity: 桶容量上限（默认 20，允许短时突发）
|
||
- time_func: 可选时钟注入，便于测试（默认 time.monotonic）
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
rate: float = DEFAULT_RATE,
|
||
capacity: int = DEFAULT_CAPACITY,
|
||
time_func: Callable[[], float] = time.monotonic,
|
||
):
|
||
self._rate = rate
|
||
self._capacity = capacity
|
||
self._tokens = float(capacity)
|
||
self._time = time_func
|
||
self._last_refill = self._time()
|
||
self._lock = asyncio.Lock()
|
||
|
||
async def acquire(self) -> AcquireResult:
|
||
async with self._lock:
|
||
now = self._time()
|
||
elapsed = now - self._last_refill
|
||
self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
|
||
self._last_refill = now
|
||
|
||
if self._tokens >= 1.0:
|
||
self._tokens -= 1.0
|
||
return AcquireResult(allowed=True)
|
||
|
||
if self._rate <= 0.0:
|
||
return AcquireResult(allowed=False, retry_after_sec=float("inf"))
|
||
|
||
retry_after = (1.0 - self._tokens) / self._rate
|
||
return AcquireResult(allowed=False, retry_after_sec=retry_after)
|
||
|
||
def reset(self) -> None:
|
||
self._tokens = float(self._capacity)
|
||
self._last_refill = self._time()
|
||
|
||
|
||
class ChannelRateLimiter:
|
||
"""单渠道账户的限流控制器，组合令牌桶 + 并发信号量。
|
||
|
||
线程安全：所有方法在 asyncio 上下文中调用。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
rate: float = DEFAULT_RATE,
|
||
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
|
||
queue_depth: int = DEFAULT_QUEUE_DEPTH,
|
||
):
|
||
self._rate = rate
|
||
self._per_account_concurrency = per_account_concurrency
|
||
self._bucket = TokenBucket(rate=rate, capacity=DEFAULT_CAPACITY)
|
||
self._semaphore = asyncio.Semaphore(per_account_concurrency)
|
||
self._queue_depth = queue_depth
|
||
self._waiting: int = 0
|
||
|
||
async def try_acquire(self) -> AcquireResult:
|
||
"""尝试获取消息处理许可。
|
||
|
||
返回 AcquireResult(allowed=..., retry_after_sec=...)。
|
||
"""
|
||
result = await self._bucket.acquire()
|
||
if result.allowed:
|
||
return result
|
||
|
||
if self._waiting >= self._queue_depth:
|
||
logger.warning(
|
||
"Rate limit queue full (depth=%d), rejecting message",
|
||
self._queue_depth,
|
||
)
|
||
return AcquireResult(allowed=False, retry_after_sec=result.retry_after_sec)
|
||
|
||
self._waiting += 1
|
||
try:
|
||
await asyncio.sleep(result.retry_after_sec)
|
||
retry = await self._bucket.acquire()
|
||
if not retry.allowed:
|
||
return AcquireResult(allowed=False, retry_after_sec=retry.retry_after_sec)
|
||
return retry
|
||
finally:
|
||
self._waiting -= 1
|
||
|
||
async def run_with_limit(self, coro) -> AsyncIterator:
|
||
"""在并发限制内执行 Agent 调用。
|
||
|
||
先获取账户级信号量，再获取全局信号量。
|
||
"""
|
||
async with self._semaphore:
|
||
async with _global_semaphore:
|
||
async for chunk in coro:
|
||
yield chunk
|
||
|
||
def reset(self) -> None:
|
||
self._bucket.reset()
|
||
|
||
@property
|
||
def waiting_count(self) -> int:
|
||
return self._waiting
|
||
|
||
@property
|
||
def config(self) -> dict:
|
||
return {
|
||
"rate": self._rate,
|
||
"per_account_concurrency": self._per_account_concurrency,
|
||
"queue_depth": self._queue_depth,
|
||
}
|
||
|
||
|
||
class RateLimitManager:
|
||
"""限流管理器 — 按账户维度管理限流器实例。
|
||
|
||
每个 (channel_type, account_id) 组合共享同一个 ChannelRateLimiter。
|
||
内置自动清理机制防止无限增长。
|
||
"""
|
||
|
||
def __init__(self, max_limiters: int = DEFAULT_MAX_LIMITERS):
|
||
self._limiters: dict[str, tuple[float, ChannelRateLimiter]] = {}
|
||
self._max_limiters = max_limiters
|
||
self._last_cleanup = time.monotonic()
|
||
self._lock = asyncio.Lock()
|
||
|
||
def _key(self, channel_type: str, account_id: str) -> str:
|
||
return f"{channel_type}:{account_id}"
|
||
|
||
async def get_limiter(
|
||
self,
|
||
channel_type: str,
|
||
account_id: str,
|
||
rate: float = DEFAULT_RATE,
|
||
per_account_concurrency: int = DEFAULT_PER_ACCOUNT_CONCURRENCY,
|
||
) -> ChannelRateLimiter:
|
||
async with self._lock:
|
||
self._maybe_cleanup()
|
||
key = self._key(channel_type, account_id)
|
||
if key not in self._limiters:
|
||
self._limiters[key] = (
|
||
time.monotonic(),
|
||
ChannelRateLimiter(
|
||
rate=rate,
|
||
per_account_concurrency=per_account_concurrency,
|
||
),
|
||
)
|
||
else:
|
||
self._limiters[key] = (time.monotonic(), self._limiters[key][1])
|
||
return self._limiters[key][1]
|
||
|
||
async def remove_limiter(self, channel_type: str, account_id: str) -> None:
|
||
async with self._lock:
|
||
key = self._key(channel_type, account_id)
|
||
self._limiters.pop(key, None)
|
||
|
||
async def clear(self) -> None:
|
||
async with self._lock:
|
||
self._limiters.clear()
|
||
|
||
async def get_stats(self) -> dict[str, dict]:
|
||
async with self._lock:
|
||
return {
|
||
key: {
|
||
"waiting": limiter.waiting_count,
|
||
**limiter.config,
|
||
}
|
||
for key, (_, limiter) in self._limiters.items()
|
||
}
|
||
|
||
def _maybe_cleanup(self) -> None:
|
||
if len(self._limiters) <= self._max_limiters:
|
||
return
|
||
now = time.monotonic()
|
||
if now - self._last_cleanup < 300:
|
||
return
|
||
self._last_cleanup = now
|
||
expired = [key for key, (last_touch, _) in self._limiters.items() if now - last_touch > 3600]
|
||
for key in expired:
|
||
self._limiters.pop(key, None)
|
||
if expired:
|
||
logger.info("Cleaned up %d stale rate limiters", len(expired))
|
||
|
||
|
||
rate_limit_manager = RateLimitManager()
|