"""5.6 推理的两段与缓存（API 版）。三个小实验，都用流式接口计时： 1. 输入变长：输出固定只要几个 token，看首 token 延迟（TTFT）怎么随输入长度涨。这一段主要是 prefill。 2. 输出变长：输入固定很短，看总时间怎么随输出长度涨，算出每个输出 token 要多久。这一段是 decode。 3. 前缀缓存：同一份长文档上连问几轮。一种写法每轮只在末尾追加；另一种在开头插一个当前时间戳。 比较每一轮的首 token 延迟。 用法： export AI_BUILDER_TOKEN=sk_... python api_timing.py # 默认 deepseek，三个实验都跑，约 5 分钟 python api_timing.py --model grok-4-fast python api_timing.py --only cache # 只跑某一个：input / output / cache python api_timing.py --doc my_notes.txt # 缓存实验换成你自己的长文档 """ import argparse import json import os import random import statistics import string import time from datetime import datetime from pathlib import Path from openai import OpenAI BASE_URL = os.getenv('AI_BUILDER_SPACE_BASE_URL', 'https://space.ai-builders.com/backend/v1') PRICES = {'deepseek': (0.15, 0.60), 'grok-4-fast': (1.25, 2.50)} # 美元/百万 token（输入, 输出），2026-10 AI Builder Space 价目 client = None def stream(model, messages, max_tokens): """发一次流式请求，返回 (首 token 秒数, 总秒数, 收到的文本块数)。文本块数约等于输出 token 数。""" t0 = time.time() first, chunks, text = None, 0, [] s = client.chat.completions.create(model=model, messages=messages, max_tokens=max_tokens, stream=True) for ch in s: if ch.choices and ch.choices[0].delta and ch.choices[0].delta.content: if first is None: first = time.time() - t0 chunks += 1 text.append(ch.choices[0].delta.content) return first, time.time() - t0, chunks, ''.join(text) def prompt_tokens(model, messages): """非流式调用一次，只为读 usage 里的输入 token 数（流式响应里没有 usage）。""" r = client.chat.completions.create(model=model, messages=messages, max_tokens=1) return r.usage.prompt_tokens def filler(n_words, seed): """一段随机英文“单词”。每次用不同的 seed，保证没有两次请求的前缀相同，避免命中缓存。""" rnd = random.Random(seed) words = [''.join(rnd.choice(string.ascii_lowercase) for _ in range(rnd.randint(3, 8))) for _ in range(n_words)] lines = [' '.join(words[i:i + 12]) for i in range(0, n_words, 12)] return '\n'.join(f'{i}: {l}' for i, l in enumerate(lines)) def exp_input(model, reps): print('\n## 实验 1：输入变长，输出只要 8 个 token\n') print('| 输入 token | 首 token 延迟（中位数） | 各次 |') print('|---|---|---|') out = [] for n_words in [500, 2500, 10000, 40000]: ttfts = [] for k in range(reps): seed = f'{time.time()}-{k}' msgs = [{'role': 'user', 'content': filler(n_words, seed) + '\n\n上面一共有几行？只回答数字。'}] ttfts.append(stream(model, msgs, 8)[0]) n_tok = prompt_tokens(model, [{'role': 'user', 'content': filler(n_words, 'count') + '\n\n上面一共有几行？只回答数字。'}]) out.append(dict(prompt_tokens=n_tok, ttft=ttfts)) print(f'| {n_tok:,} | {statistics.median(ttfts):.2f} 秒 | {", ".join(f"{x:.2f}" for x in ttfts)} |') return out def exp_output(model, reps): print('\n## 实验 2：输入很短，输出变长\n') print('| 要求的输出 token | 实际收到的块 | 首 token 延迟 | 总时间 | 每个输出 token |') print('|---|---|---|---|---|') out = [] msgs = [{'role': 'user', 'content': '从 1 开始一直往上数，用逗号分隔，不要停，不要说别的。'}] for max_tokens in [50, 200, 800]: runs = [stream(model, msgs, max_tokens) for _ in range(reps)] ttft = statistics.median(r[0] for r in runs) total = statistics.median(r[1] for r in runs) chunks = statistics.median(r[2] for r in runs) per_tok = statistics.median((r[1] - r[0]) / max(r[2] - 1, 1) for r in runs) out.append(dict(max_tokens=max_tokens, ttft=ttft, total=total, chunks=chunks, sec_per_token=per_tok)) print(f'| {max_tokens} | {chunks:.0f} | {ttft:.2f} 秒 | {total:.2f} 秒 | {per_tok * 1000:.1f} 毫秒 |') return out QUESTIONS = ['第 3 行的第一个词是什么？只回答这个词。', '第 10 行一共有几个词？只回答数字。', '第 25 行的最后一个词是什么？只回答这个词。', '第 7 行的第二个词是什么？只回答这个词。', '第 40 行的第一个词是什么？只回答这个词。', '第 12 行的第三个词是什么？只回答这个词。'] def exp_cache(model, rounds, doc_text): print(f'\n## 实验 3：同一份长文档上连问 {rounds} 轮\n') results = {} for mode in ['append_only', 'timestamp_first']: doc = doc_text or filler(30000, f'{mode}-{time.time()}') # 两种写法用不同的文档，互不沾光 history, ttfts = [], [] for i in range(rounds): system = doc if mode == 'append_only' else f'当前时间：{datetime.now().isoformat()}\n\n{doc}' q = {'role': 'user', 'content': QUESTIONS[i % len(QUESTIONS)]} ttft, _, _, reply = stream(model, [{'role': 'system', 'content': system}] + history + [q], 20) history += [q, {'role': 'assistant', 'content': reply}] ttfts.append(ttft) results[mode] = ttfts n_tok = prompt_tokens(model, [{'role': 'system', 'content': doc_text or filler(30000, 'count')}, {'role': 'user', 'content': QUESTIONS[0]}]) print(f'每轮输入约 {n_tok:,} token。\n') print('| 轮次 | 只在末尾追加 | 开头插时间戳 |') print('|---|---|---|') for i in range(rounds): print(f'| {i + 1} | {results["append_only"][i]:.2f} 秒 | {results["timestamp_first"][i]:.2f} 秒 |') a, b = results['append_only'][1:], results['timestamp_first'][1:] print(f'| 第 2 轮起的中位数 | {statistics.median(a):.2f} 秒 | {statistics.median(b):.2f} 秒 |') pin, pout = PRICES.get(model, (0, 0)) print(f'\n按 AI Builder Space 价目，每轮输入约 ${n_tok * pin / 1e6:.4f}，两种写法一样：' '它不单独给缓存命中打折，usage 里也不报命中了多少。缓存在这里省的是时间。') return dict(prompt_tokens=n_tok, **results) def main(): global client ap = argparse.ArgumentParser() ap.add_argument('--model', default='deepseek', help='要支持流式：deepseek、grok-4-fast 可以；kimi、gemini 不行') ap.add_argument('--only', choices=['input', 'output', 'cache']) ap.add_argument('--reps', type=int, default=3) ap.add_argument('--rounds', type=int, default=6) ap.add_argument('--doc', help='缓存实验用的长文档（文本文件）；不给就用随机生成的约 10 万 token 文本') ap.add_argument('--out', default='results') args = ap.parse_args() token = os.getenv('AI_BUILDER_TOKEN') if not token: raise SystemExit('请先设置环境变量 AI_BUILDER_TOKEN（你的 AI Builder Space key）') client = OpenAI(base_url=BASE_URL, api_key=token, timeout=600) doc_text = Path(args.doc).read_text(encoding='utf-8') if args.doc else None print(f'模型 {args.model}，{datetime.now():%Y-%m-%d %H:%M}') res = {'model': args.model} if args.only in (None, 'input'): res['input'] = exp_input(args.model, args.reps) if args.only in (None, 'output'): res['output'] = exp_output(args.model, args.reps) if args.only in (None, 'cache'): res['cache'] = exp_cache(args.model, args.rounds, doc_text) Path(args.out).mkdir(exist_ok=True) path = Path(args.out) / f'api_{args.model}_{datetime.now():%Y%m%d_%H%M}.json' path.write_text(json.dumps(res, ensure_ascii=False, indent=1), encoding='utf-8') print(f'\n原始计时已存到 {path}') if __name__ == '__main__': main()