flexllm 0.3.3__py3-none-any.whl

flexllm/pricing/updater.py

@@ -0,0 +1,223 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+模型定价自动更新脚本
+
+从 OpenRouter API 获取所有模型定价，更新到 pricing.json
+
+使用方法:
+    # 预览更新内容
+    python -m flexllm.pricing.updater
+
+    # 直接更新 pricing.json
+    python -m flexllm.pricing.updater --apply
+
+    # 输出 JSON 格式
+    python -m flexllm.pricing.updater --json
+"""
+
+import argparse
+import json
+import re
+import urllib.request
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+# OpenRouter API 端点
+OPENROUTER_API = "https://openrouter.ai/api/v1/models"
+
+# 定价文件路径
+PRICING_FILE = Path(__file__).parent / "data.json"
+
+# 排除的模型模式
+EXCLUDE_PATTERNS = [
+    r":free$",       # 免费模型
+    r":floor$",      # floor 模型
+    r":extended$",   # extended 模型
+]
+
+
+def fetch_models() -> List[Dict]:
+    """从 OpenRouter API 获取模型列表"""
+    req = urllib.request.Request(
+        OPENROUTER_API,
+        headers={"User-Agent": "flexllm-pricing-updater/1.0"}
+    )
+    with urllib.request.urlopen(req, timeout=30) as resp:
+        data = json.loads(resp.read().decode("utf-8"))
+    return data.get("data", [])
+
+
+def should_exclude(model_id: str) -> bool:
+    """检查模型是否应该排除"""
+    for pattern in EXCLUDE_PATTERNS:
+        if re.search(pattern, model_id):
+            return True
+    return False
+
+
+def normalize_model_id(model_id: str) -> Optional[str]:
+    """
+    将 OpenRouter 模型 ID 规范化
+
+    规则: 直接取斜杠后面的部分
+    例如: openai/gpt-4o -> gpt-4o
+          anthropic/claude-3.5-sonnet -> claude-3.5-sonnet
+    """
+    if should_exclude(model_id):
+        return None
+
+    # 移除 :thinking 等后缀
+    clean_id = re.sub(r":\w+$", "", model_id)
+
+    # 取斜杠后面的部分
+    if "/" in clean_id:
+        return clean_id.split("/", 1)[1]
+
+    return clean_id
+
+
+def parse_pricing(model: Dict) -> Optional[Tuple[float, float]]:
+    """
+    解析模型定价信息
+
+    Returns:
+        (input_price, output_price) 单位: $/1M tokens
+    """
+    pricing = model.get("pricing", {})
+
+    prompt_price = pricing.get("prompt", "0")
+    completion_price = pricing.get("completion", "0")
+
+    try:
+        # OpenRouter 返回的是 $/token，转换为 $/1M tokens
+        input_price = float(prompt_price) * 1e6
+        output_price = float(completion_price) * 1e6
+
+        # 过滤免费模型和价格异常的模型
+        if input_price == 0 and output_price == 0:
+            return None
+        if input_price < 0 or output_price < 0:
+            return None
+
+        return (round(input_price, 4), round(output_price, 4))
+    except (ValueError, TypeError):
+        return None
+
+
+def collect_pricing() -> Dict[str, Dict[str, float]]:
+    """
+    收集所有模型的定价信息
+
+    Returns:
+        {model_name: {"input": price_per_1m, "output": price_per_1m}}
+    """
+    models = fetch_models()
+    pricing_map = {}
+
+    for model in models:
+        model_id = model.get("id", "")
+        name = normalize_model_id(model_id)
+
+        if not name:
+            continue
+
+        pricing = parse_pricing(model)
+        if not pricing:
+            continue
+
+        input_price, output_price = pricing
+
+        # 如果同一模型有多个版本，保留价格最低的
+        if name in pricing_map:
+            existing = pricing_map[name]
+            if input_price >= existing["input"]:
+                continue
+
+        pricing_map[name] = {
+            "input": input_price,
+            "output": output_price,
+        }
+
+    return pricing_map
+
+
+def update_pricing_file(pricing_map: Dict[str, Dict[str, float]]) -> bool:
+    """
+    更新 data.json 文件
+
+    Args:
+        pricing_map: {model_name: {"input": price, "output": price}}
+
+    Returns:
+        是否更新成功
+    """
+    # 按模型名排序
+    sorted_models = dict(sorted(pricing_map.items()))
+
+    data = {
+        "_meta": {
+            "updated_at": datetime.now().strftime("%Y-%m-%d"),
+            "source": "OpenRouter API (https://openrouter.ai/api/v1/models)",
+            "unit": "$/1M tokens",
+        },
+        "models": sorted_models,
+    }
+
+    try:
+        with open(PRICING_FILE, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, ensure_ascii=False)
+        return True
+    except Exception as e:
+        print(f"Error writing {PRICING_FILE}: {e}")
+        return False
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="从 OpenRouter API 更新模型定价表"
+    )
+    parser.add_argument(
+        "--apply",
+        action="store_true",
+        help="直接更新 data.json（默认只预览）",
+    )
+    parser.add_argument(
+        "--json",
+        action="store_true",
+        help="输出 JSON 格式",
+    )
+    args = parser.parse_args()
+
+    print("Fetching models from OpenRouter API...")
+    pricing_map = collect_pricing()
+    print(f"Found {len(pricing_map)} models with pricing info")
+
+    if args.json:
+        print(json.dumps(pricing_map, indent=2, ensure_ascii=False))
+        return
+
+    if args.apply:
+        print(f"\nUpdating {PRICING_FILE}...")
+        if update_pricing_file(pricing_map):
+            print(f"✓ Successfully updated data.json ({len(pricing_map)} models)")
+        else:
+            print("✗ Failed to update data.json")
+            exit(1)
+    else:
+        print("\n" + "=" * 60)
+        print("Preview (use --apply to update data.json):")
+        print("=" * 60 + "\n")
+
+        for name in sorted(pricing_map.keys())[:30]:
+            p = pricing_map[name]
+            print(f"  {name:<40} ${p['input']:<10.4f} / ${p['output']:<10.4f}")
+
+        if len(pricing_map) > 30:
+            print(f"  ... and {len(pricing_map) - 30} more models")
+
+
+if __name__ == "__main__":
+    main()

flexllm/provider_router.py

@@ -0,0 +1,213 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+多 Provider 负载均衡和故障转移
+
+支持多个 API endpoint 的轮询、加权分配和自动 fallback。
+"""
+
+import time
+import random
+from dataclasses import dataclass, field
+from typing import List, Optional, Literal
+from threading import Lock
+
+
+Strategy = Literal["round_robin", "weighted", "random", "fallback"]
+
+
+@dataclass
+class ProviderConfig:
+    """
+    单个 Provider 配置
+
+    Attributes:
+        base_url: API 基础 URL
+        api_key: API 密钥
+        weight: 权重 (用于 weighted 策略)
+        model: 可选的模型覆盖
+        enabled: 是否启用
+    """
+    base_url: str
+    api_key: str = "EMPTY"
+    weight: float = 1.0
+    model: Optional[str] = None
+    enabled: bool = True
+
+
+@dataclass
+class ProviderStatus:
+    """Provider 运行时状态"""
+    config: ProviderConfig
+    failures: int = 0
+    last_failure: float = 0
+    is_healthy: bool = True
+
+
+class ProviderRouter:
+    """
+    Provider 路由器
+
+    支持策略:
+        - round_robin: 轮询
+        - weighted: 加权随机
+        - random: 随机
+        - fallback: 主备模式 (只有主挂了才用备)
+    """
+
+    def __init__(
+        self,
+        providers: List[ProviderConfig],
+        strategy: Strategy = "round_robin",
+        failure_threshold: int = 3,
+        recovery_time: float = 60.0,
+    ):
+        """
+        初始化路由器
+
+        Args:
+            providers: Provider 配置列表
+            strategy: 路由策略
+            failure_threshold: 连续失败多少次后标记为不健康
+            recovery_time: 不健康后多久尝试恢复 (秒)
+        """
+        if not providers:
+            raise ValueError("至少需要一个 provider")
+
+        self.strategy = strategy
+        self.failure_threshold = failure_threshold
+        self.recovery_time = recovery_time
+
+        self._providers = [
+            ProviderStatus(config=p) for p in providers if p.enabled
+        ]
+        self._index = 0
+        self._lock = Lock()
+
+        if not self._providers:
+            raise ValueError("没有可用的 provider")
+
+    def _get_healthy_providers(self) -> List[ProviderStatus]:
+        """获取健康的 provider 列表"""
+        now = time.time()
+        healthy = []
+
+        for p in self._providers:
+            # 尝试恢复
+            if not p.is_healthy and (now - p.last_failure) > self.recovery_time:
+                p.is_healthy = True
+                p.failures = 0
+
+            if p.is_healthy:
+                healthy.append(p)
+
+        return healthy if healthy else self._providers  # 全挂时返回所有
+
+    def get_next(self) -> ProviderConfig:
+        """
+        获取下一个可用的 provider
+
+        Returns:
+            ProviderConfig
+        """
+        with self._lock:
+            healthy = self._get_healthy_providers()
+
+            if self.strategy == "round_robin":
+                self._index = (self._index + 1) % len(healthy)
+                return healthy[self._index].config
+
+            elif self.strategy == "weighted":
+                weights = [p.config.weight for p in healthy]
+                total = sum(weights)
+                r = random.uniform(0, total)
+                cumsum = 0
+                for p in healthy:
+                    cumsum += p.config.weight
+                    if r <= cumsum:
+                        return p.config
+                return healthy[-1].config
+
+            elif self.strategy == "random":
+                return random.choice(healthy).config
+
+            elif self.strategy == "fallback":
+                # 优先使用第一个健康的
+                return healthy[0].config
+
+            else:
+                return healthy[0].config
+
+    def mark_failed(self, provider: ProviderConfig) -> None:
+        """
+        标记 provider 失败
+
+        Args:
+            provider: 失败的 provider 配置
+        """
+        with self._lock:
+            for p in self._providers:
+                if p.config.base_url == provider.base_url:
+                    p.failures += 1
+                    p.last_failure = time.time()
+                    if p.failures >= self.failure_threshold:
+                        p.is_healthy = False
+                    break
+
+    def mark_success(self, provider: ProviderConfig) -> None:
+        """
+        标记 provider 成功，重置失败计数
+
+        Args:
+            provider: 成功的 provider 配置
+        """
+        with self._lock:
+            for p in self._providers:
+                if p.config.base_url == provider.base_url:
+                    p.failures = 0
+                    p.is_healthy = True
+                    break
+
+    def get_all_healthy(self) -> List[ProviderConfig]:
+        """获取所有健康的 provider"""
+        with self._lock:
+            return [p.config for p in self._get_healthy_providers()]
+
+    @property
+    def stats(self) -> dict:
+        """返回路由器统计信息"""
+        with self._lock:
+            return {
+                "total": len(self._providers),
+                "healthy": sum(1 for p in self._providers if p.is_healthy),
+                "strategy": self.strategy,
+                "providers": [
+                    {
+                        "base_url": p.config.base_url,
+                        "healthy": p.is_healthy,
+                        "failures": p.failures,
+                    }
+                    for p in self._providers
+                ],
+            }
+
+
+def create_router_from_urls(
+    urls: List[str],
+    api_key: str = "EMPTY",
+    strategy: Strategy = "round_robin",
+) -> ProviderRouter:
+    """
+    便捷函数：从 URL 列表创建路由器
+
+    Args:
+        urls: API URL 列表
+        api_key: 统一的 API 密钥
+        strategy: 路由策略
+
+    Returns:
+        ProviderRouter 实例
+    """
+    providers = [ProviderConfig(base_url=url, api_key=api_key) for url in urls]
+    return ProviderRouter(providers, strategy=strategy)

flexllm/token_counter.py

@@ -0,0 +1,270 @@
+#! /usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+"""
+Token 计数和成本估算模块
+
+支持使用 tiktoken 精确计算，或在缺失时使用估算方法。
+定价数据从 pricing/data.json 加载，可通过 `flexllm pricing --update` 更新。
+"""
+
+import hashlib
+import json
+from typing import Union, List, Dict, Any, Optional
+
+# tiktoken 是可选依赖
+try:
+    import tiktoken
+    TIKTOKEN_AVAILABLE = True
+except ImportError:
+    TIKTOKEN_AVAILABLE = False
+
+# 从 pricing 模块导入定价功能
+from .pricing import get_pricing, reload_pricing, estimate_cost as _estimate_cost
+
+# 兼容旧 API：MODEL_PRICING 现在是动态获取的
+def _get_model_pricing():
+    return get_pricing()
+
+# 为了向后兼容，保留 MODEL_PRICING 变量
+MODEL_PRICING = _get_model_pricing()
+
+# 模型到 tiktoken 编码器的映射
+MODEL_TO_ENCODING = {
+    # GPT-5 系列
+    "gpt-5": "o200k_base",
+    "gpt-5.1": "o200k_base",
+    "gpt-5.1-codex": "o200k_base",
+    "gpt-5.2": "o200k_base",
+    "gpt-5.2-pro": "o200k_base",
+    # GPT-4o 系列
+    "gpt-4o": "o200k_base",
+    "gpt-4o-mini": "o200k_base",
+    # GPT-4.1 系列
+    "gpt-4.1": "o200k_base",
+    "gpt-4.1-mini": "o200k_base",
+    "gpt-4.1-nano": "o200k_base",
+    # GPT-4 系列
+    "gpt-4-turbo": "cl100k_base",
+    "gpt-4": "cl100k_base",
+    "gpt-3.5-turbo": "cl100k_base",
+    # o 系列推理模型
+    "o1": "o200k_base",
+    "o1-mini": "o200k_base",
+    "o1-pro": "o200k_base",
+    "o3": "o200k_base",
+    "o3-mini": "o200k_base",
+    "o4-mini": "o200k_base",
+}
+
+# 编码器缓存
+_encoding_cache: Dict[str, Any] = {}
+
+
+def _get_encoding(model: str):
+    """获取模型对应的 tiktoken 编码器"""
+    if not TIKTOKEN_AVAILABLE:
+        return None
+
+    encoding_name = MODEL_TO_ENCODING.get(model, "cl100k_base")
+    if encoding_name not in _encoding_cache:
+        _encoding_cache[encoding_name] = tiktoken.get_encoding(encoding_name)
+    return _encoding_cache[encoding_name]
+
+
+def _estimate_tokens_simple(text: str) -> int:
+    """简单估算：中文约 2 字符/token，英文约 4 字符/token"""
+    if not text:
+        return 0
+    # 粗略统计中文字符
+    chinese_chars = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
+    other_chars = len(text) - chinese_chars
+    return chinese_chars // 2 + other_chars // 4 + 1
+
+
+def count_tokens(
+    content: Union[str, List[Dict], Dict],
+    model: str = "gpt-4o"
+) -> int:
+    """
+    计算 token 数量
+
+    Args:
+        content: 文本字符串或 messages 列表
+        model: 模型名称，用于选择正确的 tokenizer
+
+    Returns:
+        token 数量
+    """
+    # 处理 messages 格式
+    if isinstance(content, list):
+        total = 0
+        for msg in content:
+            if isinstance(msg, dict):
+                # 每条消息有固定开销
+                total += 4  # role + content 标记
+                for key, value in msg.items():
+                    if isinstance(value, str):
+                        total += count_tokens(value, model)
+                    elif isinstance(value, list):
+                        # 处理多模态内容
+                        for item in value:
+                            if isinstance(item, dict) and "text" in item:
+                                total += count_tokens(item["text"], model)
+                            elif isinstance(item, dict) and "image_url" in item:
+                                # 图像 token 估算 (低分辨率约 85，高分辨率约 170*tiles)
+                                total += 85
+        return total + 2  # 结束标记
+
+    if isinstance(content, dict):
+        return count_tokens(json.dumps(content, ensure_ascii=False), model)
+
+    # 文本处理
+    text = str(content)
+    encoding = _get_encoding(model)
+
+    if encoding:
+        return len(encoding.encode(text))
+    else:
+        return _estimate_tokens_simple(text)
+
+
+def count_messages_tokens(
+    messages_list: List[List[Dict]],
+    model: str = "gpt-4o"
+) -> int:
+    """
+    批量计算 messages 的 token 总数
+
+    Args:
+        messages_list: messages 列表的列表
+        model: 模型名称
+
+    Returns:
+        总 token 数量
+    """
+    return sum(count_tokens(msgs, model) for msgs in messages_list)
+
+
+def estimate_cost(
+    input_tokens: int,
+    output_tokens: int = 0,
+    model: str = "gpt-4o"
+) -> float:
+    """
+    估算 API 调用成本
+
+    Args:
+        input_tokens: 输入 token 数
+        output_tokens: 输出 token 数 (如果未知可传 0)
+        model: 模型名称
+
+    Returns:
+        估算成本 (美元)
+    """
+    return _estimate_cost(input_tokens, output_tokens, model)
+
+
+def estimate_batch_cost(
+    messages_list: List[List[Dict]],
+    model: str = "gpt-4o",
+    avg_output_tokens: int = 500
+) -> Dict[str, Any]:
+    """
+    估算批量处理的成本
+
+    Args:
+        messages_list: messages 列表的列表
+        model: 模型名称
+        avg_output_tokens: 预估每条请求的平均输出 token 数
+
+    Returns:
+        包含详细估算信息的字典
+    """
+    input_tokens = count_messages_tokens(messages_list, model)
+    output_tokens = len(messages_list) * avg_output_tokens
+    cost = estimate_cost(input_tokens, output_tokens, model)
+
+    return {
+        "count": len(messages_list),
+        "input_tokens": input_tokens,
+        "estimated_output_tokens": output_tokens,
+        "total_tokens": input_tokens + output_tokens,
+        "estimated_cost_usd": round(cost, 4),
+        "model": model,
+    }
+
+
+def _normalize_message_for_hash(message: Dict) -> Dict:
+    """
+    规范化消息用于 hash 计算，将 base64 图片替换为其内容 hash
+
+    这样做的好处：
+    1. 减少 hash 计算的数据量（base64 可能有几 MB）
+    2. 同一张图片即使重新编码也会产生相同的缓存键
+    """
+    if not isinstance(message, dict):
+        return message
+
+    result = {}
+    for key, value in message.items():
+        if key == "content" and isinstance(value, list):
+            # 处理多模态内容（OpenAI 格式）
+            normalized_content = []
+            for item in value:
+                if isinstance(item, dict) and "image_url" in item:
+                    image_url = item["image_url"]
+                    if isinstance(image_url, dict):
+                        url = image_url.get("url", "")
+                    else:
+                        url = str(image_url)
+
+                    # 检查是否是 base64 数据
+                    if url.startswith("data:image"):
+                        # 提取 base64 部分并计算 hash
+                        base64_data = url.split(",", 1)[-1] if "," in url else url
+                        img_hash = hashlib.md5(base64_data.encode()).hexdigest()[:16]
+                        # 用短 hash 替代完整 base64
+                        normalized_item = {
+                            "type": item.get("type", "image_url"),
+                            "image_url": {"url": f"img_hash:{img_hash}"}
+                        }
+                        normalized_content.append(normalized_item)
+                    else:
+                        # URL 类型保持不变
+                        normalized_content.append(item)
+                else:
+                    normalized_content.append(item)
+            result[key] = normalized_content
+        else:
+            result[key] = value
+    return result
+
+
+def messages_hash(
+    messages: List[Dict],
+    model: str = "",
+    **kwargs
+) -> str:
+    """
+    生成 messages 的唯一哈希值，用于缓存键
+
+    Args:
+        messages: 消息列表
+        model: 模型名称
+        **kwargs: 其他影响结果的参数 (temperature, max_tokens 等)
+
+    Returns:
+        MD5 哈希字符串
+    """
+    # 规范化消息（优化 base64 图片的处理）
+    normalized_messages = [_normalize_message_for_hash(m) for m in messages]
+
+    # 构建要哈希的内容
+    cache_key_data = {
+        "messages": normalized_messages,
+        "model": model,
+        **{k: v for k, v in kwargs.items() if v is not None}
+    }
+    content = json.dumps(cache_key_data, sort_keys=True, ensure_ascii=False)
+    return hashlib.md5(content.encode()).hexdigest()

flexllm/utils/__init__.py

@@ -0,0 +1 @@
+from .core import async_retry