()
| 359 | |
| 360 | def main(): |
| 361 | parser = argparse.ArgumentParser(description="Token 花销统计工具") |
| 362 | parser.add_argument("--start-date", type=str, default=None, help="开始日期 YYYYMMDD") |
| 363 | parser.add_argument("--end-date", type=str, default=None, help="结束日期 YYYYMMDD") |
| 364 | parser.add_argument("--group-by", type=str, default="date", choices=["date", "task_type", "model", "user_id"]) |
| 365 | parser.add_argument("--mode", type=str, default="stats", choices=["stats", "benchmark", "compare"]) |
| 366 | parser.add_argument("--llm-model", type=str, default="qwen3.5-plus", help="LLM 模型名称") |
| 367 | parser.add_argument("--days", type=int, default=1, help="基准测试天数") |
| 368 | args = parser.parse_args() |
| 369 | |
| 370 | # 连接数据库获取论文和用户数据 |
| 371 | conn = sqlite3.connect(DB_PATH) |
| 372 | papers = conn.execute("SELECT title, abstract FROM papers").fetchall() |
| 373 | papers = [{"title": p[0], "abstract": p[1] or ""} for p in papers] |
| 374 | |
| 375 | profiles = conn.execute("SELECT user_id, profile_json FROM profiles").fetchall() |
| 376 | user_profiles = [{"user_id": p[0], **json.loads(p[1])} for p in profiles] |
| 377 | conn.close() |
| 378 | |
| 379 | print(f"Loaded {len(papers)} papers, {len(user_profiles)} users") |
| 380 | |
| 381 | if args.mode == "stats" and args.start_date and args.end_date: |
| 382 | # 统计已有日志 |
| 383 | stats = get_usage_stats(args.start_date, args.end_date, args.group_by) |
| 384 | print(f"\n=== Token Usage Stats ({args.start_date} to {args.end_date}) ===\n") |
| 385 | |
| 386 | for key, s in sorted(stats.items()): |
| 387 | print(f"[{key}]") |
| 388 | print(f" Embedding: {s['embedding_input']:,} tokens") |
| 389 | print(f" LLM: {s['llm_input'] + s['llm_output']:,} tokens") |
| 390 | print(f" Cost: ${s['cost']:.4f}") |
| 391 | print(f" Calls: {s['calls']}") |
| 392 | print() |
| 393 | |
| 394 | elif args.mode == "benchmark": |
| 395 | # 基准测试指定天数 |
| 396 | today = datetime.now() |
| 397 | for i in range(args.days): |
| 398 | test_date = (today - timedelta(days=i)).strftime("%Y-%m-%d") |
| 399 | print(f"\n{'='*60}") |
| 400 | print(f"Benchmarking {test_date} with {args.llm_model}...") |
| 401 | print(f"{'='*60}\n") |
| 402 | |
| 403 | result = benchmark_one_day(test_date, args.llm_model, papers, user_profiles) |
| 404 | |
| 405 | print(f"Date: {result['date']}") |
| 406 | print(f"LLM Model: {result['llm_model']}") |
| 407 | print(f"Embedding Model: {result['embedding_model']}") |
| 408 | print("\nBreakdown:") |
| 409 | for task, data in result["breakdown"].items(): |
| 410 | print(f" {task}:") |
| 411 | print(f" Count: {data.get('papers', data.get('users', data.get('reports', 0)))}") |
| 412 | print(f" Tokens: {data['total_tokens']:,} ({data['tokens_per_paper'] if 'tokens_per_paper' in data else data['tokens_per_user'] if 'tokens_per_user' in data else data['tokens_per_report']}/item)") |
| 413 | print("\nTotal:") |
| 414 | print(f" Embedding Tokens: {result['total']['embedding_tokens']:,}") |
| 415 | print(f" LLM Tokens: {result['total']['llm_tokens']:,}") |
| 416 | print(f" Total Tokens: {result['total']['total_tokens']:,}") |
| 417 | print(f" Estimated Cost: ${result['total']['estimated_cost_usd']:.4f}") |
| 418 |
no test coverage detected