AI成本预测器

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI 成本预测器（ai-cost-predictor）核心逻辑 v0.1
定位：多模型成本预算建模，输入调用量/模型分布/缓存率 -> 输出月度成本区间 + 预算里程碑 + 调价预警 + 省钱建议。
依赖：仅标准库。价格数据以各模型官方定价页为准，发布前需核验。
"""

import json
from datetime import date

CURRENCY_SYMBOL = {"USD": "$", "CNY": "¥"}

# 内置价格表：每百万 token 单价（input / output），单位见 currency
# 【数据核验红线】以下条目来自 2026-11-03 官方定价页快照，发布前必须与各厂商官方定价页逐条核验；
# 标注「*未核验」的条目不保证准确；价格表缺失或过期的模型，一律按"离线降级"处理，绝不使用臆造价格。
# 来源：Anthropic(anthropic.com/pricing)/OpenAI(openai.com)/Google(ai.google.dev)/DeepSeek(api-docs.deepseek.com 2026-11-03)/
#       Kimi(platform.moonshot.ai)/智谱(open.bigmodel.cn)/通义(阿里云百炼)/混元(腾讯云 TokenHub)
PRICE_TABLE = {
    # ---- Claude（Anthropic 2026-11-03）----
    "claude-haiku-5.5": {
        "currency": "USD", "input": 0.10, "output": 0.50,
        "note": "Claude Haiku 5.5 轻量档（≤10 万 token；缓存读 $0.01/缓存写 $0.125）",
    },
    "claude-haiku-4.5": {
        "currency": "USD", "input": 1.0, "output": 5.0,
        "note": "Claude Haiku 4.5；缓存读 $0.10/缓存写 $1.25，Batch 5 折",
    },
    "claude-sonnet-5": {
        "currency": "USD", "input": 2.0, "output": 10.0,
        "note": "Claude Sonnet 5（官方页注脚仍标 Introductory $2/$10 through Aug 31 2026; $3/$15 standard thereafter，以官方现状为准）；缓存读 $0.20/写 $2.50",
    },
    "claude-sonnet-5.5": {
        "currency": "USD", "input": 2.0, "output": 10.0,
        "note": "Claude Sonnet 5.5（与 Sonnet 5 同价，官方页未直接展示，待复核）",
    },
    "claude-opus-5": {
        "currency": "USD", "input": 5.0, "output": 25.0,
        "note": "Claude Opus 5；缓存读 $0.50/写 $6.25，Batch 输入 $2.5/输出 $12.5",
    },
    "claude-opus-5.5": {
        "currency": "USD", "input": 4.0, "output": 20.0,
        "note": "Claude Opus 5.5 旗舰（2026-09-22 发布；官方页当前未展示，待复核）",
    },
    "claude-fable-5": {
        "currency": "USD", "input": 10.0, "output": 50.0,
        "note": "Claude Fable 5 最高阶公开模型；缓存读 $1.00/写 $12.50",
    },
    # ---- OpenAI GPT（2026-11-03）----
    "gpt-5.6-sol": {
        "currency": "USD", "input": 4.0, "output": 20.0,
        "note": "GPT-5.6 Sol；缓存输入 $0.40；促销价至少至 2026-11-21",
    },
    "gpt-5.6-terra": {
        "currency": "USD", "input": 2.0, "output": 12.0,
        "note": "GPT-5.6 Terra（均衡型）；缓存输入 $0.20",
    },
    "gpt-5.6-luna": {
        "currency": "USD", "input": 0.2, "output": 1.2,
        "note": "GPT-5.6 Luna（经济档）；缓存输入 $0.02",
    },
    "gpt-6-astra": {
        "currency": "USD", "input": 10.0, "output": 50.0,
        "note": "GPT-6 Astra（1.05M 上下文）",
    },
    "gpt-5.5": {
        "currency": "USD", "input": 12.5, "output": 75.0,
        "note": "GPT-5.5 旗舰（仅短上下文档；官方搜索结果未直接展示，待复核）",
    },
    "gpt-5.4-mini": {
        "currency": "USD", "input": 1.5, "output": 9.0,
        "note": "GPT-5.4 mini",
    },
    "gpt-5-mini": {
        "currency": "USD", "input": 0.45, "output": 3.6,
        "note": "GPT-5 mini；缓存输入 $0.045",
    },
    "gpt-4.1-nano": {
        "currency": "USD", "input": 0.2, "output": 0.8,
        "note": "GPT-4.1 nano；缓存输入 $0.05",
    },
    # ---- Google Gemini ----
    "gemini-3.1-pro": {
        "currency": "USD", "input": 2.0, "output": 12.0,
        "note": "Gemini 3.1 Pro Preview（≤200K）；>200K 输入 $4/输出 $18；2027-01-01 起调价；官方搜索结果未直接展示，待复核",
    },
    "gemini-3.6-flash": {
        "currency": "USD", "input": 0.75, "output": 3.75,
        "note": "Gemini 3.6 Flash（2026-12-31 前促销价；缓存 $0.075）；2027-01-01 起 $1.5/$7.5/$0.15",
    },
    "gemini-3.5-flash-lite": {
        "currency": "USD", "input": 0.3, "output": 2.5,
        "note": "Gemini 3.5 Flash-Lite；缓存 $0.03；官方搜索结果未直接展示，待复核",
    },
    # ---- DeepSeek（2026-11-03 官方 api-docs.deepseek.com；高峰=UTC 01:00-04:00 和 06:00-10:00，其余为闲时，未提周末/假日；闲时=高峰一半；USD）----
    "deepseek-v4-pro": {
        "currency": "USD", "input_peak": 1.32, "output_peak": 3.96,
        "input_offpeak": 0.66, "output_offpeak": 1.98,
        "input_hit_peak": 0.044, "input_hit_offpeak": 0.022,
        "note": "DeepSeek-V4-Pro-0813 峰谷价（Cache Miss/输出）；Cache Hit 高峰 0.044/闲时 0.022",
    },
    "deepseek-v4-flash": {
        "currency": "USD", "input_peak": 0.44, "output_peak": 1.32,
        "input_offpeak": 0.22, "output_offpeak": 0.66,
        "input_hit_peak": 0.014, "input_hit_offpeak": 0.007,
        "note": "DeepSeek-V4-Flash-0731 峰谷价（Cache Miss/输出）；Cache Hit 高峰 0.014/闲时 0.007",
    },
    # ---- 豆包（字节火山方舟，2026-11-03 官方定价页）----
    "doubao-seed-2.1-pro": {
        "currency": "CNY", "input": 3.0, "output": 15.0,
        "note": "Doubao-Seed-2.1-Pro（输入≤1024 档）；缓存命中 ¥1.20；批量 in ¥1.5/out ¥7.5",
    },
    "doubao-seed-2.1-turbo": {
        "currency": "CNY", "input": 3.0, "output": 15.0,
        "note": "Doubao-Seed-2.1-Turbo（260628）；缓存命中 ¥0.60",
    },
    # ---- 百度文心（千帆官方定价页 2026-06-25）----
    "ernie-5.1": {
        "currency": "CNY", "input": 4.0, "output": 18.0,
        "note": "文心 ERNIE 5.1（输入≤32K；32K-128K 输入 ¥6/输出 ¥22）",
    },
    "ernie-5.0": {
        "currency": "CNY", "input": 6.0, "output": 24.0,
        "note": "文心 ERNIE 5.0（输入≤32K；32K-128K 输入 ¥10/输出 ¥40；国际站 $1.4/$5.6）",
    },
    # ---- 阶跃星辰（官方开放平台定价页）----
    "step-3.7-flash": {
        "currency": "CNY", "input": 1.35, "output": 8.1,
        "note": "阶跃星辰 Step-3.7-Flash；缓存命中 ¥0.27",
    },
    # ---- MiniMax（官方定价页 2026-08-20 验证）----
    "minimax-m3": {
        "currency": "USD", "input": 0.30, "output": 1.20,
        "note": "MiniMax M3（≤512K）；缓存命中 $0.06",
    },
    # ---- Mistral（官方定价页）----
    "mistral-large-3": {
        "currency": "USD", "input": 0.5, "output": 1.5,
        "note": "Mistral Large 3（旗舰开源权重）",
    },
    "mistral-medium-3.5": {
        "currency": "USD", "input": 1.5, "output": 7.5,
        "note": "Mistral Medium 3.5（Agentic coding）",
    },
    # ---- 国内主流 ----
    "kimi-k3": {
        "currency": "USD", "input": 3.0, "output": 15.0,
        "note": "Kimi K3 旗舰；缓存命中 $0.30",
    },
    "kimi-k2.6": {
        "currency": "USD", "input": 0.95, "output": 4.0,
        "note": "Kimi K2.6；缓存命中 $0.19/$0.30",
    },
    "glm-5.2": {
        "currency": "CNY", "input": 8.0, "output": 28.0,
        "note": "智谱 GLM-5.2 旗舰；按请求输入长度分档，缓存命中 ¥2",
    },
    "glm-4.7-flash": {
        "currency": "CNY", "input": 0.0, "output": 0.0,
        "note": "智谱 GLM-4.7-Flash 免费档（200K 上下文）",
    },
    "qwen3.8-max": {
        "currency": "USD", "input": 2.0, "output": 6.0,
        "note": "通义千问 qwen3.8-max 国际档；错峰 2 折/忙时 5 折",
    },
    "qwen3.5-plus": {
        "currency": "CNY", "input": 0.96, "output": 5.76,
        "note": "通义千问 qwen3.5-plus 中国内地档；按 128K/256K/1M 分档",
    },
    "hunyuan-hy3": {
        "currency": "CNY", "input": 1.0, "output": 4.0,
        "note": "腾讯混元 Hy3；缓存命中 ¥0.25",
    },
}

DEFAULT_CACHE_HIT = 0.20   # 默认缓存命中率
DEFAULT_BUFFER = 0.30      # 默认重试/峰值缓冲系数（悲观档额外计入）
OPTIMISTIC_FACTOR = 0.80   # 乐观档系数
PESSIMISTIC_FACTOR = 1.20  # 悲观档系数（含长输出、重试、分布偏差）


def _price_lookup(model: str):
    """从价格表取模型条目；缺失返回 None（离线降级）。"""
    return PRICE_TABLE.get(model)


def _per_call_cost(model_price, avg_in_tokens, avg_out_tokens, cache_hit):
    """单次调用成本（按输入/输出 token 计，未含缓冲）。价格缺失时返回 None。"""
    if not model_price:
        return None
    currency = model_price.get("currency", "USD")
    degrade_reason = None

    if "input_peak" in model_price:
        # 峰谷计价模型：默认按高峰价保守估算（支持 Cache Hit 单价）
        in_price = model_price["input_peak"]
        out_price = model_price["output_peak"]
        hit_price = model_price.get("input_hit_peak")
        degrade_reason = "按高峰价保守估算（降级精度），低谷价约为一半"
        if hit_price is not None:
            cost = avg_in_tokens / 1_000_000 * (in_price * (1 - cache_hit) + hit_price * cache_hit) + avg_out_tokens / 1_000_000 * out_price
        else:
            cost = avg_in_tokens / 1_000_000 * in_price * (1 - cache_hit) + avg_out_tokens / 1_000_000 * out_price
    elif model_price.get("input") is None:
        # 输入价缺失：仅按输出价估算并标记降级
        out_price = model_price.get("output_peak") or model_price.get("output")
        if out_price is None:
            return None
        cost = avg_out_tokens / 1_000_000 * out_price
        degrade_reason = "输入价未核验，仅按输出价估算（降级精度）"
    else:
        in_price = model_price["input"]
        out_price = model_price.get("output")
        if out_price is None:
            return None
        cost = avg_in_tokens / 1_000_000 * in_price * (1 - cache_hit) + avg_out_tokens / 1_000_000 * out_price

    return {"cost": cost, "currency": currency,
            "degraded": degrade_reason is not None, "degrade_reason": degrade_reason}


def monthly_cost(daily_calls=0, model_dist=None, avg_in_tokens=0, avg_out_tokens=0,
                 cache_hit=DEFAULT_CACHE_HIT, buffer=DEFAULT_BUFFER):
    """
    计算月度成本区间（乐观/预期/悲观三档）。
    model_dist: {model: 占比}，占比合计应约等于 1（内部归一化）。
    返回 dict；任一模型价格缺失时进入 degraded 状态并在 warnings 中说明。
    """
    if daily_calls <= 0:
        return {"ok": True, "monthly_calls": 0, "expected": 0.0, "optimistic": 0.0,
                "pessimistic": 0.0, "currency": "USD", "milestones": {},
                "warnings": [], "degraded": False}

    if not model_dist:
        return {"ok": False, "error": "model_dist 不能为空", "warnings": []}

    total_share = sum(model_dist.values())
    if total_share <= 0:
        return {"ok": False, "error": "model_dist 占比合计必须大于 0", "warnings": []}

    warnings = []
    degraded = False
    per_model = {}
    currency_totals = {}  # 按币种聚合的月成本（未含缓冲）

    for model, share in model_dist.items():
        share = share / total_share  # 归一化
        price = _price_lookup(model)
        unit = _per_call_cost(price, avg_in_tokens, avg_out_tokens, cache_hit)
        if unit is None:
            warnings.append(f"模型 {model} 价格数据缺失或不可用（离线降级），已跳过该模型成本估算")
            degraded = True
            continue
        if unit["degraded"]:
            warnings.append(f"模型 {model}：{unit.get('degrade_reason') or '价格数据降级精度'}")
            degraded = True
        per_model[model] = {"cost": unit["cost"], "currency": unit["currency"]}
        currency_totals[unit["currency"]] = currency_totals.get(unit["currency"], 0.0) + 
            unit["cost"] * share * daily_calls * 30

    if not currency_totals:
        # 全部模型价格缺失：不产出虚假成本
        return {"ok": True, "monthly_calls": daily_calls * 30, "expected": 0.0,
                "optimistic": 0.0, "pessimistic": 0.0, "currency": "USD",
                "currency_totals": {}, "milestones": {},
                "warnings": warnings, "degraded": True, "per_model": per_model}

    # 主币种 = 金额最大的币种；混合币种不强行相加，分列展示
    main_currency = max(currency_totals, key=currency_totals.get)
    monthly_expected = currency_totals[main_currency] * (1 + buffer)
    optimistic = monthly_expected * OPTIMISTIC_FACTOR
    pessimistic = monthly_expected * PESSIMISTIC_FACTOR

    milestones = {}
    daily_burn = monthly_expected / (30 * (1 + buffer))  # 未含缓冲的日消耗
    for pct in (0.5, 0.8, 1.0):
        if daily_burn > 0:
            day = math_ceil(monthly_expected * pct / daily_burn)
            milestones[f"{int(pct * 100)}%"] = day
        else:
            milestones[f"{int(pct * 100)}%"] = 0

    currency_totals_out = {c: round(v * (1 + buffer), 2) for c, v in currency_totals.items()}

    return {"ok": True, "monthly_calls": daily_calls * 30, "expected": round(monthly_expected, 2),
            "optimistic": round(optimistic, 2), "pessimistic": round(pessimistic, 2),
            "currency": main_currency, "currency_totals": currency_totals_out,
            "milestones": milestones,
            "warnings": warnings, "degraded": degraded, "per_model": per_model}


def math_ceil(x):
    import math
    return math.ceil(x)


def price_change_alerts(today=None):
    """调价/政策变更预警。返回即将生效的价格冲击清单。"""
    today = today or date.today()
    alerts = []
    # 注：Claude Sonnet 5 官方页注脚仍标 "Introductory $2/$10 through Aug 31 2026; $3/$15 standard thereafter"，
    # 与"涨价取消"说法冲突，以官方页现状为准、不写死转正结论，暂不产生到期涨价预警。
    # 后续如需新增调价预警，在此按"model/event/impact/action"结构追加。
    return alerts


def savings_advice(result, avg_in_tokens=0, avg_out_tokens=0, cache_hit=DEFAULT_CACHE_HIT):
    """基于估算结果给出省钱建议（规则模板）。"""
    advice = []
    if result.get("ok"):
        exp = result["expected"]
        if cache_hit < 0.5:
            advice.append(f"当前缓存命中率 {cache_hit:.0%}，提升至 50% 可显著降低输入成本（输入按 1-命中率 计费）")
        if avg_in_tokens and avg_out_tokens and avg_in_tokens > avg_out_tokens * 4:
            advice.append("输入 token 占比过高，建议精简上下文与 Prompt Caching 复用，减少重复发送")
        advice.append("如按峰值价计费的模型占比高，可将非实时任务错峰到闲时执行（如 DeepSeek V4-Pro 闲时输出低 50%）")
        advice.append("建议为预算设置 70%/80%/90% 使用率分级告警，触发后人工复核")
    return advice


def run_prediction(params: dict) -> dict:
    """统一入口：接收参数 dict，返回结构化预测结果。"""
    result = monthly_cost(
        daily_calls=params.get("daily_calls", 0),
        model_dist=params.get("model_dist"),
        avg_in_tokens=params.get("avg_in_tokens", 0),
        avg_out_tokens=params.get("avg_out_tokens", 0),
        cache_hit=params.get("cache_hit", DEFAULT_CACHE_HIT),
        buffer=params.get("buffer", DEFAULT_BUFFER),
    )
    result["alerts"] = price_change_alerts()
    result["advice"] = savings_advice(result, params.get("avg_in_tokens", 0),
                                      params.get("avg_out_tokens", 0),
                                      params.get("cache_hit", DEFAULT_CACHE_HIT))
    return result


if __name__ == "__main__":
    import sys
    demo = {
        "daily_calls": 1000,
        "model_dist": {"claude-sonnet-5": 1.0},
        "avg_in_tokens": 2000,
        "avg_out_tokens": 500,
        "cache_hit": 0.2,
        "buffer": 0.3,
    }
    if len(sys.argv) > 1:
        with open(sys.argv[1], "r", encoding="utf-8") as f:
            demo = json.load(f)
    print(json.dumps(run_prediction(demo), ensure_ascii=False, indent=2))