TokenLedgerRs 0.1.0

Token management and pricing governance CLI for AI coding agents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! CLI command handlers for benchmarks.

use crate::benchmarks::{artificial_analysis, openrouter, store::BenchmarkStore};
use crate::cli::BenchmarksArgs;
use anyhow::Result;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::info;

/// Main entry point for benchmarks CLI commands
pub fn run_benchmarks(args: BenchmarksArgs) -> Result<()> {
    match args.command {
        crate::cli::BenchmarksCommand::Refresh(args) => run_refresh(args),
        crate::cli::BenchmarksCommand::List(args) => run_list(args),
        crate::cli::BenchmarksCommand::Show(args) => run_show(args),
        crate::cli::BenchmarksCommand::Validate(args) => run_validate(args),
    }
}

/// Refresh benchmarks from configured sources
fn run_refresh(args: crate::cli::RefreshBenchmarksArgs) -> Result<()> {
    println!("Refreshing benchmarks...");

    let rt = tokio::runtime::Runtime::new()?;

    rt.block_on(async {
        let store = Arc::new(RwLock::new(BenchmarkStore::new(3600)));
        let mut total = 0;

        // Refresh from Artificial Analysis if API key provided
        if let Some(ref key) = args.aa_api_key {
            info!("Fetching benchmarks from Artificial Analysis...");
            match artificial_analysis::fetch_benchmarks(key).await {
                Ok(data) => {
                    for benchmark in &data {
                        store
                            .write()
                            .await
                            .merge(benchmark.model_id.clone(), benchmark.clone())
                            .await;
                    }
                    total += data.len();
                    println!("✓ Fetched {} models from Artificial Analysis", data.len());
                }
                Err(e) => {
                    println!("✗ Failed to fetch from Artificial Analysis: {}", e);
                }
            }
        }

        // Refresh from OpenRouter if API key provided
        if let Some(ref key) = args.openrouter_api_key {
            info!("Fetching benchmarks from OpenRouter...");
            match openrouter::fetch_benchmarks(key).await {
                Ok(data) => {
                    for benchmark in &data {
                        store
                            .write()
                            .await
                            .merge(benchmark.model_id.clone(), benchmark.clone())
                            .await;
                    }
                    total += data.len();
                    println!("✓ Fetched {} models from OpenRouter", data.len());
                }
                Err(e) => {
                    println!("✗ Failed to fetch from OpenRouter: {}", e);
                }
            }
        }

        // Output results
        if total > 0 {
            println!("\nTotal benchmarks: {}", total);

            if let Some(output_path) = &args.output {
                // Just print count to output file
                std::fs::write(output_path, format!("{{\"count\":{}}}", total))?;
                println!("✓ Written summary to: {}", output_path.display());
            }
        } else {
            println!("No benchmarks fetched. Provide API keys or use --no-fetch for cache.");
            println!("\nUsage:");
            println!("  tokenledger benchmarks refresh --aa-api-key YOUR_KEY");
            println!("  tokenledger benchmarks refresh --openrouter-api-key YOUR_KEY");
        }

        Ok(())
    })
}

/// List available benchmarks
fn run_list(args: crate::cli::ListBenchmarksArgs) -> Result<()> {
    println!("Listing benchmarks (limit: {})...", args.limit);
    // Placeholder - would load from store
    println!("Use 'tokenledger benchmarks refresh' first to fetch data.");
    Ok(())
}

/// Show specific model benchmark
fn run_show(args: crate::cli::ShowBenchmarkArgs) -> Result<()> {
    println!("Showing benchmark for: {}", args.model_id);
    // Placeholder - would lookup from store
    println!("Use 'tokenledger benchmarks refresh' first to fetch data.");
    Ok(())
}

/// Validate benchmark configuration
fn run_validate(args: crate::cli::ValidateBenchmarksArgs) -> Result<()> {
    println!("Validating benchmark configuration...");

    if let Some(config_path) = args.config {
        let content = std::fs::read_to_string(&config_path)?;
        let _config: serde_yaml::Value = serde_yaml::from_str(&content)?;
        println!("✓ Config is valid: {}", config_path.display());
    } else {
        println!("No config file provided. Checking default locations...");

        let default_paths = ["configs/benchmarks.yaml", "configs/benchmarks.example.yaml"];

        for path in default_paths {
            if std::path::Path::new(path).exists() {
                println!("✓ Found: {}", path);
            }
        }
    }

    Ok(())
}