use std::collections::BTreeMap;
use std::fmt;
use std::sync::Arc;
use serde::{Deserialize, Serialize};
use crate::ports::{Answer, AnswerValue, ClassifierProfile, Question};
const PICO_PER_USD: u64 = 1_000_000_000_000;
const CHARS_PER_TOKEN: usize = 4;
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct AdapterId(String);
impl AdapterId {
#[must_use]
pub fn new(id: impl Into<String>) -> Self {
Self(id.into())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for AdapterId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct QuestionKind(String);
impl QuestionKind {
#[must_use]
pub fn new(kind: impl Into<String>) -> Self {
Self(kind.into())
}
#[must_use]
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for QuestionKind {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str(&self.0)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
#[non_exhaustive]
pub enum TokenSource {
Estimated,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Tokens {
pub input: u64,
pub output: u64,
pub source: TokenSource,
}
impl Tokens {
#[must_use]
pub fn estimate(
profile: &ClassifierProfile,
state: &str,
questions: &BTreeMap<String, Question>,
answers: Option<&BTreeMap<String, Answer>>,
) -> Self {
let (sent, _) = profile.truncate(state);
let mut input = sent.chars().count();
for (id, question) in questions {
input += id.chars().count() + question.instructions().chars().count();
let options: Vec<(&String, &String)> = match question {
Question::Choice { criteria, .. } => criteria.iter().collect(),
Question::Score { levels, .. } => levels.iter().map(|(n, m)| (n, m)).collect(),
Question::Noul { .. } => Vec::new(),
};
for (name, meaning) in options {
input += name.chars().count() + meaning.chars().count();
}
}
let output = answers.map_or(0, |answers| {
answers
.iter()
.map(|(id, answer)| {
let probabilities: usize = answer
.probabilities
.iter()
.map(|(label, p)| label.chars().count() + format!("{p:.2}").len())
.sum();
id.chars().count() + value_text(&answer.value).chars().count() + probabilities
})
.sum()
});
Self {
input: tokens_for(input),
output: tokens_for(output),
source: TokenSource::Estimated,
}
}
}
fn tokens_for(chars: usize) -> u64 {
u64::try_from(chars.div_ceil(CHARS_PER_TOKEN)).unwrap_or(u64::MAX)
}
pub(crate) fn value_text(value: &AnswerValue) -> String {
match value {
AnswerValue::Choice(label) => label.clone(),
AnswerValue::Score(score) => score.to_string(),
AnswerValue::Noul(verdict) => verdict.to_string(),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct Price {
pub input_pico_usd_per_token: u64,
pub output_pico_usd_per_token: u64,
}
impl Price {
pub const FREE: Price = Price::per_million_tokens(0, 0);
#[must_use]
pub const fn per_million_tokens(input_micro_usd: u64, output_micro_usd: u64) -> Price {
Price {
input_pico_usd_per_token: input_micro_usd,
output_pico_usd_per_token: output_micro_usd,
}
}
#[must_use]
pub const fn cost(&self, tokens: &Tokens) -> Cost {
Cost(
tokens
.input
.saturating_mul(self.input_pico_usd_per_token)
.saturating_add(tokens.output.saturating_mul(self.output_pico_usd_per_token)),
)
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Default)]
pub struct Cost(u64);
impl Cost {
pub const ZERO: Cost = Cost(0);
#[must_use]
pub const fn from_pico_usd(pico_usd: u64) -> Cost {
Cost(pico_usd)
}
#[must_use]
pub const fn pico_usd(&self) -> u64 {
self.0
}
}
impl fmt::Display for Cost {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"${}.{:012}",
self.0 / PICO_PER_USD,
self.0 % PICO_PER_USD
)
}
}
#[derive(Debug, Clone, Default, PartialEq)]
pub struct PriceSheet(BTreeMap<AdapterId, Price>);
impl PriceSheet {
#[must_use]
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with(mut self, adapter: AdapterId, price: Price) -> Self {
self.0.insert(adapter, price);
self
}
#[must_use]
pub fn cost_of(&self, adapter: &AdapterId, tokens: &Tokens) -> Option<Cost> {
self.0.get(adapter).map(|price| price.cost(tokens))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CallRole {
Served,
Shadow,
Discarded,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum CallOutcome {
Ok,
Failed,
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub struct CallRecord {
pub adapter: AdapterId,
pub kinds: Vec<QuestionKind>,
pub tokens: Tokens,
pub cost: Option<Cost>,
pub role: CallRole,
pub outcome: CallOutcome,
}
pub trait CostLedger: Send + Sync {
fn record(&self, record: CallRecord);
}
#[derive(Clone)]
pub struct Accounting {
prices: PriceSheet,
ledger: Arc<dyn CostLedger>,
}
impl Accounting {
#[must_use]
pub fn new(prices: PriceSheet, ledger: Arc<dyn CostLedger>) -> Self {
Self { prices, ledger }
}
pub(crate) fn record(
&self,
adapter: &AdapterId,
profile: &ClassifierProfile,
state: &str,
questions: &BTreeMap<String, Question>,
answers: Option<&BTreeMap<String, Answer>>,
role: CallRole,
) {
let tokens = Tokens::estimate(profile, state, questions, answers);
self.ledger.record(CallRecord {
adapter: adapter.clone(),
kinds: questions.keys().map(QuestionKind::new).collect(),
tokens,
cost: answers.and_then(|_| self.prices.cost_of(adapter, &tokens)),
role,
outcome: if answers.is_some() {
CallOutcome::Ok
} else {
CallOutcome::Failed
},
});
}
}
impl fmt::Debug for Accounting {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Accounting")
.field("prices", &self.prices)
.finish_non_exhaustive()
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct LedgerTotals {
pub calls: u64,
pub input_tokens: u64,
pub output_tokens: u64,
pub pico_usd: u128,
pub unpriced_calls: u64,
pub failed_calls: u64,
}
impl LedgerTotals {
fn add(&mut self, record: &CallRecord) {
self.calls = self.calls.saturating_add(1);
self.input_tokens = self.input_tokens.saturating_add(record.tokens.input);
self.output_tokens = self.output_tokens.saturating_add(record.tokens.output);
match (record.cost, record.outcome) {
(Some(cost), _) => {
self.pico_usd = self.pico_usd.saturating_add(u128::from(cost.pico_usd()));
}
(None, CallOutcome::Ok) => self.unpriced_calls = self.unpriced_calls.saturating_add(1),
(None, CallOutcome::Failed) => {}
}
if record.outcome == CallOutcome::Failed {
self.failed_calls = self.failed_calls.saturating_add(1);
}
}
}
#[derive(Debug, Default)]
pub struct InMemoryLedger {
#[expect(
clippy::disallowed_types,
reason = "the ledger is an explicitly wired recording cell, not ambient request state (ADR 0007)"
)]
records: std::sync::Mutex<Vec<CallRecord>>,
}
impl InMemoryLedger {
#[must_use]
pub fn new() -> Self {
Self::default()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Vec<CallRecord>> {
self.records
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
}
#[must_use]
pub fn records(&self) -> Vec<CallRecord> {
self.lock().clone()
}
#[must_use]
pub fn totals(&self) -> BTreeMap<AdapterId, LedgerTotals> {
let mut totals: BTreeMap<AdapterId, LedgerTotals> = BTreeMap::new();
for record in self.lock().iter() {
totals
.entry(record.adapter.clone())
.or_default()
.add(record);
}
totals
}
#[must_use]
pub fn role_totals(&self) -> BTreeMap<CallRole, LedgerTotals> {
let mut totals: BTreeMap<CallRole, LedgerTotals> = BTreeMap::new();
for record in self.lock().iter() {
totals.entry(record.role).or_default().add(record);
}
totals
}
}
impl CostLedger for InMemoryLedger {
fn record(&self, record: CallRecord) {
self.lock().push(record);
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ports::Calibration;
fn tokens(input: u64, output: u64) -> Tokens {
Tokens {
input,
output,
source: TokenSource::Estimated,
}
}
#[test]
fn a_published_price_maps_one_to_one_and_displays_every_pico_usd() {
let price = Price::per_million_tokens(3_000_000, 15_000_000);
assert_eq!(
price.cost(&tokens(1_000_000, 0)).to_string(),
"$3.000000000000"
);
assert_eq!(
price.cost(&tokens(1_000_000, 1_000_000)).pico_usd(),
18 * PICO_PER_USD
);
assert_eq!(Cost::from_pico_usd(1).to_string(), "$0.000000000001");
assert_eq!(
Price::per_million_tokens(u64::MAX, u64::MAX)
.cost(&tokens(u64::MAX, u64::MAX))
.pico_usd(),
u64::MAX,
"saturates, never wraps into a credit"
);
}
#[test]
fn an_unpriced_adapter_is_unknown_and_a_free_one_is_zero() {
let free = AdapterId::new("self-hosted");
let sheet = PriceSheet::new().with(free.clone(), Price::FREE);
assert_eq!(sheet.cost_of(&free, &tokens(10, 10)), Some(Cost::ZERO));
assert_eq!(
sheet.cost_of(&AdapterId::new("mystery"), &tokens(10, 10)),
None
);
}
#[test]
fn tokens_are_estimated_from_the_truncated_state_the_questions_and_the_answers() {
let questions = BTreeMap::from([(
"ok".to_owned(),
Question::Choice {
instructions: "pick".to_owned(),
criteria: BTreeMap::from([
("a".to_owned(), "yes".to_owned()),
("b".to_owned(), "no".to_owned()),
]),
},
)]);
let profile = ClassifierProfile::new(Calibration::Classifier, 8);
let asked = Tokens::estimate(&profile, &"x".repeat(20), &questions, None);
assert_eq!(
asked,
tokens(6, 0),
"ceil(22 / 4), and no output for a failed call"
);
let answers = BTreeMap::from([(
"ok".to_owned(),
Answer::choice(
"a",
BTreeMap::from([("a".to_owned(), 0.9), ("b".to_owned(), 0.1)]),
),
)]);
let answered = Tokens::estimate(&profile, "", &questions, Some(&answers));
assert_eq!(answered.output, 4);
assert_eq!(answered.source, TokenSource::Estimated);
}
fn record(
adapter: &str,
cost: Option<Cost>,
role: CallRole,
outcome: CallOutcome,
) -> CallRecord {
CallRecord {
adapter: AdapterId::new(adapter),
kinds: vec![QuestionKind::new("topic")],
tokens: tokens(100, 10),
cost,
role,
outcome,
}
}
#[test]
fn totals_separate_unpriced_from_failed_and_do_not_saturate_at_one_calls_ceiling() {
let ledger = InMemoryLedger::new();
let huge = Some(Cost::from_pico_usd(u64::MAX));
ledger.record(record("cheap", huge, CallRole::Served, CallOutcome::Ok));
ledger.record(record("cheap", huge, CallRole::Discarded, CallOutcome::Ok));
ledger.record(record("cheap", None, CallRole::Served, CallOutcome::Ok));
ledger.record(record("cheap", None, CallRole::Shadow, CallOutcome::Failed));
let cheap = ledger.totals()[&AdapterId::new("cheap")];
assert_eq!(cheap.calls, 4);
assert_eq!(cheap.pico_usd, u128::from(u64::MAX) * 2);
assert_eq!(
cheap.unpriced_calls, 1,
"only the answered, unpriced call is a gap"
);
assert_eq!(cheap.failed_calls, 1);
assert_eq!(cheap.input_tokens, 400);
let roles = ledger.role_totals();
assert_eq!(roles[&CallRole::Discarded].calls, 1);
assert_eq!(roles[&CallRole::Served].calls, 2);
assert_eq!(ledger.records()[3].role, CallRole::Shadow, "oldest first");
}
}