Function main

experiments/token_cost/token_usage_tracker.py:361–442

()

def main():
    parser = argparse.ArgumentParser(description="Token 花销统计工具")
    parser.add_argument("--start-date", type=str, default=None, help="开始日期 YYYYMMDD")
    parser.add_argument("--end-date", type=str, default=None, help="结束日期 YYYYMMDD")
    parser.add_argument("--group-by", type=str, default="date", choices=["date", "task_type", "model", "user_id"])
    parser.add_argument("--mode", type=str, default="stats", choices=["stats", "benchmark", "compare"])
    parser.add_argument("--llm-model", type=str, default="qwen3.5-plus", help="LLM 模型名称")
    parser.add_argument("--days", type=int, default=1, help="基准测试天数")
    args = parser.parse_args()

    # 连接数据库获取论文和用户数据
    conn = sqlite3.connect(DB_PATH)
    papers = conn.execute("SELECT title, abstract FROM papers").fetchall()
    papers = [{"title": p[0], "abstract": p[1] or ""} for p in papers]

    profiles = conn.execute("SELECT user_id, profile_json FROM profiles").fetchall()
    user_profiles = [{"user_id": p[0], **json.loads(p[1])} for p in profiles]
    conn.close()

    print(f"Loaded {len(papers)} papers, {len(user_profiles)} users")

    if args.mode == "stats" and args.start_date and args.end_date:
        # 统计已有日志
        stats = get_usage_stats(args.start_date, args.end_date, args.group_by)
        print(f"\n=== Token Usage Stats ({args.start_date} to {args.end_date}) ===\n")

        for key, s in sorted(stats.items()):
            print(f"[{key}]")
            print(f"  Embedding: {s['embedding_input']:,} tokens")
            print(f"  LLM: {s['llm_input'] + s['llm_output']:,} tokens")
            print(f"  Cost: ${s['cost']:.4f}")
            print(f"  Calls: {s['calls']}")
            print()

    elif args.mode == "benchmark":
        # 基准测试指定天数
        today = datetime.now()
        for i in range(args.days):
            test_date = (today - timedelta(days=i)).strftime("%Y-%m-%d")
            print(f"\n{'='*60}")
            print(f"Benchmarking {test_date} with {args.llm_model}...")
            print(f"{'='*60}\n")

            result = benchmark_one_day(test_date, args.llm_model, papers, user_profiles)

            print(f"Date: {result['date']}")
            print(f"LLM Model: {result['llm_model']}")
            print(f"Embedding Model: {result['embedding_model']}")
            print(f"\nBreakdown:")
            for task, data in result["breakdown"].items():
                print(f"  {task}:")
                print(f"    Count: {data.get('papers', data.get('users', data.get('reports', 0)))}")
                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)")
            print(f"\nTotal:")
            print(f"  Embedding Tokens: {result['total']['embedding_tokens']:,}")
            print(f"  LLM Tokens: {result['total']['llm_tokens']:,}")
            print(f"  Total Tokens: {result['total']['total_tokens']:,}")
            print(f"  Estimated Cost: ${result['total']['estimated_cost_usd']:.4f}")