yapi-ai 0.4.0

LLM wire APIs, model catalog and registry, cost tracking, credentials and OAuth. An internal library of yapi, with no stable API.
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

//! Request cost from token usage.

use yapi_types::message::Usage;
use yapi_types::model::{Model, Pricing};

/// Prices `usage` with the model's rates and stores the result in `usage.cost`.
///
/// The highest tier whose `inputTokensAbove` is below the request's input tokens
/// (fresh, cached and written) prices the whole request. One-hour cache writes cost
/// twice the base input rate. Arithmetic follows pi operation by operation, so costs
/// print identically.
pub fn calculate_cost(model: &Model, usage: &mut Usage) {
    calculate_cost_with(&model.cost, usage);
}

/// [`calculate_cost`] with a model's `pricing`, for any model type.
pub fn calculate_cost_with(pricing: &Pricing, usage: &mut Usage) {
    let input_tokens = usage.input + usage.cache_read + usage.cache_write;
    let (mut input, mut output, mut cache_read, mut cache_write) = (
        pricing.input,
        pricing.output,
        pricing.cache_read,
        pricing.cache_write,
    );
    let mut matched: Option<u64> = None;
    for tier in pricing.tiers.iter().flatten() {
        if input_tokens > tier.input_tokens_above
            && matched.is_none_or(|threshold| tier.input_tokens_above > threshold)
        {
            (input, output, cache_read, cache_write) =
                (tier.input, tier.output, tier.cache_read, tier.cache_write);
            matched = Some(tier.input_tokens_above);
        }
    }

    let long_write = usage.cache_write_1h.unwrap_or(0);
    let short_write = usage.cache_write.saturating_sub(long_write);
    let cost = &mut usage.cost;
    cost.input = (input / 1_000_000.0) * usage.input as f64;
    cost.output = (output / 1_000_000.0) * usage.output as f64;
    cost.cache_read = (cache_read / 1_000_000.0) * usage.cache_read as f64;
    cost.cache_write =
        (cache_write * short_write as f64 + input * 2.0 * long_write as f64) / 1_000_000.0;
    cost.total = cost.input + cost.output + cost.cache_read + cost.cache_write;
}

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

    fn model(cost: &str) -> Model {
        serde_json::from_str(&format!(
            r#"{{"id":"m","name":"M","api":"a","provider":"p","baseUrl":"u","reasoning":false,
                "input":["text"],"cost":{cost},"contextWindow":1,"maxTokens":1}}"#
        ))
        .unwrap()
    }

    #[test]
    fn matches_pi_numbers() {
        // From pi's JSON output for claude-sonnet-4-5 on the mock cassette.
        let model = model(r#"{"input":3,"output":15,"cacheRead":0.3,"cacheWrite":3.75}"#);
        let mut usage = Usage {
            input: 12,
            output: 6,
            ..Usage::default()
        };
        calculate_cost(&model, &mut usage);
        assert_eq!(
            json::to_string(&usage.cost).unwrap(),
            r#"{"input":0.000036,"output":0.00009,"cacheRead":0,"cacheWrite":0,"total":0.000126}"#
        );
    }

    #[test]
    fn highest_matching_tier_wins() {
        let model = model(
            r#"{"input":1,"output":1,"cacheRead":1,"cacheWrite":1,"tiers":[
                {"inputTokensAbove":100,"input":2,"output":2,"cacheRead":2,"cacheWrite":2},
                {"inputTokensAbove":10,"input":3,"output":3,"cacheRead":3,"cacheWrite":3}]}"#,
        );
        let mut usage = Usage {
            input: 1_000_000,
            ..Usage::default()
        };
        calculate_cost(&model, &mut usage);
        assert_eq!(usage.cost.input, 2.0);

        let mut usage = Usage {
            input: 50,
            ..Usage::default()
        };
        calculate_cost(&model, &mut usage);
        assert_eq!(usage.cost.input, 3.0 / 1_000_000.0 * 50.0);
    }

    #[test]
    fn long_cache_writes_double() {
        let model = model(r#"{"input":3,"output":0,"cacheRead":0,"cacheWrite":3.75}"#);
        let mut usage = Usage {
            cache_write: 300,
            cache_write_1h: Some(100),
            ..Usage::default()
        };
        calculate_cost(&model, &mut usage);
        assert_eq!(
            usage.cost.cache_write,
            (3.75 * 200.0 + 3.0 * 2.0 * 100.0) / 1e6
        );
    }
}