Skip to content

Commit 4fe14fe

Browse files
authored
feat(ai): Conditionally set total_cost and total_tokens attributes (#4868)
Currently we always set `total_cost` and `total_tokens` attributes on every AI span, even if `input_tokens` or `output_tokens` are not present. At the moment, that is not the behavior we want, since for example we have spans that are just invoking agent and it has no cost nor any tokens are used, so setting this attributes is wrong. This also shows up then in UI and additionally confusing users
1 parent 6b0cc7f commit 4fe14fe

File tree

3 files changed

+91
-18
lines changed

3 files changed

+91
-18
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
- Derive a `sentry.description` attribute for V2 spans ([#4832](https://github.com/getsentry/relay/pull/4832))
1818
- Consider `gen_ai` also as AI span op prefix. ([#4859](https://github.com/getsentry/relay/pull/4859))
1919
- Change pii scrubbing on some AI attributes to optional ([#4860](https://github.com/getsentry/relay/pull/4860))
20+
- Conditionally set `total_cost` and `total_tokens` attributes on AI spans. ([#4868](https://github.com/getsentry/relay/pull/4868))
2021

2122
## 25.6.1
2223

relay-event-normalization/src/event.rs

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2410,6 +2410,72 @@ mod tests {
24102410
);
24112411
}
24122412

2413+
#[test]
2414+
fn test_ai_data_with_no_tokens() {
2415+
let json = r#"
2416+
{
2417+
"spans": [
2418+
{
2419+
"timestamp": 1702474613.0495,
2420+
"start_timestamp": 1702474613.0175,
2421+
"description": "OpenAI ",
2422+
"op": "gen_ai.invoke_agent",
2423+
"span_id": "9c01bd820a083e63",
2424+
"parent_span_id": "a1e13f3f06239d69",
2425+
"trace_id": "922dda2462ea4ac2b6a4b339bee90863",
2426+
"data": {
2427+
"gen_ai.request.model": "claude-2.1"
2428+
}
2429+
}
2430+
]
2431+
}
2432+
"#;
2433+
2434+
let mut event = Annotated::<Event>::from_json(json).unwrap();
2435+
2436+
normalize_event(
2437+
&mut event,
2438+
&NormalizationConfig {
2439+
ai_model_costs: Some(&ModelCosts {
2440+
version: 2,
2441+
costs: vec![],
2442+
models: HashMap::from([(
2443+
"claude-2.1".to_owned(),
2444+
ModelCostV2 {
2445+
input_per_token: 0.01,
2446+
output_per_token: 0.02,
2447+
output_reasoning_per_token: 0.03,
2448+
input_cached_per_token: 0.0,
2449+
},
2450+
)]),
2451+
}),
2452+
..NormalizationConfig::default()
2453+
},
2454+
);
2455+
2456+
let spans = event.value().unwrap().spans.value().unwrap();
2457+
2458+
assert_eq!(spans.len(), 1);
2459+
// total_cost shouldn't be set if no tokens are present on span data
2460+
assert_eq!(
2461+
spans
2462+
.first()
2463+
.and_then(|span| span.value())
2464+
.and_then(|span| span.data.value())
2465+
.and_then(|data| data.gen_ai_usage_total_cost.value()),
2466+
None
2467+
);
2468+
// total_tokens shouldn't be set if no tokens are present on span data
2469+
assert_eq!(
2470+
spans
2471+
.first()
2472+
.and_then(|span| span.value())
2473+
.and_then(|span| span.data.value())
2474+
.and_then(|data| data.gen_ai_usage_total_tokens.value()),
2475+
None
2476+
);
2477+
}
2478+
24132479
#[test]
24142480
fn test_ai_data_with_ai_op_prefix() {
24152481
let json = r#"

relay-event-normalization/src/normalize/span/ai.rs

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -11,31 +11,32 @@ fn calculate_ai_model_cost(model_cost: Option<ModelCostV2>, data: &SpanData) ->
1111
let input_tokens_used = data
1212
.gen_ai_usage_input_tokens
1313
.value()
14-
.and_then(Value::as_f64)
15-
.unwrap_or(0.0);
14+
.and_then(Value::as_f64);
1615

1716
let output_tokens_used = data
1817
.gen_ai_usage_output_tokens
1918
.value()
20-
.and_then(Value::as_f64)
21-
.unwrap_or(0.0);
19+
.and_then(Value::as_f64);
2220
let output_reasoning_tokens_used = data
2321
.gen_ai_usage_output_tokens_reasoning
2422
.value()
25-
.and_then(Value::as_f64)
26-
.unwrap_or(0.0);
23+
.and_then(Value::as_f64);
2724
let input_cached_tokens_used = data
2825
.gen_ai_usage_input_tokens_cached
2926
.value()
30-
.and_then(Value::as_f64)
31-
.unwrap_or(0.0);
27+
.and_then(Value::as_f64);
28+
29+
if input_tokens_used.is_none() && output_tokens_used.is_none() {
30+
return None;
31+
}
3232

3333
let mut result = 0.0;
3434

35-
result += cost_per_token.input_per_token * input_tokens_used;
36-
result += cost_per_token.output_per_token * output_tokens_used;
37-
result += cost_per_token.output_reasoning_per_token * output_reasoning_tokens_used;
38-
result += cost_per_token.input_cached_per_token * input_cached_tokens_used;
35+
result += cost_per_token.input_per_token * input_tokens_used.unwrap_or(0.0);
36+
result += cost_per_token.output_per_token * output_tokens_used.unwrap_or(0.0);
37+
result +=
38+
cost_per_token.output_reasoning_per_token * output_reasoning_tokens_used.unwrap_or(0.0);
39+
result += cost_per_token.input_cached_per_token * input_cached_tokens_used.unwrap_or(0.0);
3940

4041
Some(result)
4142
}
@@ -72,15 +73,20 @@ pub fn map_ai_measurements_to_data(span: &mut Span) {
7273
let input_tokens = data
7374
.gen_ai_usage_input_tokens
7475
.value()
75-
.and_then(Value::as_f64)
76-
.unwrap_or(0.0);
76+
.and_then(Value::as_f64);
7777
let output_tokens = data
7878
.gen_ai_usage_output_tokens
7979
.value()
80-
.and_then(Value::as_f64)
81-
.unwrap_or(0.0);
82-
data.gen_ai_usage_total_tokens
83-
.set_value(Value::F64(input_tokens + output_tokens).into());
80+
.and_then(Value::as_f64);
81+
82+
if input_tokens.is_none() && output_tokens.is_none() {
83+
// don't set total_tokens if there are no input nor output tokens
84+
return;
85+
}
86+
87+
data.gen_ai_usage_total_tokens.set_value(
88+
Value::F64(input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0)).into(),
89+
);
8490
}
8591
}
8692

0 commit comments

Comments
 (0)