rust_llm 2.1.1

//! Port of `lib/ruby_llm/accounting/usage.rb`: the per-attempt usage `Tracker`, the `Entry`
//! readers it relies on (the `UsageEntry` struct itself lives in `message.rs`), and the
//! `Accounting::Usage::Result` storage shared by results that carry attempt accounting.
//!
//! Ruby shares entry objects between the tracker, the caller, and the result. Here the tracker
//! owns its entries: `start` returns the entry's `id`, [`Tracker::entry`] reads it back, and
//! `succeed` attaches a snapshot of the ledger to the result.

use std::sync::Arc;

use serde_json::{Map, Value};

use crate::chat::UsageRecorder;
use crate::config::Config;
use crate::cost::{Cost, Tier};
use crate::embedding::Embedding;
use crate::error::Error;
use crate::image::Image;
use crate::message::{Message, Operation, UsageEntry, UsageStatus};
use crate::model::Model;
use crate::moderation::Moderation;
use crate::rerank::Rerank;
use crate::speech::Speech;
use crate::tokens::Tokens;
use crate::transcription::Transcription;

/// Who an attempt's usage is attributed to (`Accounting::Usage::Entry#owner`). Ruby takes any
/// object and the Rails ledger keeps it only when it is an Active Record; here a [`UsageOwner::Record`]
/// names a row (`owner_type`/`owner_id`), and [`UsageOwner::Other`] is any other value, reported in
/// `usage.rust_llm` but never written to a ledger's owner columns.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum UsageOwner {
    /// A database row, like a `User` or an `Account`: `(owner_type, owner_id)`.
    Record { owner_type: String, owner_id: i64 },
    /// Anything else, such as `"user-42"`.
    Other(String),
}

impl UsageOwner {
    /// A row of `owner_type` (the class name Rails would store, like `"User"`) with id `owner_id`.
    pub fn record(owner_type: impl Into<String>, owner_id: i64) -> UsageOwner {
        UsageOwner::Record {
            owner_type: owner_type.into(),
            owner_id,
        }
    }

    /// The owner as it appears in a `usage.rust_llm` payload: a record as `{type, id}`, anything
    /// else as its string.
    pub fn to_value(&self) -> Value {
        match self {
            UsageOwner::Record {
                owner_type,
                owner_id,
            } => serde_json::json!({ "type": owner_type, "id": owner_id }),
            UsageOwner::Other(s) => Value::String(s.clone()),
        }
    }
}

impl From<&str> for UsageOwner {
    fn from(s: &str) -> Self {
        UsageOwner::Other(s.to_string())
    }
}

impl From<String> for UsageOwner {
    fn from(s: String) -> Self {
        UsageOwner::Other(s)
    }
}

tokio::task_local! {
    static OWNER: Option<UsageOwner>;
}

/// `Accounting::Usage.owner`: the owner of the running task's `with_usage_owner` block, if any.
pub fn usage_owner() -> Option<UsageOwner> {
    OWNER.try_with(Clone::clone).ok().flatten()
}

/// `RubyLLM.with_usage_owner(owner) { ... }` (`Accounting::Usage.with_owner`): attributes the
/// usage of every operation `future` runs to `owner`. An operation's own `owner:` option wins.
/// The outer owner comes back when `future` finishes, fails, or is dropped.
///
/// Ruby keeps the owner in fiber storage, which threads started inside the block inherit. A
/// task-local is not inherited by `tokio::spawn`: wrap the spawned future in [`with_current_owner`].
pub async fn with_usage_owner<F: std::future::Future>(
    owner: impl Into<Option<UsageOwner>>,
    future: F,
) -> F::Output {
    OWNER.scope(owner.into(), future).await
}

/// `future` attributed to the current owner, for a task spawned inside a `with_usage_owner` block
/// (Ruby's threads inherit fiber storage).
pub fn with_current_owner<F: std::future::Future>(
    future: F,
) -> impl std::future::Future<Output = F::Output> {
    OWNER.scope(usage_owner(), future)
}

/// `Accounting::Usage.owned_by(owner) { ... }`: an operation's `owner:` keyword, which wins over
/// the ambient owner; `None` leaves the ambient owner in place.
pub(crate) async fn owned_by<F: std::future::Future>(
    owner: Option<UsageOwner>,
    future: F,
) -> F::Output {
    match owner {
        Some(owner) => OWNER.scope(Some(owner), future).await,
        None => future.await,
    }
}

/// `Accounting::Usage.ledger`: persists finished entries that no chat records, such as the
/// attempts of one-shot operations and finished video jobs (`rust_llm_loco::UsageLedger` writes
/// them to `rust_llm_usages`). Set it on [`Config::usage_ledger`]. A ledger logs its own write
/// failures and never fails the operation that billed the attempt.
#[async_trait::async_trait]
pub trait UsageLedger: Send + Sync {
    /// `record(entry)`.
    async fn record(&self, entry: &UsageEntry);
}

/// `Accounting::Usage.record(entry)`: hands `entry` to the configured ledger, if any.
pub(crate) async fn record(config: &Config, entry: &UsageEntry) {
    if let Some(ledger) = &config.usage_ledger {
        ledger.record(entry).await;
    }
}

/// `Accounting::Usage.report(entry, config:)` for each entry: recorded in the ledger, then
/// published as `usage.rust_llm`.
pub(crate) async fn report(config: &Arc<Config>, entries: &[UsageEntry]) {
    for entry in entries {
        record(config, entry).await;
        crate::instrumentation::usage(config, entry);
    }
}

/// `Accounting::Usage::Entry` readers (`accounting/usage.rb`).
impl UsageEntry {
    /// `Entry.new(operation:, provider:, model:)`: a pending attempt with unknown tokens and an
    /// unpriced cost. The port stores an absent model (`model: nil`) as `""`; `to_h` reports it
    /// as `null`.
    pub fn new(
        operation: Operation,
        provider: impl Into<String>,
        model: Option<&str>,
    ) -> UsageEntry {
        let tokens = Tokens::default();
        UsageEntry {
            id: UsageEntry::next_id(),
            operation,
            provider: provider.into(),
            model: model.unwrap_or_default().to_string(),
            status: UsageStatus::Pending,
            cost: Cost::new(&tokens, None, Tier::Standard),
            tokens,
            owner: usage_owner(),
        }
    }

    pub fn is_pending(&self) -> bool {
        self.status == UsageStatus::Pending
    }

    pub fn is_succeeded(&self) -> bool {
        self.status == UsageStatus::Succeeded
    }

    pub fn is_failed(&self) -> bool {
        self.status == UsageStatus::Failed
    }

    pub fn is_cancelled(&self) -> bool {
        self.status == UsageStatus::Cancelled
    }

    /// `usage_available?`: `tokens.to_h.any?`.
    pub fn usage_available(&self) -> bool {
        !self.tokens.is_empty()
    }

    /// `Entry#to_h`.
    pub fn to_h(&self) -> Value {
        let mut h = Map::new();
        h.insert("operation".into(), self.operation.as_str().into());
        h.insert("provider".into(), self.provider.clone().into());
        h.insert(
            "model".into(),
            if self.model.is_empty() {
                Value::Null
            } else {
                self.model.clone().into()
            },
        );
        h.insert("status".into(), self.status.as_str().into());
        h.insert(
            "tokens".into(),
            crate::instrumentation::tokens_h(&self.tokens),
        );
        h.insert("cost".into(), crate::instrumentation::cost_h(&self.cost));
        Value::Object(h)
    }
}

/// `Accounting::Usage::Result`: a result object that carries attempt accounting, plus what
/// `Tracker#succeed` reads from it (`billed.tokens`, `billed.cost`, and `model_info=`).
pub trait UsageResult {
    /// `ruby_llm_usage_entries=`.
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>);
    /// `billed.tokens`.
    fn usage_tokens(&self) -> Tokens;
    /// `billed.cost`.
    fn usage_cost(&self) -> Cost;
    /// `item.model_info = result_model(item) if item.respond_to?(:model_info=)`: `resolve` maps
    /// the result's model id (`nil` when absent) to the model it is priced with. Results without
    /// `model_info=` (such as `Moderation`) keep the default no-op.
    fn assign_model_info(&mut self, _resolve: &dyn Fn(Option<&str>) -> Option<Model>) {}
}

/// `result.model` as `Tracker#result_model` reads it: an empty id is `nil`.
fn model_id(model: &str) -> Option<&str> {
    (!model.is_empty()).then_some(model)
}

impl UsageResult for Message {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost(None)
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(self.model.as_deref());
    }
}

impl UsageResult for Transcription {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(model_id(&self.model));
    }
}

impl UsageResult for Embedding {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(model_id(&self.model));
    }
}

impl UsageResult for Speech {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(model_id(&self.model));
    }
}

impl UsageResult for Rerank {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(model_id(&self.model));
    }
}

impl UsageResult for Image {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        self.model_info = resolve(model_id(&self.model));
    }
}

/// A request that produced several images is billed once: the first image carries the call
/// (`billed = results.first`), and every image gets its `model_info`.
impl UsageResult for Vec<Image> {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        if let Some(first) = self.first_mut() {
            first.set_usage_entries(entries);
        }
    }
    fn usage_tokens(&self) -> Tokens {
        self.first().map(Image::tokens).unwrap_or_default()
    }
    fn usage_cost(&self) -> Cost {
        self.first()
            .map(Image::cost)
            .unwrap_or_else(|| Cost::new(&Tokens::default(), None, Tier::Standard))
    }
    fn assign_model_info(&mut self, resolve: &dyn Fn(Option<&str>) -> Option<Model>) {
        for image in self {
            image.assign_model_info(resolve);
        }
    }
}

impl UsageResult for Moderation {
    fn set_usage_entries(&mut self, entries: Vec<UsageEntry>) {
        self.usage_entries = entries;
    }
    fn usage_tokens(&self) -> Tokens {
        self.tokens()
    }
    fn usage_cost(&self) -> Cost {
        self.cost()
    }
}

/// `Accounting::Usage::Tracker`: tracks the physical transport attempts belonging to one
/// operation.
pub struct Tracker {
    operation: Operation,
    provider: String,
    model_info: Option<Model>,
    config: Arc<Config>,
    on_finish: Option<UsageRecorder>,
    entries: Vec<UsageEntry>,
    pending: Vec<u64>,
}

impl Tracker {
    /// `Tracker.new(operation:, provider:, model:, config:, on_finish:)`. `provider` is the
    /// provider's slug (Ruby reads `provider.slug`).
    pub fn new(
        operation: Operation,
        provider: impl Into<String>,
        model: Option<Model>,
        config: Arc<Config>,
        on_finish: Option<UsageRecorder>,
    ) -> Tracker {
        Tracker {
            operation,
            provider: provider.into(),
            model_info: model,
            config,
            on_finish,
            entries: Vec::new(),
            pending: Vec::new(),
        }
    }

    /// `attr_reader :entries`.
    pub fn entries(&self) -> &[UsageEntry] {
        &self.entries
    }

    /// The entry `start` returned the id of.
    pub fn entry(&self, id: u64) -> Option<&UsageEntry> {
        self.entries.iter().find(|e| e.id == id)
    }

    fn entry_mut(&mut self, id: u64) -> Option<&mut UsageEntry> {
        self.entries.iter_mut().find(|e| e.id == id)
    }

    /// `start`: records a new pending attempt and returns its id.
    pub fn start(&mut self) -> u64 {
        let entry = UsageEntry::new(
            self.operation,
            self.provider.clone(),
            self.model_info.as_ref().map(|m| m.id.as_str()),
        );
        let id = entry.id;
        self.entries.push(entry);
        self.pending.push(id);
        id
    }

    /// `observe(chunk)`: a streamed chunk's token counts update the attempt in flight; a count
    /// the chunk does not report keeps the earlier value (`merge_stream_tokens`).
    pub fn observe(&mut self, chunk: &Message) {
        self.observe_tokens(&chunk.tokens());
    }

    /// `observe_tokens(tokens)`: keeps the usage a response reported for the current attempt, so
    /// an attempt that fails after the provider answered, such as a blocked one, still records
    /// what it billed. The attempt stays pending.
    pub fn observe_tokens(&mut self, tokens: &Tokens) {
        let Some(&id) = self.pending.last() else {
            return;
        };
        let Some(entry) = self.entry_mut(id) else {
            return;
        };
        entry.tokens.merge_latest(tokens);
        entry.cost = Cost::new(&entry.tokens, None, Tier::Standard);
    }

    /// `fail_attempt(entry, error)`: ignored unless the attempt is still pending.
    pub fn fail_attempt(&mut self, entry: Option<u64>, error: &Error) {
        let Some(id) = entry else { return };
        let Some(entry) = self.entry(id).filter(|e| e.is_pending()) else {
            return;
        };
        let status = if matches!(error, Error::Cancelled) {
            UsageStatus::Cancelled
        } else {
            UsageStatus::Failed
        };
        let tokens = self.failure_tokens(entry, error);
        self.finish(id, status, tokens, None);
    }

    /// `fail_pending(error)`: fails every attempt still in flight.
    pub fn fail_pending(&mut self, error: &Error) {
        for id in self.pending.clone() {
            self.fail_attempt(Some(id), error);
        }
    }

    /// `succeed(result)`: the last pending attempt is credited with the result's tokens and cost;
    /// earlier pending attempts succeed with no tokens. The ledger is then attached to the result,
    /// even when no attempt was recorded.
    pub fn succeed<R: UsageResult>(&mut self, result: &mut R) {
        result.assign_model_info(&|id| self.result_model(id));
        let pending = self.pending.clone();
        let Some((&last, earlier)) = pending.split_last() else {
            result.set_usage_entries(self.entries.clone());
            return;
        };
        let tokens = result.usage_tokens();
        let cost = result.usage_cost();
        for &id in earlier {
            self.finish(id, UsageStatus::Succeeded, Tokens::default(), None);
        }
        self.finish(last, UsageStatus::Succeeded, tokens, Some(cost));
        result.set_usage_entries(self.entries.clone());
    }

    /// `result_model`: the requested model unless the response echoes a different id this
    /// provider's registry knows.
    fn result_model(&self, id: Option<&str>) -> Option<Model> {
        let Some(id) = id else {
            return self.model_info.clone();
        };
        if self.model_info.as_ref().is_some_and(|m| m.id == id) {
            return self.model_info.clone();
        }
        crate::models::models()
            .find(id, Some(&self.provider))
            .ok()
            .or_else(|| self.model_info.clone())
    }

    /// `failure_tokens`: without a model the tokens stay as observed; otherwise the rule the chat
    /// loop uses (`chat::failure_tokens`): observed stream tokens are kept, a request that never
    /// reached the provider or that it refused (4xx) is zero, anything else stays unknown.
    fn failure_tokens(&self, entry: &UsageEntry, error: &Error) -> Tokens {
        if self.model_info.is_none() {
            return entry.tokens.clone();
        }
        crate::chat::failure_tokens(error, entry.usage_available().then(|| entry.tokens.clone()))
    }

    /// `finish`: a supplied cost wins when it has a total; otherwise the tokens are priced for
    /// the operation's category (`CATEGORY_BY_OPERATION`).
    fn finish(&mut self, id: u64, status: UsageStatus, tokens: Tokens, cost: Option<Cost>) {
        let cost = cost
            .filter(|c| c.total().is_some())
            .unwrap_or_else(|| self.price(&tokens));
        let Some(entry) = self.entries.iter_mut().find(|e| e.id == id) else {
            return;
        };
        entry.status = status;
        entry.tokens = tokens;
        entry.cost = cost;
        self.pending.retain(|&p| p != id);
        let entry = entry.clone();
        if let Some(on_finish) = &mut self.on_finish {
            on_finish(&entry);
        }
        crate::instrumentation::usage(&self.config, &entry);
    }

    fn price(&self, tokens: &Tokens) -> Cost {
        let model = self.model_info.as_ref();
        match self.operation {
            Operation::Chat
            | Operation::Moderation
            | Operation::Ocr
            | Operation::Judgment
            | Operation::Video
            | Operation::Research => Cost::new(tokens, model, Tier::Standard),
            Operation::Embedding | Operation::Rerank => {
                crate::rerank::embeddings_cost(tokens, model)
            }
            Operation::Image => Cost::images(tokens, model, None),
            Operation::Speech | Operation::Transcription => Cost::audio(tokens, model),
        }
    }
}