agent-model-contract 0.2.0

Transport-neutral contracts for Agent Infra model APIs
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

//! Observed token counters: absent is unreported, while Some(0) is an explicit zero.

use serde::{Deserialize, Serialize};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct TokenUsage {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_tokens: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub output_tokens: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_cached_tokens: Option<u64>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub input_cache_write_tokens: Option<u64>,
}

impl TokenUsage {
    pub fn is_reported(&self) -> bool {
        self.input_tokens.is_some()
            || self.output_tokens.is_some()
            || self.input_cached_tokens.is_some()
            || self.input_cache_write_tokens.is_some()
    }

    /// SSE usage frames contain cumulative counters, not increments.
    pub fn merge(&mut self, next: &Self) {
        self.input_tokens = next.input_tokens.or(self.input_tokens);
        self.output_tokens = next.output_tokens.or(self.output_tokens);
        self.input_cached_tokens = next.input_cached_tokens.or(self.input_cached_tokens);
        self.input_cache_write_tokens = next
            .input_cache_write_tokens
            .or(self.input_cache_write_tokens);
    }

    pub fn total_tokens(&self) -> Option<u64> {
        self.input_tokens?.checked_add(self.output_tokens?)
    }
}

/// Number of invocations that actually reported each counter; denominator is invocation_count.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
#[serde(rename_all = "camelCase")]
pub struct TokenUsageCoverage {
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub input_cached_tokens: u64,
    pub input_cache_write_tokens: u64,
}

impl crate::ModelUsageSummary {
    pub fn observe_usage(&mut self, usage: &TokenUsage) {
        if let Some(tokens) = usage.input_tokens {
            self.input_tokens = self.input_tokens.saturating_add(tokens);
            self.usage_coverage.input_tokens += 1;
        }
        if let Some(tokens) = usage.output_tokens {
            self.output_tokens = self.output_tokens.saturating_add(tokens);
            self.usage_coverage.output_tokens += 1;
        }
        self.usage_coverage.input_cached_tokens += u64::from(usage.input_cached_tokens.is_some());
        self.usage_coverage.input_cache_write_tokens +=
            u64::from(usage.input_cache_write_tokens.is_some());
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use serde_json::json;

    #[test]
    fn missing_zero_and_partial_counters_survive_serialization_and_merge() {
        let mut usage: TokenUsage =
            serde_json::from_value(json!({"inputTokens":10,"inputCachedTokens":6})).unwrap();
        assert_eq!(usage.total_tokens(), None);
        usage.merge(&serde_json::from_value(json!({"outputTokens":4})).unwrap());
        usage.merge(&TokenUsage::default());
        assert_eq!(usage.total_tokens(), Some(14));
        assert_eq!(
            serde_json::to_value(&usage).unwrap(),
            json!({"inputTokens":10,"outputTokens":4,"inputCachedTokens":6})
        );
        usage.merge(&serde_json::from_value(json!({"inputCachedTokens":0})).unwrap());
        assert_eq!(usage.input_cached_tokens, Some(0));
        assert!(!TokenUsage::default().is_reported());
        assert_eq!(
            serde_json::to_value(TokenUsage::default()).unwrap(),
            json!({})
        );
    }

    #[test]
    fn summaries_count_only_observed_fields_and_total_overflow_is_absent() {
        let mut summary = crate::ModelUsageSummary::default();
        summary.observe_usage(&TokenUsage::default());
        summary.observe_usage(&TokenUsage {
            input_tokens: Some(0),
            input_cached_tokens: Some(0),
            ..Default::default()
        });
        assert_eq!(summary.usage_coverage.input_tokens, 1);
        assert_eq!(summary.usage_coverage.output_tokens, 0);
        assert_eq!(summary.usage_coverage.input_cached_tokens, 1);
        assert_eq!(
            TokenUsage {
                input_tokens: Some(u64::MAX),
                output_tokens: Some(1),
                ..Default::default()
            }
            .total_tokens(),
            None
        );
    }
}