PlantAI-rag-new/tooling/aqua_kb/cost_model.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""cost_model.py — 水产库(aqua)建库成本校准模型(T 代理,2026-09-12)。

参数化估算 LLM 抽取成本与 MinerU GPU 耗时,输出 Markdown 预算表。
只依赖标准库;所有数值可被命令行覆盖;默认值来源见
kb-libraries/aqua/qa/budget_template.md 假设清单。

用法:
  $PY tooling/aqua_kb/cost_model.py                        # 默认参数,打印预算表
  $PY tooling/aqua_kb/cost_model.py --json                 # 机器可读
  $PY tooling/aqua_kb/cost_model.py --out budget.md        # 同时写入文件
  $PY tooling/aqua_kb/cost_model.py --price-cny-m-input 1.0 --price-cny-m-output 2.0 \
       --md-chars 2109 --pdf-chars 60000 --n-md 31848 --n-pdf 13175

成本公式(人民币):
  in_tokens  = chars / chars_per_token + overhead_tokens * calls
  out_tokens = (chars / chars_per_token) * out_in_ratio
  cost       = (in_tokens * p_in + out_tokens * p_out) / 1e6 * docs
其中 p_in/p_out 为每百万 token 单价。
"""
from __future__ import annotations

import argparse
import json
import sys
from dataclasses import dataclass


@dataclass
class DocSpec:
    label: str
    n_docs: int
    chars_per_doc: float
    llm_calls_per_doc: float
    chunks_per_doc: float | None = None
    cost: float = 0.0
    in_tokens: int = 0
    out_tokens: int = 0


@dataclass
class Params:
    # 单价(¥/百万 token)——官方 deepseek-v4-flash 列表价(缓存未命中口径)
    price_cny_m_input: float = 1.0
    price_cny_m_output: float = 2.0
    # 字符/token 折算
    chars_per_token: float = 4.0      # 英文约 4 字符/token(aqua 语料英文占主导)
    cjk_chars_per_token: float = 1.5  # 中文约 1.5 字符/token(当前语料占比 <5%)
    # LLM 调用假设
    overhead_tokens_per_call: float = 350.0  # 每调用的提示词/模板开销
    out_in_ratio: float = 0.35        # 输出 token / 输入正文 token(P1 试点拟合前先估值)
    md_llm_calls: float = 2.0         # MD 卡:提取 + 复核 各 1 次
    pdf_chunk_chars: float = 8000.0   # PDF 全文分块大小
    pdf_output_per_call: float = 800.0  # PDF 每块调用输出 token(P1 实测前估值)
    # 语料规模
    n_md: int = 31848
    n_pdf: int = 13175
    n_pilot_pdf: int = 30
    n_pilot_md: int = 470
    # 实测/历史锚点
    md_chars: float = 2109.0   # 实测:source/题目摘要_MD 200 篇抽样 mean(P50=2169)
    pdf_chars: float = 60000.0  # 假设:MinerU 输出正文均值(P1 30 篇试点后回填实测)
    mineru_files_per_hour: float = 700.0  # 历史实测:单卡 3 端点并流 ~9,363 篇/11h ≈ 850/h,取保守 700
    mineru_unprocessable: float = 0.177   # 历史实测:1656/9363;不可处理篇不产生抽取成本
    gpu_hours_per_day: float = 24.0
    peak_multiplier: float = 1.0  # 高峰时段单价倍数(官方列表价口径=1.0;若按高峰 ×2 口径计费则设 2.0)


def effective_input_price(p: Params) -> float:
    # 高峰时段(工作日北京时间 09-12/14-18)约占总时长 1/4;
    # 列表价口径(peak_multiplier=1.0)下加权后与列表价一致。
    frac_peak = 0.25
    return p.price_cny_m_input * (1 - frac_peak) + p.price_cny_m_input * p.peak_multiplier * frac_peak


def model_docs(p: Params) -> tuple[list[DocSpec], list[DocSpec]]:
    """返回 (批次[P1 试点], 全量) 两组成本估算。"""
    p_in_eff = effective_input_price(p)

    def md_cost(chars: float, n: int) -> DocSpec:
        body_tokens = chars / p.chars_per_token
        in_t = body_tokens + p.overhead_tokens_per_call * p.md_llm_calls
        out_t = body_tokens * p.out_in_ratio
        c = (in_t * p_in_eff + out_t * p.price_cny_m_output) / 1e6 * n
        return DocSpec("MD 摘要卡", n, chars, p.md_llm_calls, None, c, int(in_t * n), int(out_t * n))

    def pdf_cost(chars: float, n: int) -> DocSpec:
        chunks = max(1.0, round(chars / p.pdf_chunk_chars))
        in_t = (chars / p.chars_per_token) + p.overhead_tokens_per_call * chunks
        out_t = chunks * p.pdf_output_per_call
        c = (in_t * p_in_eff + out_t * p.price_cny_m_output) / 1e6 * n
        return DocSpec("PDF 全文", n, chars, chunks, chunks, c, int(in_t * n), int(out_t * n))

    batch = [md_cost(p.md_chars, p.n_pilot_md), pdf_cost(p.pdf_chars, p.n_pilot_pdf)]
    full = [md_cost(p.md_chars, p.n_md), pdf_cost(p.pdf_chars, p.n_pdf)]
    return batch, full


def render_md(p: Params, batch: list[DocSpec], full: list[DocSpec]) -> str:
    p_in_eff = effective_input_price(p)
    L: list[str] = []
    L.append("# 水产库建库预算表(由 cost_model.py 生成)")
    L.append("")
    L.append("| 参数 | 值 | 来源 |")
    L.append("|---|---|---|")
    L.append(f"| LLM 单价(输入/输出,¥/百万 token) | {p.price_cny_m_input:.2f} / {p.price_cny_m_output:.2f} | 官方 deepseek-v4-flash 列表价(缓存未命中) |")
    L.append(f"| 高峰加权后输入单价 | {p_in_eff:.4f} | peak_multiplier={p.peak_multiplier:.1f}(1.0=列表价口径;若高峰 ×2 计费口径则设 2.0) |")
    L.append(f"| 字符/token | {p.chars_per_token}(英文) / {p.cjk_chars_per_token}(中文) | 英文实测经验;中文 1.5 待试点校准 |")
    L.append(f"| MD 卡平均字符数 | {p.md_chars:,.0f} | 实测:200 篇抽样(P50=2169) |")
    L.append(f"| PDF 正文平均字符数 | {p.pdf_chars:,.0f} | 假设,P1 30 篇试点后回填 |")
    L.append(f"| MD 篇数 / PDF 篇数(全量) | {p.n_md:,} / {p.n_pdf:,} | census.json 实测 |")
    L.append(f"| MinerU 吞吐 | {p.mineru_files_per_hour:,.0f} 篇/小时(单卡 3 端点并流) | 历史实测外推(2026-08-31 批) |")
    L.append("")
    L.append("## P1 试点批次(30 PDF + 470 MD)")
    L.append("")
    L.append("| 文档类型 | 篇数 | LLM 调用/篇 | 输入 token | 输出 token | 成本(¥) |")
    L.append("|---|---|---|---|---|---|")
    tb = sum(s.cost for s in batch)
    for s in batch:
        L.append(f"| {s.label} | {s.n_docs:,} | {s.llm_calls_per_doc:g} | {s.in_tokens:,} | {s.out_tokens:,} | {s.cost:,.2f} |")
    L.append(f"| **合计** | | | | | **{tb:,.2f}** |")
    L.append("")
    mineru_h = p.n_pilot_pdf / p.mineru_files_per_hour
    L.append(f"- MinerU GPU 耗时:30 篇约 {max(mineru_h * 60, 10):.0f} 分钟(含模型加载起步开销;历史不可处理率 {p.mineru_unprocessable:.0%} → 预计成功 ~25 篇)")
    L.append("- GPU:单卡 RTX 5090(32GB),MinerU 本地运行,无 API 成本")
    L.append("")
    L.append("## 全量批次(31,848 MD + 13,175 PDF)")
    L.append("")
    L.append("| 文档类型 | 篇数 | LLM 调用/篇 | 输入 token | 输出 token | 成本(¥) |")
    L.append("|---|---|---|---|---|---|")
    tf = sum(s.cost for s in full)
    for s in full:
        L.append(f"| {s.label} | {s.n_docs:,} | {s.llm_calls_per_doc:g} | {s.in_tokens:,} | {s.out_tokens:,} | {s.cost:,.2f} |")
    L.append(f"| **合计(中值)** | | | | | **{tf:,.2f}** |")
    lo, hi = tf * 0.7, tf * 1.3
    L.append(f"| **合计(±30% 区间)** | | | | | **{lo:,.2f} – {hi:,.2f}** |")
    L.append("")
    days = p.n_pdf / p.mineru_files_per_hour / p.gpu_hours_per_day
    L.append(f"- MinerU 全量 PDF:{p.n_pdf:,} 篇全部过一遍(其中约 {p.mineru_unprocessable:.0%} 判为不可处理),按单卡 3 端点并流历史实测吞吐(保守取 {p.mineru_files_per_hour:,.0f} 篇/时)全天跑 **{days:,.1f} 天**;加失败重试与运维余量,预算按 1–2 天排期")
    L.append("- 成本区间已含 ±30% 不确定性(pdf_chars 假设、out_in_ratio、高峰价、重试等)")
    L.append("- 假设清单与待回填项见 kb-libraries/aqua/qa/budget_template.md")
    L.append("")
    return "\n".join(L)


def main() -> int:
    ap = argparse.ArgumentParser(description="aqua 库成本校准模型")
    for k, v in Params().__dict__.items():
        if isinstance(v, bool):
            continue
        ap.add_argument(f"--{k.replace('_', '-')}", type=float, default=v)
    ap.add_argument("--json", action="store_true", help="输出机器可读 JSON")
    ap.add_argument("--out", default="", help="同时写入预算表 md 文件")
    args = ap.parse_args()
    p = Params(**{k: getattr(args, k) for k in Params().__dict__})
    batch, full = model_docs(p)
    if args.json:
        out = {
            "params": p.__dict__,
            "pilot_batch": {s.label: {"n_docs": s.n_docs, "in_tokens": s.in_tokens,
                                       "out_tokens": s.out_tokens, "cost_cny": round(s.cost, 2)}
                            for s in batch},
            "pilot_total_cny": round(sum(s.cost for s in batch), 2),
            "full_batch": {s.label: {"n_docs": s.n_docs, "in_tokens": s.in_tokens,
                                      "out_tokens": s.out_tokens, "cost_cny": round(s.cost, 2)}
                           for s in full},
            "full_total_cny": round(sum(s.cost for s in full), 2),
            "full_total_range_cny": [round(sum(s.cost for s in full) * 0.7, 2),
                                     round(sum(s.cost for s in full) * 1.3, 2)],
            "mineru_gpu_days_full": round(p.n_pdf / p.mineru_files_per_hour / 24, 2),
        }
        print(json.dumps(out, ensure_ascii=False, indent=2))
        return 0
    md = render_md(p, batch, full)
    print(md)
    if args.out:
        with open(args.out, "w", encoding="utf-8") as f:
            f.write(md)
        print(f"[written] {args.out}", file=sys.stderr)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())