agentusage 0.1.1

Local-first agent usage tracking
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

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq)]
pub enum Status {
    #[serde(rename = "OK")]
    Ok,
    #[serde(rename = "NEAR_LIMIT")]
    NearLimit,
    #[serde(rename = "LIMITED")]
    Limited,
    #[serde(rename = "AUTH_REQUIRED")]
    AuthRequired,
    #[serde(rename = "UNSUPPORTED")]
    Unsupported,
    #[serde(rename = "ERROR")]
    Error,
    #[default]
    #[serde(rename = "UNKNOWN")]
    Unknown,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct Metric {
    pub limit: Option<f64>,
    pub remaining: Option<f64>,
    pub used: Option<f64>,
    pub unit: String,
    pub window: String,
}

impl Metric {
    pub fn percent(&self) -> f64 {
        match (self.limit, self.remaining, self.used) {
            (Some(limit), Some(remaining), _) if limit > 0.0 => remaining / limit * 100.0,
            (Some(limit), _, Some(used)) if limit > 0.0 => (limit - used) / limit * 100.0,
            _ => -1.0,
        }
    }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct TokenUsage {
    pub input_tokens: Option<i64>,
    pub output_tokens: Option<i64>,
    pub reasoning_tokens: Option<i64>,
    pub cache_read_tokens: Option<i64>,
    pub cache_write_tokens: Option<i64>,
    pub total_tokens: Option<i64>,
    pub cost_usd: Option<f64>,
    pub requests: Option<i64>,
}

impl TokenUsage {
    pub fn sum_total_tokens(&mut self) {
        if self.total_tokens.is_some() {
            return;
        }
        let values = [
            self.input_tokens,
            self.output_tokens,
            self.reasoning_tokens,
            self.cache_read_tokens,
            self.cache_write_tokens,
        ];
        if values.iter().any(Option::is_some) {
            self.total_tokens = Some(values.into_iter().flatten().sum());
        }
    }
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct UsageSnapshot {
    pub provider_id: String,
    pub account_id: String,
    pub timestamp: DateTime<Utc>,
    pub status: Status,
    pub metrics: std::collections::BTreeMap<String, Metric>,
}

#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct UsageEvent {
    pub time: DateTime<Utc>,
    pub provider_id: String,
    pub model: String,
    pub project: String,
    pub session: String,
    pub input_tokens: i64,
    pub output_tokens: i64,
    pub cache_read_tokens: i64,
    pub cache_creation_tokens: i64,
    pub reasoning_tokens: i64,
    pub cost_usd: f64,
    pub has_cost: bool,
}

pub trait UsageProvider {
    fn id(&self) -> &str;
}

pub trait ItemizedUsageProvider {
    fn itemized_usage(&self) -> anyhow::Result<Vec<UsageEvent>>;
}