dwara_core/ai/
cost.rs

1//! AI cost attribution (DW-079): pricing tables and per-call cost
2//! computation.
3//!
4//! The pricing table is a COMPILED view of the `ai.pricing` config map:
5//! provider model identifier -> per-1k-token micro-USD rates. It lives
6//! on the dataplane behind an ArcSwap and is swapped on every config
7//! reload (the same pattern as the AI budget engine), so a pricing
8//! change takes effect on the NEXT request without a restart.
9//!
10//! # Cost computation
11//!
12//! `cost_micros` computes integer micro-USD for one call's usage:
13//!
14//! ```text
15//! (input_tokens  - cached_tokens) * input_per_1k_micros  / 1000
16//! + cached_tokens * cached_input_per_1k_micros           / 1000
17//! + output_tokens * output_per_1k_micros                 / 1000
18//! ```
19//!
20//! When `batch` is true, the batch rates (`batch_input_per_1k_micros`,
21//! `batch_output_per_1k_micros`) replace the standard input/output
22//! rates when declared; cached-token pricing is NOT combined with
23//! batch rates (a batch request uses the batch rate for all input
24//! tokens — the provider's batch API does not report a separate cached
25//! tier). When optional fields are absent, the standard rates apply
26//! (the fallback).
27//!
28//! Integer division truncates (the per-1k rate is the provider's
29//! published unit; sub-micro-USD fractions are lost — the honest
30//! rounding boundary, pinned by tests). Saturating adds guard against
31//! overflow. An UNKNOWN model (no pricing entry) returns `None` — the
32//! caller applies the `unknown_model_policy` (allow / fail_closed /
33//! alert, DW-AI-03). The original fail-open behavior (cost 0) is
34//! preserved by the `allow` policy at the call site.
35//!
36//! # Dependency direction
37//!
38//! `ai` depends on `config` only (see `scripts/check_deps.py`); this
39//! module reads `config::ai::AiConfig` and `ai::types::Usage`, both
40//! same-level or lower. The spend RECORD DTO lives in `analytics` (a
41//! plain struct, not importing `ai::types::Usage`) — the dataplane
42//! converts at the call site, keeping the dependency direction
43//! downward.
44
45use crate::ai::types::Usage;
46use crate::config::ai::AiConfig;
47use std::collections::BTreeMap;
48
49/// One model's per-1k-token rates (micro-USD). The optional cached and
50/// batch fields extend the standard rates (DW-AI-03); when absent the
51/// standard rates apply.
52#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
53struct Price {
54    input_per_1k_micros: u64,
55    output_per_1k_micros: u64,
56    /// Discounted rate for prompt-cached input tokens.
57    cached_input_per_1k_micros: Option<u64>,
58    /// Batch-API input rate.
59    batch_input_per_1k_micros: Option<u64>,
60    /// Batch-API output rate.
61    batch_output_per_1k_micros: Option<u64>,
62}
63
64/// The compiled pricing table (DW-079): provider model -> rates. Built
65/// at dataplane refresh from the published config; immutable once
66/// built. Stored on the dataplane behind an ArcSwap and swapped on
67/// reload, so a pricing change applies to the next request with no
68/// restart.
69#[derive(Debug, Clone, Default)]
70pub struct PricingTable {
71    prices: BTreeMap<String, Price>,
72}
73
74impl PricingTable {
75    /// Compile from the `ai:` config block's pricing map. Absent or
76    /// empty pricing yields an empty table (every model is unknown ->
77    /// the caller applies the unknown-model policy).
78    pub fn compile(cfg: Option<&AiConfig>) -> Self {
79        let prices = cfg
80            .map(|c| {
81                c.pricing
82                    .iter()
83                    .map(|(model, p)| {
84                        (
85                            model.clone(),
86                            Price {
87                                input_per_1k_micros: p.input_per_1k_micros,
88                                output_per_1k_micros: p.output_per_1k_micros,
89                                cached_input_per_1k_micros: p.cached_input_per_1k_micros,
90                                batch_input_per_1k_micros: p.batch_input_per_1k_micros,
91                                batch_output_per_1k_micros: p.batch_output_per_1k_micros,
92                            },
93                        )
94                    })
95                    .collect()
96            })
97            .unwrap_or_default();
98        PricingTable { prices }
99    }
100
101    /// Whether the table carries a pricing entry for `provider_model`.
102    /// Cheap lookup used by the dataplane to apply the unknown-model
103    /// policy (DW-AI-03) without computing the full cost.
104    pub fn has_pricing(&self, provider_model: &str) -> bool {
105        self.prices.contains_key(provider_model)
106    }
107
108    /// Micro-USD for one call's usage against this pricing table.
109    /// Returns `None` when the model has no pricing entry (the caller
110    /// applies the unknown-model policy); `Some(cost)` when it does.
111    /// Integer micro-USD, saturating. When `batch` is true, batch rates
112    /// replace the standard rates when declared. When
113    /// `usage.cached_tokens` is present and a cached rate is declared,
114    /// the cached portion of the prompt is priced at the cached rate
115    /// (not combined with batch — a batch request uses the batch rate
116    /// for all input tokens).
117    pub fn cost_micros(&self, provider_model: &str, usage: Usage, batch: bool) -> Option<u64> {
118        let price = self.prices.get(provider_model)?;
119        let input = usage.prompt_tokens.unwrap_or(0);
120        let output = usage.completion_tokens.unwrap_or(0);
121        // Batch rates replace the standard rates when declared; cached
122        // pricing is skipped under batch (the batch API does not report
123        // a separate cached tier).
124        if batch {
125            let in_rate = price
126                .batch_input_per_1k_micros
127                .unwrap_or(price.input_per_1k_micros);
128            let out_rate = price
129                .batch_output_per_1k_micros
130                .unwrap_or(price.output_per_1k_micros);
131            let input_cost = input.saturating_mul(in_rate) / 1000;
132            let output_cost = output.saturating_mul(out_rate) / 1000;
133            return Some(input_cost.saturating_add(output_cost));
134        }
135        // Standard (non-batch) path: split the prompt into cached and
136        // non-cached portions when both a cached rate and a cached
137        // token count are present.
138        let cached = usage.cached_tokens.unwrap_or(0);
139        let (input_cost, cached_cost) = if let Some(cached_rate) = price.cached_input_per_1k_micros
140        {
141            let non_cached = input.saturating_sub(cached);
142            let std_cost = non_cached.saturating_mul(price.input_per_1k_micros) / 1000;
143            let cached_cost = cached.saturating_mul(cached_rate) / 1000;
144            (std_cost, cached_cost)
145        } else {
146            let std_cost = input.saturating_mul(price.input_per_1k_micros) / 1000;
147            (std_cost, 0)
148        };
149        let output_cost = output.saturating_mul(price.output_per_1k_micros) / 1000;
150        Some(
151            input_cost
152                .saturating_add(cached_cost)
153                .saturating_add(output_cost),
154        )
155    }
156
157    /// Whether the table carries any pricing (cheap dataplane skip).
158    pub fn is_empty(&self) -> bool {
159        self.prices.is_empty()
160    }
161}
dwara_core/ai/cost.rs

dwara_core/ai/
cost.rs
