git_internal/internal/object/
run_usage.rs1use 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#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
40#[serde(deny_unknown_fields)]
41pub struct RunUsage {
42 #[serde(flatten)]
45 header: Header,
46 run_id: Uuid,
48 input_tokens: u64,
50 output_tokens: u64,
52 total_tokens: u64,
54 #[serde(default, skip_serializing_if = "Option::is_none")]
56 cost_usd: Option<f64>,
57}
58
59impl RunUsage {
60 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 pub fn header(&self) -> &Header {
80 &self.header
81 }
82
83 pub fn run_id(&self) -> Uuid {
85 self.run_id
86 }
87
88 pub fn input_tokens(&self) -> u64 {
90 self.input_tokens
91 }
92
93 pub fn output_tokens(&self) -> u64 {
95 self.output_tokens
96 }
97
98 pub fn total_tokens(&self) -> u64 {
100 self.total_tokens
101 }
102
103 pub fn cost_usd(&self) -> Option<f64> {
105 self.cost_usd
106 }
107
108 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 #[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}