Skip to main content

git_internal/internal/object/
run_usage.rs

1//! Run usage / cost event.
2//!
3//! `RunUsage` stores immutable usage totals for a `Run`.
4//!
5//! # How to use this object
6//!
7//! - Create it after the run, model call batch, or accounting phase has
8//!   produced stable token totals.
9//! - Keep it append-only; if Libra needs additional rollups, compute
10//!   them in projections.
11//!
12//! # How it works with other objects
13//!
14//! - `run_id` links usage to the owning `Run`.
15//! - `Provenance` supplies the corresponding provider/model
16//!   configuration.
17//!
18//! # How Libra should call it
19//!
20//! Libra should aggregate analytics, quotas, and billing views from
21//! stored `RunUsage` records instead of backfilling usage into
22//! `Provenance` or `Run`.
23
24use std::fmt;
25
26use serde::{Deserialize, Serialize};
27use uuid::Uuid;
28
29use crate::{
30    errors::GitError,
31    hash::ObjectHash,
32    internal::object::{
33        ObjectTrait,
34        types::{ActorRef, Header, ObjectType},
35    },
36};
37
38/// Immutable token / cost summary for one `Run`.
39#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40#[serde(deny_unknown_fields)]
41pub struct RunUsage {
42    /// Common object header carrying the immutable object id, type,
43    /// creator, and timestamps.
44    #[serde(flatten)]
45    header: Header,
46    /// Canonical owning run for this usage summary.
47    run_id: Uuid,
48    /// Input tokens consumed by the run or model-call batch.
49    input_tokens: u64,
50    /// Output tokens produced by the run or model-call batch.
51    output_tokens: u64,
52    /// Precomputed total tokens for quick reads and validation.
53    total_tokens: u64,
54    /// Optional billing estimate in USD.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    cost_usd: Option<f64>,
57}
58
59impl RunUsage {
60    /// Create a new immutable usage summary for one run.
61    pub fn new(
62        created_by: ActorRef,
63        run_id: Uuid,
64        input_tokens: u64,
65        output_tokens: u64,
66        cost_usd: Option<f64>,
67    ) -> Result<Self, String> {
68        Ok(Self {
69            header: Header::new(ObjectType::RunUsage, created_by)?,
70            run_id,
71            input_tokens,
72            output_tokens,
73            total_tokens: input_tokens + output_tokens,
74            cost_usd,
75        })
76    }
77
78    /// Return the immutable header for this usage record.
79    pub fn header(&self) -> &Header {
80        &self.header
81    }
82
83    /// Return the canonical owning run id.
84    pub fn run_id(&self) -> Uuid {
85        self.run_id
86    }
87
88    /// Return the input token count.
89    pub fn input_tokens(&self) -> u64 {
90        self.input_tokens
91    }
92
93    /// Return the output token count.
94    pub fn output_tokens(&self) -> u64 {
95        self.output_tokens
96    }
97
98    /// Return the total token count.
99    pub fn total_tokens(&self) -> u64 {
100        self.total_tokens
101    }
102
103    /// Return the billing estimate in USD, if present.
104    pub fn cost_usd(&self) -> Option<f64> {
105        self.cost_usd
106    }
107
108    /// Validate that the stored total matches input plus output.
109    pub fn is_consistent(&self) -> bool {
110        self.total_tokens == self.input_tokens + self.output_tokens
111    }
112}
113
114impl fmt::Display for RunUsage {
115    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
116        write!(f, "RunUsage: {}", self.header.object_id())
117    }
118}
119
120impl ObjectTrait for RunUsage {
121    fn from_bytes(data: &[u8], _hash: ObjectHash) -> Result<Self, GitError>
122    where
123        Self: Sized,
124    {
125        serde_json::from_slice(data).map_err(|e| GitError::InvalidObjectInfo(e.to_string()))
126    }
127
128    fn get_type(&self) -> ObjectType {
129        ObjectType::RunUsage
130    }
131
132    fn get_size(&self) -> usize {
133        match serde_json::to_vec(self) {
134            Ok(v) => v.len(),
135            Err(e) => {
136                tracing::warn!("failed to compute RunUsage size: {}", e);
137                0
138            }
139        }
140    }
141
142    fn to_data(&self) -> Result<Vec<u8>, GitError> {
143        serde_json::to_vec(self).map_err(|e| GitError::InvalidObjectInfo(e.to_string()))
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    // Coverage:
152    // - usage summary totals
153    // - consistency check
154    // - optional billing estimate storage
155
156    #[test]
157    fn test_run_usage_fields() {
158        let actor = ActorRef::agent("planner").expect("actor");
159        let usage = RunUsage::new(actor, Uuid::from_u128(0x1), 100, 40, Some(0.12)).expect("usage");
160
161        assert_eq!(usage.input_tokens(), 100);
162        assert_eq!(usage.output_tokens(), 40);
163        assert_eq!(usage.total_tokens(), 140);
164        assert!(usage.is_consistent());
165        assert_eq!(usage.cost_usd(), Some(0.12));
166    }
167}