diff --git a/relay-event-normalization/src/normalize/mod.rs b/relay-event-normalization/src/normalize/mod.rs index a8637234b28..bfd49c5e1e9 100644 --- a/relay-event-normalization/src/normalize/mod.rs +++ b/relay-event-normalization/src/normalize/mod.rs @@ -261,27 +261,24 @@ impl ModelCosts { } /// Gets the cost per token, if defined for the given model. - pub fn cost_per_token(&self, model_id: &str) -> Option { + pub fn cost_per_token(&self, model_id: &str) -> Option<&ModelCostV2> { if !self.is_enabled() { return None; } // First try exact match by creating a Pattern from the model_id - let exact_key = Pattern::new(model_id).ok()?; - if self.models.contains_key(&exact_key) { - return self.models.get(&exact_key).copied(); + if let Some(value) = self.models.get(model_id) { + return Some(value); } // if there is not a direct match, try to find the match using a pattern - let cost = self.models.iter().find_map(|(key, value)| { + self.models.iter().find_map(|(key, value)| { if key.is_match(model_id) { Some(value) } else { None } - }); - - cost.copied() + }) } } @@ -481,7 +478,7 @@ mod tests { let cost = v2_config.cost_per_token("gpt-4").unwrap(); assert_eq!( cost, - ModelCostV2 { + &ModelCostV2 { input_per_token: 0.03, output_per_token: 0.06, output_reasoning_per_token: 0.12, @@ -523,7 +520,7 @@ mod tests { let cost = v2_config.cost_per_token("gpt-4-v1").unwrap(); assert_eq!( cost, - ModelCostV2 { + &ModelCostV2 { input_per_token: 0.03, output_per_token: 0.06, output_reasoning_per_token: 0.12, @@ -534,7 +531,7 @@ mod tests { let cost = v2_config.cost_per_token("gpt-4-2xxx").unwrap(); assert_eq!( cost, - ModelCostV2 { + &ModelCostV2 { input_per_token: 0.0007, output_per_token: 0.0008, output_reasoning_per_token: 0.0016, diff --git a/relay-event-normalization/src/normalize/span/ai.rs b/relay-event-normalization/src/normalize/span/ai.rs index 7041f7b2bba..fb4148cec2a 100644 --- a/relay-event-normalization/src/normalize/span/ai.rs +++ b/relay-event-normalization/src/normalize/span/ai.rs @@ -7,7 +7,7 @@ use relay_protocol::{Annotated, Getter, Value}; /// Calculates the cost of an AI model based on the model cost and the tokens used. /// Calculated cost is in US dollars. -fn extract_ai_model_cost_data(model_cost: Option, data: &mut SpanData) { +fn extract_ai_model_cost_data(model_cost: Option<&ModelCostV2>, data: &mut SpanData) { let cost_per_token = match model_cost { Some(v) => v, None => return, @@ -123,9 +123,7 @@ fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) { .and_then(|v| v.as_f64()) .unwrap_or(0.0); - let Some(data) = span.data.value_mut() else { - return; - }; + let data = span.data.get_or_insert_with(SpanData::default); // Extracts the response tokens per second if data.gen_ai_response_tokens_per_second.value().is_none()