Agent限流与熔断：从令牌桶到自适应限流

传统Web应用流量              Agent系统流量
     │                            │
     │   ╱╲                      │   ╱╲
     │  ╱  ╲    ╱╲              │  ╱  ╲╱╲╱╲
     │ ╱    ╲__╱  ╲___          │ ╱        ╲___
     │________________           │________________
     时间                        时间
     
     相对均匀的请求量              突发性强、长尾明显
     每请求资源消耗相近           单请求资源消耗差异巨大（10-100x）

import asyncio
import time
from dataclasses import dataclass

@dataclass
class TokenBucket:
    """令牌桶限流器"""
    
    capacity: float        # 桶容量
    refill_rate: float     # 每秒补充令牌数
    tokens: float = 0      # 当前令牌数
    last_refill: float = 0 # 上次补充时间
    
    def __post_init__(self):
        self.tokens = self.capacity
        self.last_refill = time.monotonic()
    
    async def acquire(self, tokens: int = 1) -> bool:
        """获取令牌"""
        while True:
            self._refill()
            
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            
            # 令牌不足，等待补充
            wait_time = (tokens - self.tokens) / self.refill_rate
            await asyncio.sleep(min(wait_time, 0.1))
    
    def _refill(self):
        """补充令牌"""
        now = time.monotonic()
        elapsed = now - self.last_refill
        self.tokens = min(
            self.capacity,
            self.tokens + elapsed * self.refill_rate
        )
        self.last_refill = now
    
    def try_acquire(self, tokens: int = 1) -> bool:
        """非阻塞获取令牌"""
        self._refill()
        if self.tokens >= tokens:
            self.tokens -= tokens
            return True
        return False

class AgentRateLimiter:
    """Agent多维度限流器"""
    
    def __init__(self):
        # 不同维度的限流器
        self.limiters = {
            "requests_per_minute": TokenBucket(capacity=100, refill_rate=100/60),
            "tokens_per_minute": TokenBucket(capacity=50000, refill_rate=50000/60),
            "tool_calls_per_minute": TokenBucket(capacity=50, refill_rate=50/60),
            "concurrent_sessions": SemaphoreLimiter(max_concurrent=20),
            "gpu_concurrent": SemaphoreLimiter(max_concurrent=4),
        }
    
    async def check_request(self, request: dict) -> dict:
        """检查请求是否被允许"""
        required = {
            "requests_per_minute": 1,
            "concurrent_sessions": 1,
        }
        
        # 根据请求类型添加额外限制
        if request.get("estimated_tokens"):
            required["tokens_per_minute"] = request["estimated_tokens"]
        
        if request.get("tool_calls"):
            required["tool_calls_per_minute"] = len(request["tool_calls"])
        
        if request.get("requires_gpu"):
            required["gpu_concurrent"] = 1
        
        # 检查所有维度
        for dim, amount in required.items():
            limiter = self.limiters[dim]
            if not limiter.try_acquire(amount):
                return {
                    "allowed": False,
                    "limited_dimension": dim,
                    "retry_after_ms": self._get_retry_after(dim, amount),
                    "message": f"Rate limit exceeded: {dim}"
                }
        
        return {"allowed": True}
    
    def _get_retry_after(self, dimension: str, amount: int) -> int:
        """计算重试等待时间"""
        limiter = self.limiters[dimension]
        if hasattr(limiter, 'refill_rate'):
            return int((amount - limiter.tokens) / limiter.refill_rate * 1000)
        return 1000  # 默认1秒

class AdaptiveRateLimiter:
    """自适应限流器——基于系统负载动态调整"""
    
    def __init__(self, config: dict):
        self.min_limit = config.get("min_limit", 10)
        self.max_limit = config.get("max_limit", 1000)
        self.current_limit = config.get("initial_limit", 100)
        
        # AIMD参数
        self.additive_increase = config.get("additive_increase", 1)
        self.multiplicative_decrease = config.get("multiplicative_decrease", 0.5)
        
        # 系统指标
        self.metrics = {
            "cpu_utilization": 0,
            "memory_utilization": 0,
            "gpu_utilization": 0,
            "p99_latency_ms": 0,
            "error_rate": 0,
        }
        
        # 调整周期
        self.adjustment_interval = 10  # 秒
        self.window_requests = 0
        self.window_errors = 0
    
    async def should_allow(self, request: dict) -> bool:
        """判断是否允许请求"""
        current_rps = self.window_requests / self.adjustment_interval
        
        if current_rps >= self.current_limit:
            return False
        
        self.window_requests += 1
        return True
    
    async def adjust_limits(self):
        """定期调整限流阈值"""
        while True:
            await asyncio.sleep(self.adjustment_interval)
            
            # 计算系统健康度
            health_score = self._calculate_health()
            
            if health_score > 0.8:
                # 系统健康，增加限流（AIMD的AI）
                self.current_limit = min(
                    self.current_limit + self.additive_increase,
                    self.max_limit
                )
                logger.info(
                    f"Increasing rate limit to {self.current_limit} "
                    f"(health: {health_score:.2f})"
                )
            elif health_score < 0.5:
                # 系统不健康，减少限流（AIMD的MD）
                self.current_limit = max(
                    int(self.current_limit * self.multiplicative_decrease),
                    self.min_limit
                )
                logger.warning(
                    f"Decreasing rate limit to {self.current_limit} "
                    f"(health: {health_score:.2f})"
                )
            
            # 重置窗口
            self.window_requests = 0
            self.window_errors = 0
    
    def _calculate_health(self) -> float:
        """计算系统健康度（0-1）"""
        weights = {
            "cpu_utilization": 0.2,
            "memory_utilization": 0.15,
            "gpu_utilization": 0.25,
            "p99_latency_ms": 0.25,
            "error_rate": 0.15,
        }
        
        thresholds = {
            "cpu_utilization": 0.8,
            "memory_utilization": 0.85,
            "gpu_utilization": 0.9,
            "p99_latency_ms": 2000,
            "error_rate": 0.05,
        }
        
        score = 1.0
        for metric, weight in weights.items():
            value = self.metrics[metric]
            threshold = thresholds[metric]
            ratio = value / threshold if threshold > 0 else 0
            
            if ratio > 1:
                # 超过阈值，扣分
                score -= weight * min(ratio - 1, 1)
        
        return max(0, score)

class CircuitBreaker:
    """熔断器"""
    
    class State:
        CLOSED = "closed"      # 正常工作
        OPEN = "open"          # 熔断，拒绝请求
        HALF_OPEN = "half_open" # 半开，试探恢复
    
    def __init__(
        self,
        failure_threshold: int = 10,
        failure_rate_threshold: float = 0.5,
        recovery_timeout: float = 30.0,
        half_open_max_calls: int = 3
    ):
        self.state = self.State.CLOSED
        self.failure_threshold = failure_threshold
        self.failure_rate_threshold = failure_rate_threshold
        self.recovery_timeout = recovery_timeout
        self.half_open_max_calls = half_open_max_calls
        
        self.failure_count = 0
        self.success_count = 0
        self.total_count = 0
        self.last_failure_time = None
        self.half_open_calls = 0
    
    async def call(self, func, *args, **kwargs):
        """通过熔断器调用函数"""
        if self.state == self.State.OPEN:
            if self._should_attempt_recovery():
                self.state = self.State.HALF_OPEN
                self.half_open_calls = 0
                logger.info("Circuit breaker entering half-open state")
            else:
                raise CircuitBreakerOpenError(
                    f"Circuit breaker is open. "
                    f"Retry after {self._recovery_remaining():.0f}s"
                )
        
        if self.state == self.State.HALF_OPEN:
            if self.half_open_calls >= self.half_open_max_calls:
                raise CircuitBreakerOpenError(
                    "Half-open: max probe calls reached"
                )
            self.half_open_calls += 1
        
        try:
            result = await func(*args, **kwargs)
            await self._on_success()
            return result
        except Exception as e:
            await self._on_failure()
            raise
    
    async def _on_success(self):
        self.success_count += 1
        self.total_count += 1
        
        if self.state == self.State.HALF_OPEN:
            if self.half_open_calls >= self.half_open_max_calls:
                self.state = self.State.CLOSED
                self.failure_count = 0
                logger.info("Circuit breaker recovered, closing")
    
    async def _on_failure(self):
        self.failure_count += 1
        self.total_count += 1
        self.last_failure_time = time.monotonic()
        
        if self.state == self.State.HALF_OPEN:
            self.state = self.State.OPEN
            logger.warning("Circuit breaker re-opened from half-open")
        elif self.state == self.State.CLOSED:
            failure_rate = self.failure_count / max(self.total_count, 1)
            if (self.failure_count >= self.failure_threshold or
                failure_rate >= self.failure_rate_threshold):
                self.state = self.State.OPEN
                logger.error(
                    f"Circuit breaker opened: "
                    f"{self.failure_count}/{self.total_count} failures "
                    f"({failure_rate:.1%})"
                )

class MultiLevelCircuitBreaker:
    """多级熔断器"""
    
    def __init__(self):
        self.breakers = {
            "llm_inference": CircuitBreaker(
                failure_threshold=5,
                failure_rate_threshold=0.3,
                recovery_timeout=30
            ),
            "tool_search": CircuitBreaker(
                failure_threshold=10,
                failure_rate_threshold=0.5,
                recovery_timeout=15
            ),
            "vector_db": CircuitBreaker(
                failure_threshold=8,
                failure_rate_threshold=0.4,
                recovery_timeout=10
            ),
            "external_api": CircuitBreaker(
                failure_threshold=3,
                failure_rate_threshold=0.2,
                recovery_timeout=60
            ),
        }
    
    async def call_with_fallback(
        self,
        primary: callable,
        fallbacks: list,
        circuit_key: str
    ):
        """带降级的熔断调用"""
        breaker = self.breakers.get(circuit_key)
        
        try:
            if breaker:
                return await breaker.call(primary)
            return await primary()
        except CircuitBreakerOpenError:
            # 主路径熔断，尝试降级
            for i, fallback in enumerate(fallbacks):
                try:
                    logger.info(f"Trying fallback {i+1}/{len(fallbacks)}")
                    return await fallback()
                except Exception as e:
                    logger.warning(f"Fallback {i+1} failed: {e}")
                    continue
            
            # 所有降级都失败
            raise AllFallbacksFailedError(
                f"All fallbacks failed for {circuit_key}"
            )

class DegradationStrategy:
    """Agent降级策略"""
    
    LEVELS = {
        "normal": {
            "model": "gpt-4o",
            "max_tools": 10,
            "max_context": 128000,
            "streaming": True,
        },
        "degraded_1": {
            "model": "gpt-4o-mini",  # 降级到小模型
            "max_tools": 5,
            "max_context": 32000,
            "streaming": True,
        },
        "degraded_2": {
            "model": "gpt-4o-mini",
            "max_tools": 2,  # 仅保留核心工具
            "max_context": 8000,
            "streaming": False,
        },
        "emergency": {
            "model": "cached-response",  # 使用缓存或模板回复
            "max_tools": 0,
            "max_context": 1000,
            "streaming": False,
        }
    }
    
    def get_current_level(self, system_load: float) -> str:
        """根据系统负载获取降级级别"""
        if system_load < 0.7:
            return "normal"
        elif system_load < 0.85:
            return "degraded_1"
        elif system_load < 0.95:
            return "degraded_2"
        else:
            return "emergency"