AI Gateway：模型路由与推理优化

from dataclasses import dataclass, field
from typing import Callable, Optional, List, Dict, Any
from enum import Enum
import re
import hashlib


class TaskType(Enum):
    """任务类型枚举"""
    CODE_GENERATION = "code_generation"
    CODE_REVIEW = "code_review"
    TEXT_SUMMARIZATION = "text_summarization"
    QUESTION_ANSWERING = "question_answering"
    CREATIVE_WRITING = "creative_writing"
    REASONING = "reasoning"
    GENERAL = "general"


@dataclass
class RoutingRule:
    """路由规则定义"""
    name: str
    priority: int = 0  # 优先级，数字越大优先级越高
    task_type: Optional[TaskType] = None
    prompt_pattern: Optional[str] = None  # 正则表达式模式
    min_tokens: int = 0  # 最小输入Token数
    max_tokens: int = 100000  # 最大输入Token数
    required_keywords: List[str] = field(default_factory=list)  # 必须包含的关键词
    excluded_keywords: List[str] = field(default_factory=list)  # 不能包含的关键词
    model_selector: Callable[["RoutingContext"], str] = field(default=None)  # 自定义选择器
    fallback_models: List[str] = field(default_factory=list)  # 备用模型列表
    
    def __post_init__(self):
        if self.prompt_pattern:
            self._compiled_pattern = re.compile(self.prompt_pattern)
        else:
            self._compiled_pattern = None
    
    def matches(self, context: "RoutingContext") -> bool:
        """检查规则是否匹配给定的上下文"""
        # 检查任务类型
        if self.task_type and context.task_type != self.task_type:
            return False
        
        # 检查Token数量范围
        if not (self.min_tokens <= context.input_tokens <= self.max_tokens):
            return False
        
        # 检查Prompt模式
        if self._compiled_pattern and not self._compiled_pattern.search(context.prompt):
            return False
        
        # 检查必需关键词
        for keyword in self.required_keywords:
            if keyword.lower() not in context.prompt.lower():
                return False
        
        # 检查排除关键词
        for keyword in self.excluded_keywords:
            if keyword.lower() in context.prompt.lower():
                return False
        
        return True


@dataclass
class RoutingContext:
    """路由上下文，包含进行路由决策所需的所有信息"""
    prompt: str
    input_tokens: int
    task_type: TaskType = TaskType.GENERAL
    user_id: str = ""
    project_id: str = ""
    priority: int = 0  # 请求优先级，0-100
    metadata: Dict[str, Any] = field(default_factory=dict)
    
    @property
    def prompt_hash(self) -> str:
        """返回Prompt的哈希值，用于缓存键生成"""
        return hashlib.sha256(self.prompt.encode()).hexdigest()[:16]


class RuleBasedRouter:
    """基于规则的路由器"""
    
    def __init__(self):
        self._rules: List[RoutingRule] = []
        self._model_registry: Dict[str, Dict[str, Any]] = {}
    
    def add_rule(self, rule: RoutingRule) -> "RuleBasedRouter":
        """添加路由规则"""
        self._rules.append(rule)
        self._rules.sort(key=lambda r: r.priority, reverse=True)
        return self
    
    def register_model(self, model_id: str, capability: List[TaskType], 
                      cost_per_1k_input: float, cost_per_1k_output: float,
                      avg_latency_ms: float = 1000, max_concurrency: int = 10) -> "RuleBasedRouter":
        """注册模型信息"""
        self._model_registry[model_id] = {
            "capabilities": capability,
            "cost_input": cost_per_1k_input,
            "cost_output": cost_per_1k_output,
            "avg_latency_ms": avg_latency_ms,
            "max_concurrency": max_concurrency,
            "current_load": 0
        }
        return self
    
    def route(self, context: RoutingContext) -> str:
        """执行路由，返回选中的模型ID"""
        # 按优先级遍历规则，找到第一条匹配的规则
        for rule in self._rules:
            if rule.matches(context):
                # 如果有自定义选择器，使用自定义选择器
                if rule.model_selector:
                    selected = rule.model_selector(context)
                    if selected and selected in self._model_registry:
                        return selected
                
                # 否则根据规则配置选择模型
                # 优先选择支持该任务类型且负载最低的模型
                suitable_models = [
                    model_id for model_id, info in self._model_registry.items()
                    if context.task_type in info["capabilities"]
                ]
                
                if suitable_models:
                    # 选择负载最低的模型（简单的负载均衡）
                    return min(suitable_models, 
                              key=lambda m: self._model_registry[m]["current_load"])
        
        # 如果没有匹配的规则，返回默认模型
        return "gpt-3.5-turbo"
    
    def get_fallback_chain(self, primary_model: str) -> List[str]:
        """获取备用模型链"""
        for rule in self._rules:
            if rule.fallback_models and primary_model in rule.fallback_models:
                return rule.fallback_models
        
        # 默认的备用模型链
        if primary_model == "gpt-4-turbo":
            return ["gpt-4", "gpt-3.5-turbo"]
        elif primary_model == "claude-3-opus":
            return ["claude-3-sonnet", "claude-2"]
        else:
            return ["gpt-3.5-turbo"]

from dataclasses import dataclass, field
from typing import List, Dict, Optional, Tuple
from datetime import datetime, timedelta
import asyncio
import aiohttp
import json


@dataclass
class ModelMetrics:
    """模型性能指标"""
    model_id: str
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    total_latency_ms: float = 0.0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    last_updated: datetime = field(default_factory=datetime.now)
    
    @property
    def success_rate(self) -> float:
        if self.total_requests == 0:
            return 1.0
        return self.successful_requests / self.total_requests
    
    @property
    def avg_latency_ms(self) -> float:
        if self.successful_requests == 0:
            return float('inf')
        return self.total_latency_ms / self.successful_requests
    
    @property
    def avg_cost_per_1k_tokens(self) -> float:
        total_tokens = self.total_input_tokens + self.total_output_tokens
        if total_tokens == 0:
            return 0.0
        return (self.total_input_tokens * 0.5 + self.total_output_tokens * 0.5) / total_tokens


@dataclass
class ModelInfo:
    """模型信息"""
    model_id: str
    display_name: str
    provider: str
    context_window: int
    capabilities: List[str]
    cost_input_per_1k: float
    cost_output_per_1k: float
    max_rpm: int  # 每分钟请求数限制
    max_tpm: int  # 每分钟Token数限制
    is_available: bool = True
    is_degraded: bool = False
    degradation_reason: Optional[str] = None


class LLMRouter:
    """
    LLM驱动的智能路由器
    
    该路由器综合考虑以下因素进行模型选择：
    1. 任务复杂度估计
    2. 模型能力匹配度
    3. 当前负载和可用性
    4. 成本效率
    5. 历史性能表现
    """
    
    def __init__(self, config: Dict):
        self.config = config
        self.models: Dict[str, ModelInfo] = {}
        self.metrics: Dict[str, ModelMetrics] = {}
        self.routing_prompt_template = config.get("routing_prompt_template", self._default_routing_prompt())
        self._session: Optional[aiohttp.ClientSession] = None
    
    def _default_routing_prompt(self) -> str:
        """默认的路由提示词"""
        return """你是一个AI模型路由专家。请根据以下请求信息，从候选模型中选择最合适的模型。

请求信息：
- 任务类型：{task_type}
- 输入Token数（估计）：{input_tokens}
- 输出Token数（估计）：{output_tokens}
- 优先级：{priority}

候选模型：
{candidate_models}

请从以下维度评估每个模型：
1. 能力匹配度（是否能胜任该任务）
2. 成本效率（是否经济实惠）
3. 当前负载（是否响应迅速）
4. 性能历史（是否有良好的表现）

请以JSON格式返回结果：
{{
    "selected_model": "模型ID",
    "reasoning": "选择理由",
    "confidence": 0.95,
    "alternatives": ["备选模型1", "备选模型2"]
}}
"""
    
    async def initialize(self):
        """初始化路由器"""
        self._session = aiohttp.ClientSession()
    
    async def close(self):
        """关闭路由器"""
        if self._session:
            await self._session.close()
    
    def register_model(self, model: ModelInfo):
        """注册模型"""
        self.models[model.model_id] = model
        if model.model_id not in self.metrics:
            self.metrics[model.model_id] = ModelMetrics(model_id=model.model_id)
    
    def update_metrics(self, model_id: str, latency_ms: float, 
                      input_tokens: int, output_tokens: int, success: bool):
        """更新模型指标"""
        if model_id not in self.metrics:
            self.metrics[model_id] = ModelMetrics(model_id=model_id)
        
        m = self.metrics[model_id]
        m.total_requests += 1
        m.total_latency_ms += latency_ms
        m.total_input_tokens += input_tokens
        m.total_output_tokens += output_tokens
        m.last_updated = datetime.now()
        
        if success:
            m.successful_requests += 1
        else:
            m.failed_requests += 1
    
    async def estimate_task_complexity(self, prompt: str, task_type: str) -> float:
        """
        估计任务复杂度
        
        返回值范围：0.0 - 1.0
        - 0.0-0.3: 简单任务
        - 0.3-0.6: 中等任务
        - 0.6-1.0: 复杂任务
        """
        # 简化的复杂度估计
        complexity_indicators = [
            len(prompt) > 1000,  # 长文本
            any(kw in prompt.lower() for kw in ["分析", "推理", "证明", "比较"]),
            any(kw in prompt.lower() for kw in ["代码", "程序", "函数", "算法"]),
            prompt.count("\n") > 10,
            len(set(prompt)) / len(prompt) < 0.3,
        ]
        
        score = sum(complexity_indicators) / len(complexity_indicators)
        
        if task_type in ["reasoning", "code_generation", "code_review"]:
            score = min(1.0, score * 1.3)
        
        return min(1.0, score)
    
    def calculate_model_score(self, model_id: str, complexity: float,
                            task_type: str) -> Tuple[float, str]:
        """计算模型对当前任务的适合度分数"""
        model = self.models.get(model_id)
        metrics = self.metrics.get(model_id, ModelMetrics(model_id=model_id))
        
        if not model or not model.is_available:
            return 0.0, "模型不可用"
        
        if model.is_degraded:
            return 0.0, f"模型降级: {model.degradation_reason}"
        
        recent_requests = metrics.total_requests
        if recent_requests >= model.max_rpm * 0.9:
            return 0.3, "接近RPM限制"
        
        # 能力匹配度 (40%权重)
        capability_score = 0.0
        if task_type in model.capabilities:
            capability_score = 1.0
        elif "general" in model.capabilities:
            capability_score = 0.5
        else:
            capability_score = 0.2
        
        # 成本效率 (20%权重)
        min_cost = min(m.cost_input_per_1k for m in self.models.values() if m.is_available)
        max_cost = max(m.cost_input_per_1k for m in self.models.values() if m.is_available)
        if max_cost > min_cost:
            cost_score = 1.0 - (model.cost_input_per_1k - min_cost) / (max_cost - min_cost)
        else:
            cost_score = 1.0
        
        cost_weight = 0.2 * (1.0 - complexity)
        capability_weight = 0.4 + 0.2 * complexity
        
        # 性能表现 (25%权重)
        performance_score = metrics.success_rate * 0.6 + \
                           (1.0 - min(1.0, metrics.avg_latency_ms / 5000)) * 0.4
        
        # 负载状况 (15%权重)
        load_score = 1.0 - min(1.0, metrics.total_requests / (model.max_rpm * 0.8))
        
        # 综合得分
        total_score = (capability_score * capability_weight +
                      cost_score * cost_weight +
                      performance_score * 0.25 +
                      load_score * 0.15)
        
        reasons = [
            f"能力匹配度: {capability_score:.2f}",
            f"成本效率: {cost_score:.2f}",
            f"性能表现: {performance_score:.2f}",
            f"负载状况: {load_score:.2f}",
        ]
        
        return total_score, "; ".join(reasons)
    
    async def route(self, task_type: str, prompt: str, 
                   estimated_input_tokens: int,
                   estimated_output_tokens: int = 1000,
                   priority: int = 50) -> str:
        """执行智能路由"""
        complexity = await self.estimate_task_complexity(prompt, task_type)
        
        available_models = [
            m for m in self.models.values() 
            if m.is_available and not m.is_degraded
        ]
        
        if not available_models:
            raise ValueError("没有可用的模型")
        
        model_scores = []
        for model in available_models:
            score, reasons = self.calculate_model_score(
                model.model_id, complexity, task_type
            )
            model_scores.append((model.model_id, score, reasons))
        
        model_scores.sort(key=lambda x: x[1], reverse=True)
        
        best_model = model_scores[0][0]
        self.metrics[best_model].total_requests += 1
        
        return best_model

from dataclasses import dataclass, field
from typing import Dict, List, Optional
import threading
import time


@dataclass
class DynamicWeight:
    """动态权重配置"""
    model_id: str
    base_weight: float = 1.0
    current_weight: float = 1.0
    min_weight: float = 0.1
    max_weight: float = 10.0
    last_adjustment: float = field(default_factory=time.time)
    adjustment_count: int = 0
    recent_latencies: List[float] = field(default_factory=list)
    recent_errors: List[bool] = field(default_factory=list)
    window_size: int = 100
    
    def record_latency(self, latency_ms: float):
        self.recent_latencies.append(latency_ms)
        if len(self.recent_latencies) > self.window_size:
            self.recent_latencies.pop(0)
    
    def record_success(self, success: bool):
        self.recent_errors.append(not success)
        if len(self.recent_errors) > self.window_size:
            self.recent_errors.pop(0)
    
    @property
    def avg_latency(self) -> float:
        if not self.recent_latencies:
            return 1000.0
        return sum(self.recent_latencies) / len(self.recent_latencies)
    
    @property
    def error_rate(self) -> float:
        if not self.recent_errors:
            return 0.0
        return sum(self.recent_errors) / len(self.recent_errors)


class DynamicWeightRouter:
    """
    动态权重路由器
    
    通过实时监控模型性能，动态调整各模型的权重，
    实现自动负载均衡和故障转移。
    """
    
    def __init__(self, health_check_interval: float = 30.0,
                 weight_decay: float = 0.95,
                 adjustment_rate: float = 0.1):
        self._weights: Dict[str, DynamicWeight] = {}
        self._lock = threading.RLock()
        self._health_check_interval = health_check_interval
        self._weight_decay = weight_decay
        self._adjustment_rate = adjustment_rate
        self._running = False
        self._health_check_thread: Optional[threading.Thread] = None
    
    def register_model(self, model_id: str, base_weight: float = 1.0):
        with self._lock:
            self._weights[model_id] = DynamicWeight(
                model_id=model_id,
                base_weight=base_weight,
                current_weight=base_weight
            )
    
    def record_result(self, model_id: str, latency_ms: float, success: bool):
        with self._lock:
            if model_id not in self._weights:
                return
            
            weight = self._weights[model_id]
            weight.record_latency(latency_ms)
            weight.record_success(success)
    
    def adjust_weights(self):
        """调整模型权重"""
        with self._lock:
            for model_id, weight in self._weights.items():
                latency_score = 1.0 - min(1.0, weight.avg_latency / 5000)
                error_score = 1.0 - weight.error_rate
                perf_score = latency_score * 0.7 + error_score * 0.3
                target_weight = weight.base_weight * perf_score
                
                weight.current_weight = (
                    weight.current_weight * (1 - self._adjustment_rate) +
                    target_weight * self._adjustment_rate
                )
                
                weight.current_weight = max(
                    weight.min_weight,
                    min(weight.max_weight, weight.current_weight)
                )
                
                weight.last_adjustment = time.time()
                weight.adjustment_count += 1
    
    def select_model(self, exclude_models: Optional[List[str]] = None) -> str:
        """使用加权随机算法选择模型"""
        with self._lock:
            available = [
                (m_id, w) for m_id, w in self._weights.items()
                if m_id not in (exclude_models or [])
            ]
            
            if not available:
                raise ValueError("没有可用的模型")
            
            total_weight = sum(w.current_weight for _, w in available)
            
            import random
            r = random.uniform(0, total_weight)
            cumulative = 0.0
            
            for model_id, weight in available:
                cumulative += weight.current_weight
                if r <= cumulative:
                    return model_id
            
            return available[0][0]
    
    def get_weight_status(self) -> Dict[str, Dict]:
        with self._lock:
            return {
                m_id: {
                    "base_weight": w.base_weight,
                    "current_weight": round(w.current_weight, 4),
                    "avg_latency_ms": round(w.avg_latency, 2),
                    "error_rate": round(w.error_rate, 4),
                    "adjustment_count": w.adjustment_count,
                    "health_status": self._calculate_health_status(w)
                }
                for m_id, w in self._weights.items()
            }
    
    def _calculate_health_status(self, weight: DynamicWeight) -> str:
        if weight.error_rate > 0.1:
            return "critical"
        elif weight.error_rate > 0.05:
            return "degraded"
        elif weight.avg_latency > 3000:
            return "slow"
        else:
            return "healthy"

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Callable
from collections import defaultdict
from datetime import datetime, timedelta
import threading
import time
import heapq
import random


@dataclass
class ModelInstance:
    """模型实例"""
    instance_id: str
    model_id: str
    endpoint: str
    max_concurrency: int = 100
    current_load: int = 0
    is_healthy: bool = True
    last_health_check: datetime = field(default_factory=datetime.now)
    avg_response_time: float = 1000.0
    total_requests: int = 0
    failed_requests: int = 0
    
    @property
    def available_concurrency(self) -> int:
        return max(0, self.max_concurrency - self.current_load)
    
    @property
    def health_score(self) -> float:
        if not self.is_healthy:
            return 0.0
        time_score = max(0, 1.0 - self.avg_response_time / 10000)
        fail_score = 1.0 - (self.failed_requests / max(1, self.total_requests))
        return time_score * 0.7 + fail_score * 0.3


class LoadBalancer(ABC):
    """负载均衡器抽象基类"""
    
    @abstractmethod
    def select_instance(self, instances: List[ModelInstance]) -> Optional[ModelInstance]:
        pass
    
    @abstractmethod
    def record_result(self, instance_id: str, latency_ms: float, success: bool):
        pass


class RoundRobinLB(LoadBalancer):
    """轮询负载均衡器"""
    
    def __init__(self):
        self._counter: int = 0
        self._lock = threading.Lock()
    
    def select_instance(self, instances: List[ModelInstance]) -> Optional[ModelInstance]:
        healthy = [i for i in instances if i.is_healthy and i.available_concurrency > 0]
        if not healthy:
            return None
        
        with self._lock:
            idx = self._counter % len(healthy)
            self._counter += 1
            return healthy[idx]
    
    def record_result(self, instance_id: str, latency_ms: float, success: bool):
        pass


class WeightedRoundRobinLB(LoadBalancer):
    """加权轮询负载均衡器"""
    
    def __init__(self):
        self._weights: Dict[str, int] = {}
        self._counters: Dict[str, int] = {}
        self._lock = threading.Lock()
    
    def set_weight(self, instance_id: str, weight: int):
        with self._lock:
            self._weights[instance_id] = weight
            self._counters[instance_id] = 0
    
    def select_instance(self, instances: List[ModelInstance]) -> Optional[ModelInstance]:
        healthy = [i for i in instances if i.is_healthy and i.available_concurrency > 0]
        if not healthy:
            return None
        
        with self._lock:
            def effective_weight(inst: ModelInstance) -> float:
                weight = self._weights.get(inst.instance_id, 1)
                counter = self._counters.get(inst.instance_id, 0)
                return weight / max(1, counter + 1)
            
            selected = min(healthy, key=effective_weight)
            self._counters[selected.instance_id] = self._counters.get(selected.instance_id, 0) + 1
            return selected
    
    def record_result(self, instance_id: str, latency_ms: float, success: bool):
        with self._lock:
            if not success:
                self._counters[instance_id] = 0


class LeastConnectionsLB(LoadBalancer):
    """最少连接负载均衡器"""
    
    def __init__(self):
        self._lock = threading.Lock()
    
    def select_instance(self, instances: List[ModelInstance]) -> Optional[ModelInstance]:
        healthy = [i for i in instances if i.is_healthy and i.available_concurrency > 0]
        if not healthy:
            return None
        
        with self._lock:
            return min(healthy, key=lambda i: i.current_load)
    
    def record_result(self, instance_id: str, latency_ms: float, success: bool):
        pass


class LatencyAwareLB(LoadBalancer):
    """
    延迟感知负载均衡器
    
    该策略综合考虑：
    1. 实例的响应延迟（越低越好）
    2. 实例的当前负载（越低越好）
    3. 实例的健康状态
    """
    
    def __init__(self, latency_weight: float = 0.6, load_weight: float = 0.4):
        self._latency_weight = latency_weight
        self._load_weight = load_weight
        self._scores: Dict[str, float] = {}
        self._lock = threading.Lock()
    
    def select_instance(self, instances: List[ModelInstance]) -> Optional[ModelInstance]:
        healthy = [i for i in instances if i.is_healthy and i.available_concurrency > 0]
        if not healthy:
            return None
        
        with self._lock:
            max_latency = max(i.avg_response_time for i in healthy) or 1
            min_latency = min(i.avg_response_time for i in healthy) or 1
            
            for inst in healthy:
                latency_score = 1.0 - (inst.avg_response_time - min_latency) / (max_latency - min_latency + 1)
                load_score = 1.0 - (inst.current_load / max(1, inst.max_concurrency))
                
                score = (latency_score * self._latency_weight + 
                        load_score * self._load_weight)
                
                self._scores[inst.instance_id] = score
            
            return max(healthy, key=lambda i: self._scores.get(i.instance_id, 0))
    
    def record_result(self, instance_id: str, latency_ms: float, success: bool):
        with self._lock:
            if instance_id in self._scores:
                alpha = 0.3
                self._scores[instance_id] = (
                    alpha * self._scores[instance_id] +
                    (1 - alpha) * (1.0 - min(1.0, latency_ms / 5000))
                )


class LoadBalancerFactory:
    """负载均衡器工厂"""
    
    _strategies = {
        "round_robin": RoundRobinLB,
        "weighted_round_robin": WeightedRoundRobinLB,
        "least_connections": LeastConnectionsLB,
        "latency_aware": LatencyAwareLB,
    }
    
    @classmethod
    def create(cls, strategy: str, **kwargs) -> LoadBalancer:
        strategy_class = cls._strategies.get(strategy)
        if not strategy_class:
            raise ValueError(f"Unknown strategy: {strategy}")
        return strategy_class(**kwargs)

import hashlib
import bisect
from typing import List, Dict, Optional, Tuple


class ConsistentHashRouter:
    """
    一致性哈希路由
    
    特点：
    1. 相同的key（用户ID/会话ID）总是路由到相同的节点
    2. 当节点发生变化时，只有少量key的映射会改变
    3. 支持虚拟节点，提高负载均衡性
    """
    
    def __init__(self, virtual_nodes: int = 150):
        self._virtual_nodes = virtual_nodes
        self._hash_ring: List[Tuple[int, str]] = []
        self._node_map: Dict[str, str] = {}
        self._lock = threading.Lock()
    
    def _hash(self, key: str) -> int:
        return int(hashlib.md5(key.encode()).hexdigest(), 16)
    
    def add_node(self, instance_id: str):
        with self._lock:
            for i in range(self._virtual_nodes):
                virtual_id = f"{instance_id}#VN{i}"
                h = self._hash(virtual_id)
                bisect.insort(self._hash_ring, (h, virtual_id))
                self._node_map[virtual_id] = instance_id
    
    def remove_node(self, instance_id: str):
        with self._lock:
            virtual_ids = [f"{instance_id}#VN{i}" for i in range(self._virtual_nodes)]
            self._hash_ring = [
                (h, vid) for h, vid in self._hash_ring
                if vid not in virtual_ids
            ]
            for vid in virtual_ids:
                self._node_map.pop(vid, None)
    
    def get_node(self, key: str) -> Optional[str]:
        if not self._hash_ring:
            return None
        
        h = self._hash(key)
        idx = bisect.bisect_left(self._hash_ring, (h, ""))
        
        if idx >= len(self._hash_ring):
            idx = 0
        
        virtual_id = self._hash_ring[idx][1]
        return self._node_map.get(virtual_id)

from dataclasses import dataclass, field
from typing import Dict, Optional, List, Callable
from datetime import datetime, timedelta, date
from enum import Enum
import threading
import json


class BudgetPeriod(Enum):
    DAILY = "daily"
    WEEKLY = "weekly"
    MONTHLY = "monthly"
    QUARTERLY = "quarterly"


class BudgetAction(Enum):
    BLOCK = "block"
    QUEUE = "queue"
    DEGRADE = "degrade"
    WARN = "warn"


@dataclass
class TokenUsage:
    user_id: str
    date: date
    input_tokens: int = 0
    output_tokens: int = 0
    request_count: int = 0
    cost: float = 0.0
    
    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class Budget:
    entity_id: str
    entity_type: str
    period: BudgetPeriod = BudgetPeriod.MONTHLY
    token_limit: int = 1_000_000
    cost_limit: float = 1000.0
    action_when_exceeded: BudgetAction = BudgetAction.WARN
    model_limits: Dict[str, int] = field(default_factory=dict)
    soft_limit_percent: float = 0.8


@dataclass
class BudgetStatus:
    budget: Budget
    current_usage: TokenUsage
    percent_used: float = 0.0
    percent_cost_used: float = 0.0
    remaining_tokens: int = 0
    remaining_cost: float = 0.0
    is_exceeded: bool = False
    is_soft_limit_exceeded: bool = False


class BudgetManager:
    """
    预算管理器
    
    职责：
    1. 跟踪用户/项目/组织的Token使用量
    2. 实施预算控制策略
    3. 提供预算预警
    4. 生成成本报告
    """
    
    def __init__(self, storage_path: Optional[str] = None):
        self._budgets: Dict[str, Budget] = {}
        self._usage: Dict[str, List[TokenUsage]] = {}
        self._warnings: Dict[str, List[datetime]] = {}
        self._lock = threading.RLock()
        self._storage_path = storage_path
        self._cost_rates: Dict[str, Dict[str, float]] = {
            "gpt-4-turbo": {"input": 0.01, "output": 0.03},
            "gpt-4": {"input": 0.03, "output": 0.06},
            "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015},
            "claude-3-opus": {"input": 0.015, "output": 0.075},
            "claude-3-sonnet": {"input": 0.003, "output": 0.015},
        }
        
        self._on_budget_exceeded: Optional[Callable] = None
        self._on_soft_limit: Optional[Callable] = None
    
    def set_cost_rate(self, model_id: str, input_rate: float, output_rate: float):
        self._cost_rates[model_id] = {"input": input_rate, "output": output_rate}
    
    def create_budget(self, budget: Budget):
        with self._lock:
            key = f"{budget.entity_type}:{budget.entity_id}"
            self._budgets[key] = budget
            self._usage.setdefault(key, [])
            self._warnings.setdefault(key, [])
    
    def get_budget(self, entity_id: str, entity_type: str) -> Optional[Budget]:
        key = f"{entity_type}:{entity_id}"
        return self._budgets.get(key)
    
    def _get_current_period(self, period: BudgetPeriod) -> Tuple[date, date]:
        today = date.today()
        
        if period == BudgetPeriod.DAILY:
            return today, today
        elif period == BudgetPeriod.WEEKLY:
            start = today - timedelta(days=today.weekday())
            end = start + timedelta(days=6)
            return start, end
        elif period == BudgetPeriod.MONTHLY:
            start = today.replace(day=1)
            if start.month == 12:
                end = start.replace(year=start.year + 1, month=1, day=1) - timedelta(days=1)
            else:
                end = start.replace(month=start.month + 1, day=1) - timedelta(days=1)
            return start, end
        elif period == BudgetPeriod.QUARTERLY:
            quarter = (today.month - 1) // 3
            start = today.replace(month=quarter * 3 + 1, day=1)
            end_month = (quarter + 1) * 3 + 1
            if end_month > 12:
                end = today.replace(year=start.year + 1, month=end_month - 12, day=1) - timedelta(days=1)
            else:
                end = start.replace(month=end_month, day=1) - timedelta(days=1)
            return start, end
    
    def _get_usage_for_period(self, entity_id: str, entity_type: str, 
                              period: BudgetPeriod) -> TokenUsage:
        key = f"{entity_type}:{entity_id}"
        start_date, end_date = self._get_current_period(period)
        
        total_input = 0
        total_output = 0
        total_requests = 0
        total_cost = 0.0
        
        with self._lock:
            for usage in self._usage.get(key, []):
                if start_date <= usage.date <= end_date:
                    total_input += usage.input_tokens
                    total_output += usage.output_tokens
                    total_requests += usage.request_count
                    total_cost += usage.cost
        
        return TokenUsage(
            user_id=entity_id,
            date=end_date,
            input_tokens=total_input,
            output_tokens=total_output,
            request_count=total_requests,
            cost=total_cost
        )
    
    def record_usage(self, entity_id: str, entity_type: str,
                    model_id: str, input_tokens: int, output_tokens: int):
        key = f"{entity_type}:{entity_id}"
        
        rates = self._cost_rates.get(model_id, {"input": 0.01, "output": 0.03})
        cost = (input_tokens * rates["input"] + output_tokens * rates["output"]) / 1000
        
        usage = TokenUsage(
            user_id=entity_id,
            date=date.today(),
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            request_count=1,
            cost=cost
        )
        
        with self._lock:
            self._usage.setdefault(key, []).append(usage)
            
            cutoff = date.today() - timedelta(days=90)
            self._usage[key] = [
                u for u in self._usage[key] if u.date >= cutoff
            ]
        
        budget = self.get_budget(entity_id, entity_type)
        if budget:
            status = self.check_budget(entity_id, entity_type)
            self._handle_budget_status(entity_id, entity_type, status)
    
    def check_budget(self, entity_id: str, entity_type: str) -> BudgetStatus:
        budget = self.get_budget(entity_id, entity_type)
        if not budget:
            return None
        
        current_usage = self._get_usage_for_period(entity_id, entity_type, budget.period)
        
        percent_used = current_usage.total_tokens / budget.token_limit if budget.token_limit > 0 else 0
        percent_cost = current_usage.cost / budget.cost_limit if budget.cost_limit > 0 else 0
        
        return BudgetStatus(
            budget=budget,
            current_usage=current_usage,
            percent_used=percent_used,
            percent_cost_used=percent_cost,
            remaining_tokens=max(0, budget.token_limit - current_usage.total_tokens),
            remaining_cost=max(0, budget.cost_limit - current_usage.cost),
            is_exceeded=percent_used >= 1.0 or percent_cost >= 1.0,
            is_soft_limit_exceeded=percent_used >= budget.soft_limit_percent or percent_cost >= budget.soft_limit_percent
        )
    
    def can_proceed(self, entity_id: str, entity_type: str,
                   estimated_tokens: int = 0, estimated_cost: float = 0.0) -> bool:
        status = self.check_budget(entity_id, entity_type)
        if not status:
            return True
        
        if status.is_exceeded:
            return False
        
        if (status.current_usage.total_tokens + estimated_tokens > status.budget.token_limit or
            status.current_usage.cost + estimated_cost > status.budget.cost_limit):
            return False
        
        return True
    
    def get_cost_report(self, entity_id: str, entity_type: str,
                       days: int = 30) -> Dict:
        key = f"{entity_type}:{entity_id}"
        cutoff = date.today() - timedelta(days=days)
        
        daily_costs: Dict[date, float] = {}
        total_cost = 0.0
        total_tokens = 0
        total_requests = 0
        
        with self._lock:
            for usage in self._usage.get(key, []):
                if usage.date >= cutoff:
                    daily_costs[usage.date] = usage.cost
                    total_cost += usage.cost
                    total_tokens += usage.total_tokens
                    total_requests += usage.request_count
        
        return {
            "entity_id": entity_id,
            "entity_type": entity_type,
            "period_days": days,
            "total_cost": round(total_cost, 4),
            "total_tokens": total_tokens,
            "total_requests": total_requests,
            "avg_cost_per_day": round(total_cost / days, 4),
            "avg_tokens_per_day": total_tokens // days,
            "daily_costs": {str(k): round(v, 4) for k, v in sorted(daily_costs.items())}
        }

import asyncio
import time
from dataclasses import dataclass, field
from typing import List, Optional, Callable, Dict, Any
from datetime import datetime, timedelta
from collections import defaultdict
import threading
import uuid


@dataclass
class BatchRequest:
    request_id: str
    prompt: str
    max_tokens: int
    temperature: float = 0.7
    metadata: Dict[str, Any] = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)
    future: asyncio.Future = field(default_factory=None)
    
    @property
    def priority(self) -> int:
        return self.metadata.get("priority", 50)
    
    @property
    def age_ms(self) -> float:
        return (datetime.now() - self.created_at).total_seconds() * 1000


@dataclass
class BatchResponse:
    request_id: str
    content: Optional[str]
    error: Optional[str]
    latency_ms: float
    input_tokens: int = 0
    output_tokens: int = 0


class BatchProcessor:
    """
    批处理处理器
    
    核心策略：
    1. 收集一定数量或时间的请求后统一处理
    2. 优先处理高优先级或等待时间长的请求
    3. 动态调整批次大小以优化吞吐量
    """
    
    def __init__(self, model_api: Callable,
                 max_batch_size: int = 32,
                 max_wait_ms: float = 100.0,
                 max_tokens_per_request: int = 2000):
        self._model_api = model_api
        self._max_batch_size = max_batch_size
        self._max_wait_ms = max_wait_ms
        self._max_tokens_per_request = max_tokens_per_request
        
        self._pending_requests: List[BatchRequest] = []
        self._lock = threading.RLock()
        self._processing = False
        
        self._stats = {
            "total_requests": 0,
            "total_batches": 0,
            "total_tokens": 0,
            "avg_batch_size": 0.0,
            "avg_latency_ms": 0.0
        }
    
    async def add_request(self, prompt: str, max_tokens: int = 1000,
                         temperature: float = 0.7,
                         priority: int = 50,
                         metadata: Optional[Dict] = None) -> str:
        request_id = str(uuid.uuid4())
        
        future = asyncio.Future()
        request = BatchRequest(
            request_id=request_id,
            prompt=prompt,
            max_tokens=min(max_tokens, self._max_tokens_per_request),
            temperature=temperature,
            metadata=metadata or {"priority": priority},
            future=future
        )
        
        with self._lock:
            self._pending_requests.append(request)
            self._stats["total_requests"] += 1
        
        asyncio.create_task(self._try_process_batch())
        
        return request_id
    
    async def _try_process_batch(self):
        with self._lock:
            should_process = (
                len(self._pending_requests) >= self._max_batch_size or
                self._processing
            )
            
            if should_process or not self._pending_requests:
                return
            
            oldest_age = self._pending_requests[0].age_ms if self._pending_requests else 0
            if oldest_age < self._max_wait_ms and len(self._pending_requests) < self._max_batch_size:
                return
            
            self._processing = True
            batch = self._prepare_batch()
        
        if batch:
            await self._process_batch(batch)
        
        with self._lock:
            self._processing = False
    
    def _prepare_batch(self) -> List[BatchRequest]:
        if not self._pending_requests:
            return []
        
        sorted_requests = sorted(
            self._pending_requests,
            key=lambda r: (-r.priority, -r.age_ms)
        )
        
        batch = sorted_requests[:self._max_batch_size]
        
        for request in batch:
            self._pending_requests.remove(request)
        
        return batch
    
    async def _process_batch(self, batch: List[BatchRequest]):
        if not batch:
            return
        
        start_time = time.time()
        
        try:
            prompts = [r.prompt for r in batch]
            responses = await self._model_api(prompts)
            
            for i, request in enumerate(batch):
                response = responses[i] if i < len(responses) else None
                
                if request.future and not request.future.done():
                    if response and not response.get("error"):
                        request.future.set_result(response)
                    else:
                        request.future.set_exception(
                            Exception(response.get("error", "Unknown error") if response else "No response")
                        )
        
        except Exception as e:
            for request in batch:
                if request.future and not request.future.done():
                    request.future.set_exception(e)
        
        elapsed_ms = (time.time() - start_time) * 1000
        with self._lock:
            self._stats["total_batches"] += 1
            self._stats["avg_batch_size"] = (
                (self._stats["avg_batch_size"] * (self._stats["total_batches"] - 1) + len(batch)) /
                self._stats["total_batches"]
            )
            self._stats["avg_latency_ms"] = (
                (self._stats["avg_latency_ms"] * (self._stats["total_batches"] - 1) + elapsed_ms) /
                self._stats["total_batches"]
            )
    
    def get_stats(self) -> Dict:
        with self._lock:
            pending = len(self._pending_requests)
        
        return {
            **self._stats,
            "pending_requests": pending,
            "processing": self._processing
        }

import asyncio
import json
from typing import AsyncIterator, Optional, Callable, Dict, Any
from dataclasses import dataclass, field
from datetime import datetime
import uuid


@dataclass
class StreamChunk:
    chunk_id: str
    content: str
    is_final: bool = False
    delta: str = ""
    index: int = 0
    usage: Optional[Dict] = None
    error: Optional[str] = None


class StreamingHandler:
    """流式响应处理器"""
    
    def __init__(self):
        self._callbacks: Dict[str, Callable] = {}
    
    async def create_stream(self, prompt: str, 
                           model_api: Callable,
                           stream_format: str = "sse") -> AsyncIterator[StreamChunk]:
        stream_id = str(uuid.uuid4())
        accumulated_content = ""
        index = 0
        
        async for event in model_api(prompt, stream=True):
            if event.get("error"):
                yield StreamChunk(
                    chunk_id=stream_id,
                    content="",
                    is_final=True,
                    index=index,
                    error=event["error"]
                )
                break
            
            delta = event.get("delta", "")
            
            if delta:
                accumulated_content += delta
                
                yield StreamChunk(
                    chunk_id=stream_id,
                    content=accumulated_content,
                    is_final=False,
                    delta=delta,
                    index=index
                )
                index += 1
            
            if event.get("done", False):
                yield StreamChunk(
                    chunk_id=stream_id,
                    content=accumulated_content,
                    is_final=True,
                    index=index,
                    usage=event.get("usage")
                )
                break

class PromptCache:
    """
    Prompt缓存
    
    对于相同的Prompt，直接返回缓存的响应，
    无需再次调用模型API。
    """
    
    def __init__(self, max_size: int = 10000, ttl_seconds: int = 3600):
        self._cache: Dict[str, Dict] = {}
        self._access_order: list = []
        self._max_size = max_size
        self._ttl_seconds = ttl_seconds
        self._lock = threading.RLock()
        
        self._stats = {
            "hits": 0,
            "misses": 0,
            "evictions": 0
        }
    
    def _generate_key(self, prompt: str, model_id: str, params: Dict) -> str:
        import hashlib
        content = f"{model_id}:{json.dumps(params, sort_keys=True)}:{prompt}"
        return hashlib.sha256(content.encode()).hexdigest()
    
    def get(self, prompt: str, model_id: str, params: Dict) -> Optional[Dict]:
        key = self._generate_key(prompt, model_id, params)
        
        with self._lock:
            entry = self._cache.get(key)
            
            if not entry:
                self._stats["misses"] += 1
                return None
            
            if datetime.now() - entry["cached_at"] > timedelta(seconds=self._ttl_seconds):
                del self._cache[key]
                self._stats["misses"] += 1
                return None
            
            if key in self._access_order:
                self._access_order.remove(key)
            self._access_order.append(key)
            
            self._stats["hits"] += 1
            return entry["response"]
    
    def put(self, prompt: str, model_id: str, params: Dict, response: Dict):
        key = self._generate_key(prompt, model_id, params)
        
        with self._lock:
            if len(self._cache) >= self._max_size and key not in self._cache:
                oldest_key = self._access_order.pop(0)
                del self._cache[oldest_key]
                self._stats["evictions"] += 1
            
            self._cache[key] = {
                "response": response,
                "cached_at": datetime.now(),
                "prompt": prompt,
                "model_id": model_id
            }
            
            if key in self._access_order:
                self._access_order.remove(key)
            self._access_order.append(key)
    
    def get_stats(self) -> Dict:
        with self._lock:
            total = self._stats["hits"] + self._stats["misses"]
            hit_rate = self._stats["hits"] / total if total > 0 else 0.0
            
            return {
                **self._stats,
                "size": len(self._cache),
                "max_size": self._max_size,
                "hit_rate": hit_rate
            }

from enum import Enum
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from typing import Callable, Optional, List, Dict, Any
import threading
import time


class CircuitState(Enum):
    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half_open"


@dataclass
class CircuitBreakerConfig:
    failure_threshold: float = 0.5
    success_threshold: int = 3
    timeout_seconds: float = 30.0
    min_requests: int = 10
    window_seconds: float = 60.0


@dataclass
class CircuitBreakerMetrics:
    total_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0
    rejected_requests: int = 0
    consecutive_failures: int = 0
    consecutive_successes: int = 0
    last_failure_time: Optional[datetime] = None
    state_changes: List[Dict] = field(default_factory=list)
    
    @property
    def failure_rate(self) -> float:
        if self.total_requests < 1:
            return 0.0
        return self.failed_requests / self.total_requests


class CircuitBreaker:
    """
    熔断器
    
    工作原理：
    1. 正常状态下（CLOSED），所有请求都正常通过
    2. 当失败率达到阈值时，切换到OPEN状态，拒绝所有请求
    3. 等待超时后，切换到HALF_OPEN状态，允许部分请求通过
    4. 如果这些请求成功，切换回CLOSED；否则重新切换到OPEN
    """
    
    def __init__(self, name: str, config: Optional[CircuitBreakerConfig] = None):
        self._name = name
        self._config = config or CircuitBreakerConfig()
        self._state = CircuitState.CLOSED
        self._metrics = CircuitBreakerMetrics()
        self._lock = threading.RLock()
        self._window_start = datetime.now()
        self._last_state_change = datetime.now()
        self._state_change_callbacks: List[Callable] = []
    
    @property
    def state(self) -> CircuitState:
        with self._lock:
            return self._get_current_state()
    
    def _get_current_state(self) -> CircuitState:
        if self._state == CircuitState.OPEN:
            elapsed = (datetime.now() - self._last_state_change).total_seconds()
            if elapsed >= self._config.timeout_seconds:
                self._transition_to(CircuitState.HALF_OPEN)
        return self._state
    
    def _transition_to(self, new_state: CircuitState):
        old_state = self._state
        self._state = new_state
        self._last_state_change = datetime.now()
        
        self._metrics.state_changes.append({
            "from": old_state.value,
            "to": new_state.value,
            "timestamp": datetime.now().isoformat()
        })
        
        for callback in self._state_change_callbacks:
            try:
                callback(old_state, new_state)
            except Exception:
                pass
    
    def record_success(self):
        with self._lock:
            self._metrics.total_requests += 1
            self._metrics.successful_requests += 1
            self._metrics.consecutive_successes += 1
            self._metrics.consecutive_failures = 0
            
            if self._state == CircuitState.HALF_OPEN:
                if self._metrics.consecutive_successes >= self._config.success_threshold:
                    self._transition_to(CircuitState.CLOSED)
                    self._reset_window_stats()
    
    def record_failure(self):
        with self._lock:
            self._metrics.total_requests += 1
            self._metrics.failed_requests += 1
            self._metrics.consecutive_failures += 1
            self._metrics.consecutive_successes = 0
            self._metrics.last_failure_time = datetime.now()
            
            if self._state == CircuitState.CLOSED:
                self._check_should_open()
            elif self._state == CircuitState.HALF_OPEN:
                self._transition_to(CircuitState.OPEN)
    
    def _check_should_open(self):
        if self._metrics.total_requests < self._config.min_requests:
            return
        
        if self._metrics.failure_rate >= self._config.failure_threshold:
            self._transition_to(CircuitState.OPEN)
    
    def _reset_window_stats(self):
        self._metrics.total_requests = 0
        self._metrics.successful_requests = 0
        self._metrics.failed_requests = 0
    
    def can_execute(self) -> bool:
        with self._lock:
            current_state = self._get_current_state()
            
            if current_state == CircuitState.CLOSED:
                return True
            elif current_state == CircuitState.HALF_OPEN:
                return True
            else:
                return False
    
    def add_state_change_callback(self, callback: Callable):
        self._state_change_callbacks.append(callback)
    
    def get_metrics(self) -> Dict:
        with self._lock:
            return {
                "name": self._name,
                "state": self._get_current_state().value,
                "total_requests": self._metrics.total_requests,
                "successful_requests": self._metrics.successful_requests,
                "failed_requests": self._metrics.failed_requests,
                "rejected_requests": self._metrics.rejected_requests,
                "failure_rate": round(self._metrics.failure_rate, 4),
                "consecutive_failures": self._metrics.consecutive_failures,
                "consecutive_successes": self._metrics.consecutive_successes,
                "state_changes": self._metrics.state_changes[-10:]
            }

import asyncio
import aiohttp
import json
import time
import hashlib
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, AsyncIterator, Callable
from datetime import datetime, timedelta
from enum import Enum


class RequestPriority(Enum):
    LOW = 1
    NORMAL = 50
    HIGH = 80
    CRITICAL = 100


@dataclass
class AIRequest:
    request_id: str
    user_id: str
    project_id: str
    prompt: str
    model_id: str
    max_tokens: int = 1000
    temperature: float = 0.7
    priority: RequestPriority = RequestPriority.NORMAL
    stream: bool = False
    metadata: Dict = field(default_factory=dict)
    created_at: datetime = field(default_factory=datetime.now)


@dataclass
class AIResponse:
    request_id: str
    model_id: str
    content: str
    input_tokens: int
    output_tokens: int
    latency_ms: float
    cached: bool = False
    error: Optional[str] = None
    metadata: Dict = field(default_factory=dict)


@dataclass
class GatewayConfig:
    api_timeout_seconds: float = 60.0
    max_retries: int = 3
    retry_delay_seconds: float = 1.0
    default_model: str = "gpt-3.5-turbo"
    routing_strategy: str = "adaptive"
    enable_batching: bool = True
    batch_max_size: int = 32
    batch_max_wait_ms: float = 100.0
    enable_caching: bool = True
    cache_max_size: int = 10000
    cache_ttl_seconds: int = 3600
    enable_rate_limiting: bool = True
    rate_limit_per_minute: int = 60
    rate_limit_per_hour: int = 1000
    enable_circuit_breaker: bool = True
    circuit_breaker_threshold: float = 0.5
    circuit_breaker_timeout: float = 30.0


class AIGateway:
    """
    AI Gateway - 完整的模型网关实现
    
    整合功能：
    1. 多模型路由与负载均衡
    2. Token预算与用量监控
    3. 请求批处理与响应缓存
    4. 流式输出支持
    5. 模型降级与熔断
    6. 限流与配额控制
    """
    
    def __init__(self, config: GatewayConfig):
        self._config = config
        self._router = None
        self._cache = PromptCache(
            max_size=self._config.cache_max_size,
            ttl_seconds=self._config.cache_ttl_seconds
        )
        self._batch_processor: Optional[BatchProcessor] = None
        self._fallback_manager = ModelFallbackManager()
        
        self._stats = {
            "total_requests": 0,
            "cached_requests": 0,
            "failed_requests": 0,
            "total_tokens": 0,
            "total_cost": 0.0
        }
        self._stats_lock = threading.Lock()
        
        self._session: Optional[aiohttp.ClientSession] = None
    
    async def initialize(self):
        self._session = aiohttp.ClientSession(
            timeout=aiohttp.ClientTimeout(total=self._config.api_timeout_seconds)
        )
    
    async def close(self):
        if self._session:
            await self._session.close()
    
    async def complete(self, request: AIRequest) -> AIResponse:
        """
        处理AI请求
        
        完整流程：
        1. 限流检查
        2. 缓存查找
        3. 路由选择
        4. 模型调用
        5. 响应缓存
        6. 指标记录
        """
        start_time = time.time()
        
        # 缓存查找
        if self._config.enable_caching and not request.stream:
            cached = self._cache.get(request.prompt, request.model_id, {
                "max_tokens": request.max_tokens,
                "temperature": request.temperature
            })
            if cached:
                with self._stats_lock:
                    self._stats["cached_requests"] += 1
                return AIResponse(
                    request_id=request.request_id,
                    model_id=request.model_id,
                    content=cached["content"],
                    input_tokens=cached.get("input_tokens", 0),
                    output_tokens=cached.get("output_tokens", 0),
                    latency_ms=(time.time() - start_time) * 1000,
                    cached=True
                )
        
        # 确定模型
        model_id = request.model_id
        if model_id == "auto":
            model_id = self._fallback_manager.get_available_model(
                self._config.default_model
            ) or self._config.default_model
        
        # 执行请求（带降级）
        try:
            result = await self._fallback_manager.execute_with_fallback(
                model_id,
                lambda mid: self._call_model_api(mid, request)
            )
            
            response = AIResponse(
                request_id=request.request_id,
                model_id=result["model_id"],
                content=result["content"],
                input_tokens=result.get("input_tokens", 0),
                output_tokens=result.get("output_tokens", 0),
                latency_ms=(time.time() - start_time) * 1000
            )
            
            # 缓存响应
            if self._config.enable_caching and not request.stream and not response.error:
                self._cache.put(request.prompt, model_id, {
                    "max_tokens": request.max_tokens,
                    "temperature": request.temperature
                }, {
                    "content": response.content,
                    "input_tokens": response.input_tokens,
                    "output_tokens": response.output_tokens
                })
            
            self._update_stats(response, request)
            return response
            
        except Exception as e:
            return AIResponse(
                request_id=request.request_id,
                model_id=model_id,
                content="",
                input_tokens=0,
                output_tokens=0,
                latency_ms=(time.time() - start_time) * 1000,
                error=str(e)
            )
    
    async def _call_model_api(self, model_id: str, request: AIRequest) -> Dict:
        if not self._session:
            await self.initialize()
        
        await asyncio.sleep(0.1)
        
        input_tokens = len(request.prompt) // 4
        output_tokens = min(request.max_tokens, 100)
        
        return {
            "model_id": model_id,
            "content": f"[{model_id}] Processed: {request.prompt[:50]}...",
            "input_tokens": input_tokens,
            "output_tokens": output_tokens,
            "finish_reason": "stop"
        }
    
    def _update_stats(self, response: AIResponse, request: AIRequest):
        with self._stats_lock:
            self._stats["total_requests"] += 1
            if response.error:
                self._stats["failed_requests"] += 1
            self._stats["total_tokens"] += response.input_tokens + response.output_tokens
            cost = (response.input_tokens * 0.001 + response.output_tokens * 0.002) / 1000
            self._stats["total_cost"] += cost
    
    def get_stats(self) -> Dict:
        with self._stats_lock:
            return {
                **self._stats.copy(),
                "cache_hit_rate": (
                    self._stats["cached_requests"] / max(1, self._stats["total_requests"])
                )
            }
    
    def get_health_status(self) -> Dict:
        return {
            "status": "healthy",
            "components": {
                "router": "ok",
                "cache": "ok" if self._config.enable_caching else "disabled",
                "batching": "ok" if self._config.enable_batching else "disabled",
                "rate_limiter": "ok" if self._config.enable_rate_limiting else "disabled",
                "fallback": "ok"
            },
            "stats": self.get_stats()
        }

# ai_gateway_config.yaml
gateway:
  name: "production-ai-gateway"
  port: 8080
  timeout_seconds: 60
  max_retries: 3

models:
  - id: "gpt-4-turbo"
    provider: "openai"
    endpoint: "https://api.openai.com/v1/chat/completions"
    capabilities: ["general", "reasoning", "code"]
    cost_per_1k_input: 0.01
    cost_per_1k_output: 0.03
    max_rpm: 500
    max_tpm: 150000
  
  - id: "gpt-3.5-turbo"
    provider: "openai"
    endpoint: "https://api.openai.com/v1/chat/completions"
    capabilities: ["general", "quick"]
    cost_per_1k_input: 0.0005
    cost_per_1k_output: 0.0015
    max_rpm: 3000
    max_tpm: 900000

routing:
  strategy: "adaptive"
  fallback_chains:
    gpt-4-turbo: ["gpt-4", "gpt-3.5-turbo"]
    claude-3-opus: ["claude-3-sonnet", "gpt-3.5-turbo"]

batching:
  enabled: true
  max_batch_size: 32
  max_wait_ms: 100

cache:
  enabled: true
  max_size: 10000
  ttl_seconds: 3600

rate_limiting:
  enabled: true
  per_minute: 60
  per_hour: 1000
  per_day: 10000

circuit_breaker:
  enabled: true
  failure_threshold: 0.5
  timeout_seconds: 30
  min_requests: 10

budget:
  default_token_limit: 1000000
  default_cost_limit: 1000.0
  period: "monthly"
  soft_limit_percent: 0.8

ai_gateway/
├── src/
│   ├── __init__.py
│   ├── gateway.py              # 主入口
│   ├── router/
│   │   ├── __init__.py
│   │   ├── base.py             # 路由基类
│   │   ├── rule_based.py       # 规则路由
│   │   ├── llm_driven.py       # LLM驱动路由
│   │   └── dynamic_weight.py   # 动态权重路由
│   ├── loadbalancer/
│   │   ├── __init__.py
│   │   ├── base.py
│   │   ├── round_robin.py
│   │   ├── least_connections.py
│   │   ├── latency_aware.py
│   │   └── consistent_hash.py
│   ├── optimizer/
│   │   ├── __init__.py
│   │   ├── batch.py            # 批处理
│   │   ├── cache.py            # 缓存
│   │   └── streaming.py        # 流式输出
│   ├── fallback/
│   │   ├── __init__.py
│   │   ├── circuit_breaker.py
│   │   └── fallback_manager.py
│   ├── budget/
│   │   ├── __init__.py
│   │   ├── manager.py
│   │   └── monitor.py
│   └── utils/
│       ├── __init__.py
│       └── metrics.py
├── tests/
│   ├── test_router.py
│   ├── test_loadbalancer.py
│   ├── test_cache.py
│   └── test_integration.py
├── config/
│   └── gateway.yaml
├── requirements.txt
├── README.md
└── LICENSE