feat(ai): data & measurement normalization

CHANGELOG.md

@@ -5,6 +5,7 @@
 **Internal**:
 
 - Remove the "unspecified" variant of `SpanKind`. ([#4774](https://github.com/getsentry/relay/pull/4774))
+- 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))
 - Switch `sysinfo` dependency back to upstream and update to 0.35.1. ([#4776](https://github.com/getsentry/relay/pull/4776))
 
 ## 25.5.1

relay-dynamic-config/src/defaults.rs

@@ -834,7 +834,7 @@ pub fn hardcoded_span_metrics() -> Vec<(GroupKey, Vec<MetricSpec>, Vec<TagMappin
                 MetricSpec {
                     category: DataCategory::Span,
                     mri: "c:spans/ai.total_tokens.used@none".into(),
-                    field: Some("span.measurements.ai_total_tokens_used.value".into()),
+                    field: Some("span.data.gen_ai\\.usage\\.total_tokens".into()),
                     condition: Some(is_ai.clone()),
                     tags: vec![
                         Tag::with_key("span.op")
@@ -869,7 +869,7 @@ pub fn hardcoded_span_metrics() -> Vec<(GroupKey, Vec<MetricSpec>, Vec<TagMappin
                 MetricSpec {
                     category: DataCategory::Span,
                     mri: "c:spans/ai.total_cost@usd".into(),
-                    field: Some("span.measurements.ai_total_cost.value".into()),
+                    field: Some("span.data.gen_ai\\.usage\\.total_cost".into()),
                     condition: Some(is_ai.clone()),
                     tags: vec![
                         Tag::with_key("span.op")

relay-event-normalization/src/event.rs

@@ -27,7 +27,7 @@ use smallvec::SmallVec;
 use uuid::Uuid;
 
 use crate::normalize::request;
-use crate::span::ai::normalize_ai_measurements;
+use crate::span::ai::enrich_ai_span_data;
 use crate::span::tag_extraction::extract_span_tags_from_event;
 use crate::utils::{self, MAX_DURATION_MOBILE_MS, get_event_user_tag};
 use crate::{
@@ -322,7 +322,7 @@ fn normalize(event: &mut Event, meta: &mut Meta, config: &NormalizationConfig) {
             .get_or_default::<PerformanceScoreContext>()
             .score_profile_version = Annotated::new(version);
     }
-    normalize_ai_measurements(event, config.ai_model_costs);
+    enrich_ai_span_data(event, config.ai_model_costs);
     normalize_breakdowns(event, config.breakdowns_config); // Breakdowns are part of the metric extraction too
     normalize_default_attributes(event, meta, config);
     normalize_trace_context_tags(event);
@@ -2190,7 +2190,7 @@ mod tests {
     }
 
     #[test]
-    fn test_ai_measurements() {
+    fn test_ai_legacy_measurements() {
         let json = r#"
             {
                 "spans": [
@@ -2271,26 +2271,113 @@ mod tests {
         assert_eq!(
             spans
                 .first()
-                .unwrap()
-                .value()
-                .unwrap()
-                .measurements
-                .value()
-                .unwrap()
-                .get_value("ai_total_cost"),
-            Some(1.23)
+                .and_then(|span| span.value())
+                .and_then(|span| span.data.value())
+                .and_then(|data| data.gen_ai_usage_total_cost.value()),
+            Some(&Value::F64(1.23))
         );
         assert_eq!(
             spans
                 .get(1)
-                .unwrap()
-                .value()
-                .unwrap()
-                .measurements
-                .value()
-                .unwrap()
-                .get_value("ai_total_cost"),
-            Some(20.0 * 2.0 + 2.0)
+                .and_then(|span| span.value())
+                .and_then(|span| span.data.value())
+                .and_then(|data| data.gen_ai_usage_total_cost.value()),
+            Some(&Value::F64(20.0 * 2.0 + 2.0))
+        );
+    }
+
+    #[test]
+    fn test_ai_data() {
+        let json = r#"
+            {
+                "spans": [
+                    {
+                        "timestamp": 1702474613.0495,
+                        "start_timestamp": 1702474613.0175,
+                        "description": "OpenAI ",
+                        "op": "ai.chat_completions.openai",
+                        "span_id": "9c01bd820a083e63",
+                        "parent_span_id": "a1e13f3f06239d69",
+                        "trace_id": "922dda2462ea4ac2b6a4b339bee90863",
+                        "data": {
+                            "gen_ai.usage.total_tokens": 1230,
+                            "ai.pipeline.name": "Autofix Pipeline",
+                            "ai.model_id": "claude-2.1"
+                        }
+                    },
+                    {
+                        "timestamp": 1702474613.0495,
+                        "start_timestamp": 1702474613.0175,
+                        "description": "OpenAI ",
+                        "op": "ai.chat_completions.openai",
+                        "span_id": "ac01bd820a083e63",
+                        "parent_span_id": "a1e13f3f06239d69",
+                        "trace_id": "922dda2462ea4ac2b6a4b339bee90863",
+                        "data": {
+                            "gen_ai.usage.input_tokens": 1000,
+                            "gen_ai.usage.output_tokens": 2000,
+                            "ai.pipeline.name": "Autofix Pipeline",
+                            "ai.model_id": "gpt4-21-04"
+                        }
+                    }
+                ]
+            }
+        "#;
+
+        let mut event = Annotated::<Event>::from_json(json).unwrap();
+
+        normalize_event(
+            &mut event,
+            &NormalizationConfig {
+                ai_model_costs: Some(&ModelCosts {
+                    version: 1,
+                    costs: vec![
+                        ModelCost {
+                            model_id: LazyGlob::new("claude-2*"),
+                            for_completion: false,
+                            cost_per_1k_tokens: 1.0,
+                        },
+                        ModelCost {
+                            model_id: LazyGlob::new("gpt4-21*"),
+                            for_completion: false,
+                            cost_per_1k_tokens: 2.0,
+                        },
+                        ModelCost {
+                            model_id: LazyGlob::new("gpt4-21*"),
+                            for_completion: true,
+                            cost_per_1k_tokens: 20.0,
+                        },
+                    ],
+                }),
+                ..NormalizationConfig::default()
+            },
+        );
+
+        let spans = event.value().unwrap().spans.value().unwrap();
+        assert_eq!(spans.len(), 2);
+        assert_eq!(
+            spans
+                .first()
+                .and_then(|span| span.value())
+                .and_then(|span| span.data.value())
+                .and_then(|data| data.gen_ai_usage_total_cost.value()),
+            Some(&Value::F64(1.23))
+        );
+        assert_eq!(
+            spans
+                .get(1)
+                .and_then(|span| span.value())
+                .and_then(|span| span.data.value())
+                .and_then(|data| data.gen_ai_usage_total_cost.value()),
+            Some(&Value::F64(20.0 * 2.0 + 2.0))
+        );
+        assert_eq!(
+            spans
+                .get(1)
+                .and_then(|span| span.value())
+                .and_then(|span| span.data.value())
+                .and_then(|data| data.gen_ai_usage_total_tokens.value()),
+            Some(&Value::F64(3000.0))
         );
     }
 

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

@@ -1,8 +1,8 @@
 //! AI cost calculation.
 
 use crate::ModelCosts;
-use relay_base_schema::metrics::MetricUnit;
-use relay_event_schema::protocol::{Event, Measurement, Span};
+use relay_event_schema::protocol::{Event, Span, SpanData};
+use relay_protocol::{Annotated, Value};
 
 /// Calculated cost is in US dollars.
 fn calculate_ai_model_cost(
@@ -33,59 +33,96 @@ fn calculate_ai_model_cost(
     }
 }
 
-/// Extract the ai_total_cost measurement into the span.
-pub fn extract_ai_measurements(span: &mut Span, ai_model_costs: &ModelCosts) {
-    let Some(span_op) = span.op.value() else {
+/// Maps AI-related measurements (legacy) to span data.
+pub fn map_ai_measurements_to_data(span: &mut Span) {
+    if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
         return;
     };
 
-    if !span_op.starts_with("ai.") {
+    let measurements = span.measurements.value();
+    let data = span.data.get_or_insert_with(SpanData::default);
+
+    let set_field_from_measurement = |target_field: &mut Annotated<Value>,
+                                      measurement_key: &str| {
+        if let Some(measurements) = measurements {
+            if target_field.value().is_none() {
+                if let Some(value) = measurements.get_value(measurement_key) {
+                    target_field.set_value(Value::F64(value).into());
+                }
+            }
+        }
+    };
+
+    set_field_from_measurement(&mut data.gen_ai_usage_total_tokens, "ai_total_tokens_used");
+    set_field_from_measurement(&mut data.gen_ai_usage_input_tokens, "ai_prompt_tokens_used");
+    set_field_from_measurement(
+        &mut data.gen_ai_usage_output_tokens,
+        "ai_completion_tokens_used",
+    );
+
+    // It might be that 'total_tokens' is not set in which case we need to calculate it
+    if data.gen_ai_usage_total_tokens.value().is_none() {
+        let input_tokens = data
+            .gen_ai_usage_input_tokens
+            .value()
+            .and_then(Value::as_f64)
+            .unwrap_or(0.0);
+        let output_tokens = data
+            .gen_ai_usage_output_tokens
+            .value()
+            .and_then(Value::as_f64)
+            .unwrap_or(0.0);
+        data.gen_ai_usage_total_tokens
+            .set_value(Value::F64(input_tokens + output_tokens).into());
+    }
+}
+
+/// Extract the gen_ai_usage_total_cost data into the span
+pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) {
+    if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
         return;
     }
 
-    let Some(measurements) = span.measurements.value() else {
+    let Some(data) = span.data.value_mut() else {
         return;
     };
 
-    let total_tokens_used = measurements.get_value("ai_total_tokens_used");
-    let prompt_tokens_used = measurements.get_value("ai_prompt_tokens_used");
-    let completion_tokens_used = measurements.get_value("ai_completion_tokens_used");
-    if let Some(model_id) = span
-        .data
+    let total_tokens_used = data
+        .gen_ai_usage_total_tokens
+        .value()
+        .and_then(Value::as_f64);
+    let prompt_tokens_used = data
+        .gen_ai_usage_input_tokens
+        .value()
+        .and_then(Value::as_f64);
+    let completion_tokens_used = data
+        .gen_ai_usage_output_tokens
         .value()
-        .and_then(|d| d.ai_model_id.value())
-        .and_then(|val| val.as_str())
-    {
+        .and_then(Value::as_f64);
+
+    if let Some(model_id) = data.ai_model_id.value().and_then(|val| val.as_str()) {
         if let Some(total_cost) = calculate_ai_model_cost(
             model_id,
             prompt_tokens_used,
             completion_tokens_used,
             total_tokens_used,
             ai_model_costs,
         ) {
-            span.measurements
-                .get_or_insert_with(Default::default)
-                .insert(
-                    "ai_total_cost".to_owned(),
-                    Measurement {
-                        value: total_cost.into(),
-                        unit: MetricUnit::None.into(),
-                    }
-                    .into(),
-                );
+            data.gen_ai_usage_total_cost
+                .set_value(Value::F64(total_cost).into());
         }
     }
 }
 
-/// Extract the ai_total_cost measurements from all of an event's spans
-pub fn normalize_ai_measurements(event: &mut Event, model_costs: Option<&ModelCosts>) {
-    if let Some(model_costs) = model_costs {
-        if let Some(spans) = event.spans.value_mut() {
-            for span in spans {
-                if let Some(mut_span) = span.value_mut() {
-                    extract_ai_measurements(mut_span, model_costs);
-                }
-            }
+/// Extract the ai data from all of an event's spans
+pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) {
+    let spans = event.spans.value_mut().iter_mut().flatten();
+    let spans = spans.filter_map(|span| span.value_mut().as_mut());
+
+    for span in spans {
+        map_ai_measurements_to_data(span);
+        if let Some(model_costs) = model_costs {
+            extract_ai_data(span, model_costs);
         }
     }
 }