feat(spanv2): Implement AI normalizations

Author: Dav1dde

relay-conventions/src/consts.rs

@@ -21,7 +21,18 @@ convention_attributes!(
     DB_SYSTEM_NAME => "db.system.name",
     DESCRIPTION => "sentry.description",
     FAAS_TRIGGER => "faas.trigger",
+    GEN_AI_COST_INPUT_TOKENS => "gen_ai.cost.input_tokens",
+    GEN_AI_COST_OUTPUT_TOKENS => "gen_ai.cost.output_tokens",
+    GEN_AI_COST_TOTAL_TOKENS => "gen_ai.cost.total_tokens",
+    GEN_AI_REQUEST_MODEL => "gen_ai.request.model",
+    GEN_AI_RESPONSE_MODEL => "gen_ai.response.model",
+    GEN_AI_RESPONSE_TPS => "gen_ai.response.tokens_per_second",
     GEN_AI_SYSTEM => "gen_ai.system",
+    GEN_AI_USAGE_INPUT_CACHED_TOKENS => "gen_ai.usage.input_tokens.cached",
+    GEN_AI_USAGE_INPUT_TOKENS => "gen_ai.usage.input_tokens",
+    GEN_AI_USAGE_OUTPUT_REASONING_TOKENS => "gen_ai.usage.output_tokens.reasoning",
+    GEN_AI_USAGE_OUTPUT_TOKENS => "gen_ai.usage.output_tokens",
+    GEN_AI_USAGE_TOTAL_TOKENS => "gen_ai.usage.total_tokens",
     HTTP_PREFETCH => "sentry.http.prefetch",
     HTTP_REQUEST_METHOD => "http.request.method",
     HTTP_RESPONSE_STATUS_CODE => "http.response.status_code",

relay-event-normalization/src/eap/ai.rs

@@ -0,0 +1,362 @@
+use std::time::Duration;
+
+use relay_conventions::{
+    GEN_AI_COST_INPUT_TOKENS, GEN_AI_COST_OUTPUT_TOKENS, GEN_AI_COST_TOTAL_TOKENS,
+    GEN_AI_REQUEST_MODEL, GEN_AI_RESPONSE_MODEL, GEN_AI_RESPONSE_TPS,
+    GEN_AI_USAGE_INPUT_CACHED_TOKENS, GEN_AI_USAGE_INPUT_TOKENS,
+    GEN_AI_USAGE_OUTPUT_REASONING_TOKENS, GEN_AI_USAGE_OUTPUT_TOKENS, GEN_AI_USAGE_TOTAL_TOKENS,
+};
+use relay_event_schema::protocol::Attributes;
+use relay_protocol::Annotated;
+
+use crate::ModelCosts;
+use crate::span::ai;
+
+/// Normalizes AI attributes.
+pub fn normalize_ai(
+    attributes: &mut Annotated<Attributes>,
+    duration: Option<Duration>,
+    costs: Option<&ModelCosts>,
+) {
+    let Some(attributes) = attributes.value_mut() else {
+        return;
+    };
+
+    normalize_total_tokens(attributes);
+    normalize_tokens_per_second(attributes, duration);
+    normalize_ai_costs(attributes, costs);
+}
+
+/// Calculates the [`GEN_AI_USAGE_TOTAL_TOKENS`] attribute.
+fn normalize_total_tokens(attributes: &mut Attributes) {
+    if attributes.contains_key(GEN_AI_USAGE_TOTAL_TOKENS) {
+        return;
+    }
+
+    let input_tokens = attributes
+        .get_value(GEN_AI_USAGE_INPUT_TOKENS)
+        .and_then(|v| v.as_u64());
+
+    let output_tokens = attributes
+        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
+        .and_then(|v| v.as_u64());
+
+    if input_tokens.is_none() && output_tokens.is_none() {
+        return;
+    }
+
+    let total_tokens = input_tokens.unwrap_or(0) + output_tokens.unwrap_or(0);
+    if let Ok(total_tokens) = i64::try_from(total_tokens) {
+        attributes.insert(GEN_AI_USAGE_TOTAL_TOKENS, total_tokens);
+    }
+}
+
+/// Calculates the [`GEN_AI_RESPONSE_TPS`] attribute.
+fn normalize_tokens_per_second(attributes: &mut Attributes, duration: Option<Duration>) {
+    let Some(duration) = duration.filter(|d| !d.is_zero()) else {
+        return;
+    };
+
+    if attributes.contains_key(GEN_AI_RESPONSE_TPS) {
+        return;
+    }
+
+    let output_tokens = attributes
+        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
+        .and_then(|v| v.as_f64())
+        .filter(|v| *v > 0.0);
+
+    if let Some(output_tokens) = output_tokens {
+        let tps = output_tokens / duration.as_secs_f64();
+        attributes.insert(GEN_AI_RESPONSE_TPS, tps);
+    }
+}
+
+/// Calculates model costs and serializes them into attributes.
+fn normalize_ai_costs(attributes: &mut Attributes, model_costs: Option<&ModelCosts>) {
+    if attributes.contains_key(GEN_AI_COST_TOTAL_TOKENS) {
+        return;
+    }
+
+    let model_cost = attributes
+        .get_value(GEN_AI_REQUEST_MODEL)
+        .or_else(|| attributes.get_value(GEN_AI_RESPONSE_MODEL))
+        .and_then(|v| v.as_str())
+        .and_then(|model| model_costs?.cost_per_token(model));
+
+    let Some(model_cost) = model_cost else { return };
+
+    let get_tokens = |key| {
+        attributes
+            .get_value(key)
+            .and_then(|v| v.as_f64())
+            .unwrap_or(0.0)
+    };
+
+    let tokens = ai::UsedTokens {
+        input_tokens: get_tokens(GEN_AI_USAGE_INPUT_TOKENS),
+        input_cached_tokens: get_tokens(GEN_AI_USAGE_INPUT_CACHED_TOKENS),
+        output_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_TOKENS),
+        output_reasoning_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_REASONING_TOKENS),
+    };
+
+    let Some(costs) = ai::calculate_costs(model_cost, tokens) else {
+        return;
+    };
+
+    attributes.insert(GEN_AI_COST_INPUT_TOKENS, costs.input);
+    attributes.insert(GEN_AI_COST_OUTPUT_TOKENS, costs.output);
+    attributes.insert(GEN_AI_COST_TOTAL_TOKENS, costs.total());
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use relay_pattern::Pattern;
+    use relay_protocol::assert_annotated_snapshot;
+
+    use crate::ModelCostV2;
+
+    use super::*;
+
+    macro_rules! attributes {
+        ($($key:expr => $value:expr),* $(,)?) => {
+            Attributes::from([
+                $(($key.into(), Annotated::new($value.into())),)*
+            ])
+        };
+    }
+
+    fn model_costs() -> ModelCosts {
+        ModelCosts {
+            version: 2,
+            models: HashMap::from([
+                (
+                    Pattern::new("claude-2.1").unwrap(),
+                    ModelCostV2 {
+                        input_per_token: 0.01,
+                        output_per_token: 0.02,
+                        output_reasoning_per_token: 0.03,
+                        input_cached_per_token: 0.04,
+                    },
+                ),
+                (
+                    Pattern::new("gpt4-21-04").unwrap(),
+                    ModelCostV2 {
+                        input_per_token: 0.09,
+                        output_per_token: 0.05,
+                        output_reasoning_per_token: 0.0,
+                        input_cached_per_token: 0.0,
+                    },
+                ),
+            ]),
+        }
+    }
+
+    #[test]
+    fn test_normalize_ai_all_tokens() {
+        let mut attributes = Annotated::new(attributes! {
+            "gen_ai.usage.input_tokens" => 1000,
+            "gen_ai.usage.output_tokens" => 2000,
+            "gen_ai.usage.output_tokens.reasoning" => 1000,
+            "gen_ai.usage.input_tokens.cached" => 500,
+            "gen_ai.request.model" => "claude-2.1".to_owned(),
+        });
+
+        normalize_ai(
+            &mut attributes,
+            Some(Duration::from_secs(1)),
+            Some(&model_costs()),
+        );
+
+        assert_annotated_snapshot!(attributes, @r#"
+        {
+          "gen_ai.cost.input_tokens": {
+            "type": "double",
+            "value": 25.0
+          },
+          "gen_ai.cost.output_tokens": {
+            "type": "double",
+            "value": 50.0
+          },
+          "gen_ai.cost.total_tokens": {
+            "type": "double",
+            "value": 75.0
+          },
+          "gen_ai.request.model": {
+            "type": "string",
+            "value": "claude-2.1"
+          },
+          "gen_ai.response.tokens_per_second": {
+            "type": "double",
+            "value": 2000.0
+          },
+          "gen_ai.usage.input_tokens": {
+            "type": "integer",
+            "value": 1000
+          },
+          "gen_ai.usage.input_tokens.cached": {
+            "type": "integer",
+            "value": 500
+          },
+          "gen_ai.usage.output_tokens": {
+            "type": "integer",
+            "value": 2000
+          },
+          "gen_ai.usage.output_tokens.reasoning": {
+            "type": "integer",
+            "value": 1000
+          },
+          "gen_ai.usage.total_tokens": {
+            "type": "integer",
+            "value": 3000
+          }
+        }
+        "#);
+    }
+
+    #[test]
+    fn test_normalize_ai_basic_tokens() {
+        let mut attributes = Annotated::new(attributes! {
+            "gen_ai.usage.input_tokens" => 1000,
+            "gen_ai.usage.output_tokens" => 2000,
+            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
+        });
+
+        normalize_ai(
+            &mut attributes,
+            Some(Duration::from_millis(500)),
+            Some(&model_costs()),
+        );
+
+        assert_annotated_snapshot!(attributes, @r#"
+        {
+          "gen_ai.cost.input_tokens": {
+            "type": "double",
+            "value": 90.0
+          },
+          "gen_ai.cost.output_tokens": {
+            "type": "double",
+            "value": 100.0
+          },
+          "gen_ai.cost.total_tokens": {
+            "type": "double",
+            "value": 190.0
+          },
+          "gen_ai.request.model": {
+            "type": "string",
+            "value": "gpt4-21-04"
+          },
+          "gen_ai.response.tokens_per_second": {
+            "type": "double",
+            "value": 4000.0
+          },
+          "gen_ai.usage.input_tokens": {
+            "type": "integer",
+            "value": 1000
+          },
+          "gen_ai.usage.output_tokens": {
+            "type": "integer",
+            "value": 2000
+          },
+          "gen_ai.usage.total_tokens": {
+            "type": "integer",
+            "value": 3000
+          }
+        }
+        "#);
+    }
+
+    #[test]
+    fn test_normalize_ai_basic_tokens_no_duration_no_cost() {
+        let mut attributes = Annotated::new(attributes! {
+            "gen_ai.usage.input_tokens" => 1000,
+            "gen_ai.usage.output_tokens" => 2000,
+            "gen_ai.request.model" => "unknown".to_owned(),
+        });
+
+        normalize_ai(&mut attributes, Some(Duration::ZERO), Some(&model_costs()));
+
+        assert_annotated_snapshot!(attributes, @r#"
+        {
+          "gen_ai.request.model": {
+            "type": "string",
+            "value": "unknown"
+          },
+          "gen_ai.usage.input_tokens": {
+            "type": "integer",
+            "value": 1000
+          },
+          "gen_ai.usage.output_tokens": {
+            "type": "integer",
+            "value": 2000
+          },
+          "gen_ai.usage.total_tokens": {
+            "type": "integer",
+            "value": 3000
+          }
+        }
+        "#);
+    }
+
+    #[test]
+    fn test_normalize_ai_does_not_override() {
+        let mut attributes = Annotated::new(attributes! {
+            "gen_ai.usage.input_tokens" => 1000,
+            "gen_ai.usage.output_tokens" => 2000,
+            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
+
+            "gen_ai.usage.total_tokens" => 1337,
+
+            "gen_ai.cost.input_tokens" => 99.0,
+            "gen_ai.cost.output_tokens" => 99.0,
+            "gen_ai.cost.total_tokens" => 123.0,
+
+            "gen_ai.response.tokens_per_second" => 42.0,
+        });
+
+        normalize_ai(
+            &mut attributes,
+            Some(Duration::from_millis(500)),
+            Some(&model_costs()),
+        );
+
+        assert_annotated_snapshot!(attributes, @r#"
+        {
+          "gen_ai.cost.input_tokens": {
+            "type": "double",
+            "value": 99.0
+          },
+          "gen_ai.cost.output_tokens": {
+            "type": "double",
+            "value": 99.0
+          },
+          "gen_ai.cost.total_tokens": {
+            "type": "double",
+            "value": 123.0
+          },
+          "gen_ai.request.model": {
+            "type": "string",
+            "value": "gpt4-21-04"
+          },
+          "gen_ai.response.tokens_per_second": {
+            "type": "double",
+            "value": 42.0
+          },
+          "gen_ai.usage.input_tokens": {
+            "type": "integer",
+            "value": 1000
+          },
+          "gen_ai.usage.output_tokens": {
+            "type": "integer",
+            "value": 2000
+          },
+          "gen_ai.usage.total_tokens": {
+            "type": "integer",
+            "value": 1337
+          }
+        }
+        "#);
+    }
+}

relay-event-normalization/src/eap/mod.rs

@@ -14,6 +14,10 @@ use relay_sampling::DynamicSamplingContext;
 
 use crate::{ClientHints, FromUserAgentInfo as _, RawUserAgentInfo};
 
+mod ai;
+
+pub use self::ai::normalize_ai;
+
 /// Normalizes/validates all attribute types.
 ///
 /// Removes and marks all attributes with an error for which the specified [`AttributeType`]