PlantAI-dev/check_cost.py

#!/usr/bin/env python3
"""
成本追踪查询脚本
"""

import requests
import json
from datetime import datetime

API_BASE = "http://localhost:7865"


def get_cost_summary():
    """获取成本摘要"""
    try:
        r = requests.get(f"{API_BASE}/cost/summary", timeout=10)
        if r.status_code == 200:
            return r.json()
        else:
            print(f"错误: HTTP {r.status_code}")
            return None
    except Exception as e:
        print(f"请求失败: {e}")
        return None


def get_cost_report():
    """获取成本报告"""
    try:
        r = requests.get(f"{API_BASE}/cost/report", timeout=10)
        if r.status_code == 200:
            return r.json().get("report", "")
        else:
            print(f"错误: HTTP {r.status_code}")
            return None
    except Exception as e:
        print(f"请求失败: {e}")
        return None


def get_agent_cost(agent_name):
    """获取特定Agent成本"""
    try:
        r = requests.get(f"{API_BASE}/cost/agents/{agent_name}", timeout=10)
        if r.status_code == 200:
            return r.json()
        else:
            print(f"错误: HTTP {r.status_code}")
            return None
    except Exception as e:
        print(f"请求失败: {e}")
        return None


def get_suggestions():
    """获取优化建议"""
    try:
        r = requests.get(f"{API_BASE}/cost/suggestions", timeout=10)
        if r.status_code == 200:
            return r.json().get("suggestions", [])
        else:
            print(f"错误: HTTP {r.status_code}")
            return []
    except Exception as e:
        print(f"请求失败: {e}")
        return []


def print_summary():
    """打印成本摘要"""
    print("=" * 60)
    print("BBAV20311 成本追踪摘要")
    print("=" * 60)
    print(f"查询时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print()
    
    summary = get_cost_summary()
    if not summary:
        print("暂无成本数据")
        return
    
    print("【总体统计】")
    print(f"  总调用次数: {summary.get('total_calls', 0)}")
    print(f"  总Token数: {summary.get('total_tokens', 0):,}")
    print(f"  总成本: ${summary.get('total_cost', 0):.6f}")
    print(f"  平均每次成本: ${summary.get('avg_cost_per_call', 0):.6f}")
    print(f"  平均每次Token: {summary.get('avg_tokens_per_call', 0):,}")
    print()
    
    agents = summary.get('agents', {})
    if agents:
        print("【Agent统计】")
        for name, stats in sorted(agents.items(), key=lambda x: x[1]['cost'], reverse=True):
            print(f"  {name}:")
            print(f"    调用: {stats['calls']}")
            print(f"    Token: {stats['tokens']:,}")
            print(f"    成本: ${stats['cost']:.6f}")
            print(f"    平均: ${stats['avg_cost']:.6f}/call")
        print()
    
    suggestions = get_suggestions()
    if suggestions:
        print("【优化建议】")
        for s in suggestions:
            print(f"  • {s}")
        print()
    
    print("=" * 60)


def print_report():
    """打印完整报告"""
    report = get_cost_report()
    if report:
        print(report)
    else:
        print("无法获取报告")


def main():
    import sys
    
    if len(sys.argv) > 1:
        command = sys.argv[1]
        
        if command == "summary":
            print_summary()
        elif command == "report":
            print_report()
        elif command == "agent" and len(sys.argv) > 2:
            agent_name = sys.argv[2]
            data = get_agent_cost(agent_name)
            if data:
                print(json.dumps(data, indent=2, ensure_ascii=False))
            else:
                print(f"Agent '{agent_name}' 未找到")
        elif command == "suggestions":
            suggestions = get_suggestions()
            print("优化建议:")
            for s in suggestions:
                print(f"  • {s}")
        else:
            print("用法:")
            print("  python check_cost.py summary      # 显示摘要")
            print("  python check_cost.py report       # 显示完整报告")
            print("  python check_cost.py agent <name> # 显示Agent成本")
            print("  python check_cost.py suggestions  # 显示优化建议")
    else:
        # 默认显示摘要
        print_summary()


if __name__ == "__main__":
    main()