use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct Usage {
pub completion_tokens: u64,
pub prompt_tokens: u64,
pub total_tokens: u64,
#[serde(skip_serializing_if = "Option::is_none")]
pub completion_tokens_details:
Option<objectiveai_sdk::agent::completions::response::CompletionTokensDetails>,
#[serde(skip_serializing_if = "Option::is_none")]
pub prompt_tokens_details:
Option<objectiveai_sdk::agent::completions::response::PromptTokensDetails>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cost: Option<rust_decimal::Decimal>,
#[serde(skip_serializing_if = "Option::is_none")]
pub cost_details: Option<CostDetails>,
}
impl Usage {
pub fn into_downstream(
self,
is_byok: bool,
cost_multiplier: rust_decimal::Decimal,
) -> objectiveai_sdk::agent::completions::response::UpstreamUsage {
let upstream_inference_cost = self.cost.unwrap_or_default();
let upstream_upstream_inference_cost = self
.cost_details
.unwrap_or_default()
.upstream_inference_cost
.unwrap_or_default();
let upstream_total_cost = upstream_inference_cost + upstream_upstream_inference_cost;
let total_cost = upstream_total_cost * cost_multiplier;
let (cost, cost_details, total_cost) = if is_byok {
(
total_cost - upstream_total_cost,
Some(objectiveai_sdk::agent::completions::response::CostDetails {
upstream_inference_cost,
upstream_upstream_inference_cost,
}),
total_cost,
)
} else {
(total_cost, None, total_cost)
};
objectiveai_sdk::agent::completions::response::UpstreamUsage {
completion_tokens: self.completion_tokens,
prompt_tokens: self.prompt_tokens,
total_tokens: self.total_tokens,
completion_tokens_details: self.completion_tokens_details,
prompt_tokens_details: self.prompt_tokens_details,
cost,
cost_details,
total_cost,
cost_multiplier,
is_byok,
}
}
pub fn push(&mut self, other: &Usage) {
self.completion_tokens += other.completion_tokens;
self.prompt_tokens += other.prompt_tokens;
self.total_tokens += other.total_tokens;
match (
&mut self.completion_tokens_details,
&other.completion_tokens_details,
) {
(Some(self_details), Some(other_details)) => {
self_details.push(other_details);
}
(None, Some(other_details)) => {
self.completion_tokens_details = Some(other_details.clone());
}
_ => {}
}
match (
&mut self.prompt_tokens_details,
&other.prompt_tokens_details,
) {
(Some(self_details), Some(other_details)) => {
self_details.push(other_details);
}
(None, Some(other_details)) => {
self.prompt_tokens_details = Some(other_details.clone());
}
_ => {}
}
objectiveai_sdk::agent::completions::response::util::push_option_decimal(
&mut self.cost,
&other.cost,
);
match (&mut self.cost_details, &other.cost_details) {
(Some(self_cost_details), Some(other_cost_details)) => {
self_cost_details.push(other_cost_details);
}
(None, Some(other_cost_details)) => {
self.cost_details = Some(other_cost_details.clone());
}
_ => {}
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct CostDetails {
#[serde(skip_serializing_if = "Option::is_none")]
pub upstream_inference_cost: Option<rust_decimal::Decimal>,
}
impl CostDetails {
pub fn push(&mut self, other: &CostDetails) {
objectiveai_sdk::agent::completions::response::util::push_option_decimal(
&mut self.upstream_inference_cost,
&other.upstream_inference_cost,
);
}
}