Skip to content

Commit 4081fe2

Browse files
tobias-wilferttobias-wilfert
authored
feat(ai): data & measurement normalization (#4768)
- Maps old and depreciated `measurements` and `data` to new fields in line with our conventions. - Updates cost calculation and metric extraction to use `data` rather than `measurements`.
1 parent 1a6a4c4 commit 4081fe2

File tree

13 files changed

+320
-7720
lines changed

13 files changed

+320
-7720
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
**Internal**:
1212

1313
- Remove the "unspecified" variant of `SpanKind`. ([#4774](https://github.com/getsentry/relay/pull/4774))
14+
- Normalize AI data and measurements into new OTEL compatible fields and extracting metrics out of said fields. ([#4768](https://github.com/getsentry/relay/pull/4768))
1415
- Switch `sysinfo` dependency back to upstream and update to 0.35.1. ([#4776](https://github.com/getsentry/relay/pull/4776))
1516

1617
## 25.5.1

‎relay-dynamic-config/src/defaults.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -834,7 +834,7 @@ pub fn hardcoded_span_metrics() -> Vec<(GroupKey, Vec<MetricSpec>, Vec<TagMappin
834834
MetricSpec {
835835
category: DataCategory::Span,
836836
mri: "c:spans/ai.total_tokens.used@none".into(),
837-
field: Some("span.measurements.ai_total_tokens_used.value".into()),
837+
field: Some("span.data.gen_ai\\.usage\\.total_tokens".into()),
838838
condition: Some(is_ai.clone()),
839839
tags: vec![
840840
Tag::with_key("span.op")
@@ -869,7 +869,7 @@ pub fn hardcoded_span_metrics() -> Vec<(GroupKey, Vec<MetricSpec>, Vec<TagMappin
869869
MetricSpec {
870870
category: DataCategory::Span,
871871
mri: "c:spans/ai.total_cost@usd".into(),
872-
field: Some("span.measurements.ai_total_cost.value".into()),
872+
field: Some("span.data.gen_ai\\.usage\\.total_cost".into()),
873873
condition: Some(is_ai.clone()),
874874
tags: vec![
875875
Tag::with_key("span.op")

‎relay-event-normalization/src/event.rs‎

Lines changed: 106 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use smallvec::SmallVec;
2727
use uuid::Uuid;
2828

2929
use crate::normalize::request;
30-
use crate::span::ai::normalize_ai_measurements;
30+
use crate::span::ai::enrich_ai_span_data;
3131
use crate::span::tag_extraction::extract_span_tags_from_event;
3232
use crate::utils::{self, MAX_DURATION_MOBILE_MS, get_event_user_tag};
3333
use crate::{
@@ -322,7 +322,7 @@ fn normalize(event: &mut Event, meta: &mut Meta, config: &NormalizationConfig) {
322322
.get_or_default::<PerformanceScoreContext>()
323323
.score_profile_version = Annotated::new(version);
324324
}
325-
normalize_ai_measurements(event, config.ai_model_costs);
325+
enrich_ai_span_data(event, config.ai_model_costs);
326326
normalize_breakdowns(event, config.breakdowns_config); // Breakdowns are part of the metric extraction too
327327
normalize_default_attributes(event, meta, config);
328328
normalize_trace_context_tags(event);
@@ -2190,7 +2190,7 @@ mod tests {
21902190
}
21912191

21922192
#[test]
2193-
fn test_ai_measurements() {
2193+
fn test_ai_legacy_measurements() {
21942194
let json = r#"
21952195
{
21962196
"spans": [
@@ -2271,26 +2271,113 @@ mod tests {
22712271
assert_eq!(
22722272
spans
22732273
.first()
2274-
.unwrap()
2275-
.value()
2276-
.unwrap()
2277-
.measurements
2278-
.value()
2279-
.unwrap()
2280-
.get_value("ai_total_cost"),
2281-
Some(1.23)
2274+
.and_then(|span| span.value())
2275+
.and_then(|span| span.data.value())
2276+
.and_then(|data| data.gen_ai_usage_total_cost.value()),
2277+
Some(&Value::F64(1.23))
22822278
);
22832279
assert_eq!(
22842280
spans
22852281
.get(1)
2286-
.unwrap()
2287-
.value()
2288-
.unwrap()
2289-
.measurements
2290-
.value()
2291-
.unwrap()
2292-
.get_value("ai_total_cost"),
2293-
Some(20.0 * 2.0 + 2.0)
2282+
.and_then(|span| span.value())
2283+
.and_then(|span| span.data.value())
2284+
.and_then(|data| data.gen_ai_usage_total_cost.value()),
2285+
Some(&Value::F64(20.0 * 2.0 + 2.0))
2286+
);
2287+
}
2288+
2289+
#[test]
2290+
fn test_ai_data() {
2291+
let json = r#"
2292+
{
2293+
"spans": [
2294+
{
2295+
"timestamp": 1702474613.0495,
2296+
"start_timestamp": 1702474613.0175,
2297+
"description": "OpenAI ",
2298+
"op": "ai.chat_completions.openai",
2299+
"span_id": "9c01bd820a083e63",
2300+
"parent_span_id": "a1e13f3f06239d69",
2301+
"trace_id": "922dda2462ea4ac2b6a4b339bee90863",
2302+
"data": {
2303+
"gen_ai.usage.total_tokens": 1230,
2304+
"ai.pipeline.name": "Autofix Pipeline",
2305+
"ai.model_id": "claude-2.1"
2306+
}
2307+
},
2308+
{
2309+
"timestamp": 1702474613.0495,
2310+
"start_timestamp": 1702474613.0175,
2311+
"description": "OpenAI ",
2312+
"op": "ai.chat_completions.openai",
2313+
"span_id": "ac01bd820a083e63",
2314+
"parent_span_id": "a1e13f3f06239d69",
2315+
"trace_id": "922dda2462ea4ac2b6a4b339bee90863",
2316+
"data": {
2317+
"gen_ai.usage.input_tokens": 1000,
2318+
"gen_ai.usage.output_tokens": 2000,
2319+
"ai.pipeline.name": "Autofix Pipeline",
2320+
"ai.model_id": "gpt4-21-04"
2321+
}
2322+
}
2323+
]
2324+
}
2325+
"#;
2326+
2327+
let mut event = Annotated::<Event>::from_json(json).unwrap();
2328+
2329+
normalize_event(
2330+
&mut event,
2331+
&NormalizationConfig {
2332+
ai_model_costs: Some(&ModelCosts {
2333+
version: 1,
2334+
costs: vec![
2335+
ModelCost {
2336+
model_id: LazyGlob::new("claude-2*"),
2337+
for_completion: false,
2338+
cost_per_1k_tokens: 1.0,
2339+
},
2340+
ModelCost {
2341+
model_id: LazyGlob::new("gpt4-21*"),
2342+
for_completion: false,
2343+
cost_per_1k_tokens: 2.0,
2344+
},
2345+
ModelCost {
2346+
model_id: LazyGlob::new("gpt4-21*"),
2347+
for_completion: true,
2348+
cost_per_1k_tokens: 20.0,
2349+
},
2350+
],
2351+
}),
2352+
..NormalizationConfig::default()
2353+
},
2354+
);
2355+
2356+
let spans = event.value().unwrap().spans.value().unwrap();
2357+
assert_eq!(spans.len(), 2);
2358+
assert_eq!(
2359+
spans
2360+
.first()
2361+
.and_then(|span| span.value())
2362+
.and_then(|span| span.data.value())
2363+
.and_then(|data| data.gen_ai_usage_total_cost.value()),
2364+
Some(&Value::F64(1.23))
2365+
);
2366+
assert_eq!(
2367+
spans
2368+
.get(1)
2369+
.and_then(|span| span.value())
2370+
.and_then(|span| span.data.value())
2371+
.and_then(|data| data.gen_ai_usage_total_cost.value()),
2372+
Some(&Value::F64(20.0 * 2.0 + 2.0))
2373+
);
2374+
assert_eq!(
2375+
spans
2376+
.get(1)
2377+
.and_then(|span| span.value())
2378+
.and_then(|span| span.data.value())
2379+
.and_then(|data| data.gen_ai_usage_total_tokens.value()),
2380+
Some(&Value::F64(3000.0))
22942381
);
22952382
}
22962383

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

Lines changed: 71 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
//! AI cost calculation.
22
33
use crate::ModelCosts;
4-
use relay_base_schema::metrics::MetricUnit;
5-
use relay_event_schema::protocol::{Event, Measurement, Span};
4+
use relay_event_schema::protocol::{Event, Span, SpanData};
5+
use relay_protocol::{Annotated, Value};
66

77
/// Calculated cost is in US dollars.
88
fn calculate_ai_model_cost(
@@ -33,59 +33,96 @@ fn calculate_ai_model_cost(
3333
}
3434
}
3535

36-
/// Extract the ai_total_cost measurement into the span.
37-
pub fn extract_ai_measurements(span: &mut Span, ai_model_costs: &ModelCosts) {
38-
let Some(span_op) = span.op.value() else {
36+
/// Maps AI-related measurements (legacy) to span data.
37+
pub fn map_ai_measurements_to_data(span: &mut Span) {
38+
if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
3939
return;
4040
};
4141

42-
if !span_op.starts_with("ai.") {
42+
let measurements = span.measurements.value();
43+
let data = span.data.get_or_insert_with(SpanData::default);
44+
45+
let set_field_from_measurement = |target_field: &mut Annotated<Value>,
46+
measurement_key: &str| {
47+
if let Some(measurements) = measurements {
48+
if target_field.value().is_none() {
49+
if let Some(value) = measurements.get_value(measurement_key) {
50+
target_field.set_value(Value::F64(value).into());
51+
}
52+
}
53+
}
54+
};
55+
56+
set_field_from_measurement(&mut data.gen_ai_usage_total_tokens, "ai_total_tokens_used");
57+
set_field_from_measurement(&mut data.gen_ai_usage_input_tokens, "ai_prompt_tokens_used");
58+
set_field_from_measurement(
59+
&mut data.gen_ai_usage_output_tokens,
60+
"ai_completion_tokens_used",
61+
);
62+
63+
// It might be that 'total_tokens' is not set in which case we need to calculate it
64+
if data.gen_ai_usage_total_tokens.value().is_none() {
65+
let input_tokens = data
66+
.gen_ai_usage_input_tokens
67+
.value()
68+
.and_then(Value::as_f64)
69+
.unwrap_or(0.0);
70+
let output_tokens = data
71+
.gen_ai_usage_output_tokens
72+
.value()
73+
.and_then(Value::as_f64)
74+
.unwrap_or(0.0);
75+
data.gen_ai_usage_total_tokens
76+
.set_value(Value::F64(input_tokens + output_tokens).into());
77+
}
78+
}
79+
80+
/// Extract the gen_ai_usage_total_cost data into the span
81+
pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) {
82+
if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
4383
return;
4484
}
4585

46-
let Some(measurements) = span.measurements.value() else {
86+
let Some(data) = span.data.value_mut() else {
4787
return;
4888
};
4989

50-
let total_tokens_used = measurements.get_value("ai_total_tokens_used");
51-
let prompt_tokens_used = measurements.get_value("ai_prompt_tokens_used");
52-
let completion_tokens_used = measurements.get_value("ai_completion_tokens_used");
53-
if let Some(model_id) = span
54-
.data
90+
let total_tokens_used = data
91+
.gen_ai_usage_total_tokens
92+
.value()
93+
.and_then(Value::as_f64);
94+
let prompt_tokens_used = data
95+
.gen_ai_usage_input_tokens
96+
.value()
97+
.and_then(Value::as_f64);
98+
let completion_tokens_used = data
99+
.gen_ai_usage_output_tokens
55100
.value()
56-
.and_then(|d| d.ai_model_id.value())
57-
.and_then(|val| val.as_str())
58-
{
101+
.and_then(Value::as_f64);
102+
103+
if let Some(model_id) = data.ai_model_id.value().and_then(|val| val.as_str()) {
59104
if let Some(total_cost) = calculate_ai_model_cost(
60105
model_id,
61106
prompt_tokens_used,
62107
completion_tokens_used,
63108
total_tokens_used,
64109
ai_model_costs,
65110
) {
66-
span.measurements
67-
.get_or_insert_with(Default::default)
68-
.insert(
69-
"ai_total_cost".to_owned(),
70-
Measurement {
71-
value: total_cost.into(),
72-
unit: MetricUnit::None.into(),
73-
}
74-
.into(),
75-
);
111+
data.gen_ai_usage_total_cost
112+
.set_value(Value::F64(total_cost).into());
76113
}
77114
}
78115
}
79116

80-
/// Extract the ai_total_cost measurements from all of an event's spans
81-
pub fn normalize_ai_measurements(event: &mut Event, model_costs: Option<&ModelCosts>) {
82-
if let Some(model_costs) = model_costs {
83-
if let Some(spans) = event.spans.value_mut() {
84-
for span in spans {
85-
if let Some(mut_span) = span.value_mut() {
86-
extract_ai_measurements(mut_span, model_costs);
87-
}
88-
}
117+
/// Extract the ai data from all of an event's spans
118+
pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) {
119+
let spans = event.spans.value_mut().iter_mut().flatten();
120+
let spans = spans.filter_map(|span| span.value_mut().as_mut());
121+
122+
for span in spans {
123+
map_ai_measurements_to_data(span);
124+
if let Some(model_costs) = model_costs {
125+
extract_ai_data(span, model_costs);
89126
}
90127
}
91128
}

0 commit comments

Comments
 (0)