relay_event_normalization/eap/
ai.rs

1use std::time::Duration;
2
3use relay_conventions::consts::*;
4use relay_event_schema::protocol::Attributes;
5use relay_protocol::Annotated;
6
7use crate::ModelCosts;
8use crate::span::ai;
9use crate::statsd::{Counters, map_origin_to_integration, platform_tag};
10
11/// Normalizes AI attributes.
12///
13/// This aggressively overwrites existing AI attributes, in order to guarantee a consistent data
14/// set for the AI product module.
15///
16/// As an example, an OTeL user may be manually instrumenting AI request costs on spans but in a
17/// local currency. Sentry's AI model requires a consistent cost value, independent of local
18/// currencies.
19///
20/// Callers may choose to only run this normalization in processing mode to not have the
21/// normalization run multiple times.
22pub fn normalize_ai(
23    attributes: &mut Annotated<Attributes>,
24    duration: Option<Duration>,
25    costs: Option<&ModelCosts>,
26) {
27    let Some(attributes) = attributes.value_mut() else {
28        return;
29    };
30
31    // Specifically only apply normalizations if the item is recognized as an AI item by the
32    // product.
33    if !is_ai_item(attributes) {
34        return;
35    }
36
37    normalize_model(attributes);
38    normalize_ai_type(attributes);
39    normalize_total_tokens(attributes);
40    normalize_tokens_per_second(attributes, duration);
41    normalize_ai_costs(attributes, costs);
42}
43
44/// Returns whether the item is should have AI normalizations applied.
45fn is_ai_item(attributes: &mut Attributes) -> bool {
46    // The product indicator whether we consider an item to be an EAP item.
47    if attributes.get_value(GEN_AI_OPERATION_TYPE).is_some() {
48        return true;
49    }
50
51    // We use the operation name to infer the operation type.
52    if attributes.get_value(GEN_AI_OPERATION_NAME).is_some() {
53        return true;
54    }
55
56    // Older SDKs may only send a (span) op which we also use to infer the operation type.
57    let op = attributes.get_value(OP).and_then(|op| op.as_str());
58    if op.is_some_and(|op| op.starts_with("gen_ai.") || op.starts_with("ai.")) {
59        return true;
60    }
61
62    false
63}
64
65/// Normalizes the [`GEN_AI_RESPONSE_MODEL`] attribute by defaulting to the [`GEN_AI_REQUEST_MODEL`] if it is missing.
66fn normalize_model(attributes: &mut Attributes) {
67    if attributes.contains_key(GEN_AI_RESPONSE_MODEL) {
68        return;
69    }
70    let Some(model) = attributes
71        .get_value(GEN_AI_REQUEST_MODEL)
72        .and_then(|v| v.as_str())
73    else {
74        return;
75    };
76    attributes.insert(GEN_AI_RESPONSE_MODEL, model.to_owned());
77}
78
79/// Normalizes the [`GEN_AI_OPERATION_TYPE`] and infers it from the AI operation if it is missing.
80fn normalize_ai_type(attributes: &mut Attributes) {
81    let op_name = attributes
82        .get_value(GEN_AI_OPERATION_NAME)
83        .or_else(|| attributes.get_value(OP))
84        .and_then(|op| op.as_str())
85        .and_then(|op| ai::infer_ai_operation_type(op))
86        // This is fine, this normalization only happens for known AI spans.
87        .unwrap_or(ai::DEFAULT_AI_OPERATION);
88
89    attributes.insert(GEN_AI_OPERATION_TYPE, op_name.to_owned());
90}
91
92/// Calculates the [`GEN_AI_USAGE_TOTAL_TOKENS`] attribute.
93fn normalize_total_tokens(attributes: &mut Attributes) {
94    let input_tokens = attributes
95        .get_value(GEN_AI_USAGE_INPUT_TOKENS)
96        .and_then(|v| v.as_f64());
97
98    let output_tokens = attributes
99        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
100        .and_then(|v| v.as_f64());
101
102    if input_tokens.is_none() && output_tokens.is_none() {
103        return;
104    }
105
106    let total_tokens = input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0);
107    attributes.insert(GEN_AI_USAGE_TOTAL_TOKENS, total_tokens);
108}
109
110/// Calculates the [`GEN_AI_RESPONSE_TPS`] attribute.
111fn normalize_tokens_per_second(attributes: &mut Attributes, duration: Option<Duration>) {
112    let Some(duration) = duration.filter(|d| !d.is_zero()) else {
113        return;
114    };
115
116    let output_tokens = attributes
117        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
118        .and_then(|v| v.as_f64())
119        .filter(|v| *v > 0.0);
120
121    if let Some(output_tokens) = output_tokens {
122        let tps = output_tokens / duration.as_secs_f64();
123        attributes.insert(GEN_AI_RESPONSE_TPS, tps);
124    }
125}
126
127/// Calculates model costs and serializes them into attributes.
128fn normalize_ai_costs(attributes: &mut Attributes, model_costs: Option<&ModelCosts>) {
129    let origin = extract_string_value(attributes, ORIGIN);
130    let platform = extract_string_value(attributes, PLATFORM);
131
132    let integration = map_origin_to_integration(origin);
133    let platform_tag = platform_tag(platform);
134
135    let Some(model_id) = attributes
136        .get_value(GEN_AI_RESPONSE_MODEL)
137        .and_then(|v| v.as_str())
138    else {
139        relay_statsd::metric!(
140            counter(Counters::GenAiCostCalculationResult) += 1,
141            result = "calculation_no_model_id_available",
142            integration = integration,
143            platform = platform_tag,
144        );
145        return;
146    };
147
148    let Some(model_cost) = model_costs.and_then(|c| c.cost_per_token(model_id)) else {
149        relay_statsd::metric!(
150            counter(Counters::GenAiCostCalculationResult) += 1,
151            result = "calculation_no_model_cost_available",
152            integration = integration,
153            platform = platform_tag,
154        );
155        return;
156    };
157
158    let get_tokens = |key| {
159        attributes
160            .get_value(key)
161            .and_then(|v| v.as_f64())
162            .unwrap_or(0.0)
163    };
164
165    let tokens = ai::UsedTokens {
166        input_tokens: get_tokens(GEN_AI_USAGE_INPUT_TOKENS),
167        input_cached_tokens: get_tokens(GEN_AI_USAGE_INPUT_CACHED_TOKENS),
168        input_cache_write_tokens: get_tokens(GEN_AI_USAGE_INPUT_CACHE_WRITE_TOKENS),
169        output_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_TOKENS),
170        output_reasoning_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_REASONING_TOKENS),
171    };
172
173    let Some(costs) = ai::calculate_costs(model_cost, tokens, integration, platform_tag) else {
174        return;
175    };
176
177    // Overwrite all values, the attributes should reflect the values we used to calculate the total.
178    attributes.insert(GEN_AI_COST_INPUT_TOKENS, costs.input);
179    attributes.insert(GEN_AI_COST_OUTPUT_TOKENS, costs.output);
180    attributes.insert(GEN_AI_COST_TOTAL_TOKENS, costs.total());
181}
182
183fn extract_string_value<'a>(attributes: &'a Attributes, key: &str) -> Option<&'a str> {
184    attributes.get_value(key).and_then(|v| v.as_str())
185}
186
187#[cfg(test)]
188mod tests {
189    use std::collections::HashMap;
190
191    use relay_pattern::Pattern;
192    use relay_protocol::{Empty, assert_annotated_snapshot};
193
194    use crate::ModelCostV2;
195
196    use super::*;
197
198    macro_rules! attributes {
199        ($($key:expr => $value:expr),* $(,)?) => {
200            Attributes::from([
201                $(($key.into(), Annotated::new($value.into())),)*
202            ])
203        };
204    }
205
206    fn model_costs() -> ModelCosts {
207        ModelCosts {
208            version: 2,
209            models: HashMap::from([
210                (
211                    Pattern::new("claude-2.1").unwrap(),
212                    ModelCostV2 {
213                        input_per_token: 0.01,
214                        output_per_token: 0.02,
215                        output_reasoning_per_token: 0.03,
216                        input_cached_per_token: 0.04,
217                        input_cache_write_per_token: 0.0,
218                    },
219                ),
220                (
221                    Pattern::new("gpt4-21-04").unwrap(),
222                    ModelCostV2 {
223                        input_per_token: 0.09,
224                        output_per_token: 0.05,
225                        output_reasoning_per_token: 0.0,
226                        input_cached_per_token: 0.0,
227                        input_cache_write_per_token: 0.0,
228                    },
229                ),
230            ]),
231        }
232    }
233
234    #[test]
235    fn test_normalize_ai_all_tokens() {
236        let mut attributes = Annotated::new(attributes! {
237            "gen_ai.operation.type" => "ai_client".to_owned(),
238            "gen_ai.usage.input_tokens" => 1000,
239            "gen_ai.usage.output_tokens" => 2000,
240            "gen_ai.usage.output_tokens.reasoning" => 1000,
241            "gen_ai.usage.input_tokens.cached" => 500,
242            "gen_ai.request.model" => "claude-2.1".to_owned(),
243        });
244
245        normalize_ai(
246            &mut attributes,
247            Some(Duration::from_secs(1)),
248            Some(&model_costs()),
249        );
250
251        assert_annotated_snapshot!(attributes, @r#"
252        {
253          "gen_ai.cost.input_tokens": {
254            "type": "double",
255            "value": 25.0
256          },
257          "gen_ai.cost.output_tokens": {
258            "type": "double",
259            "value": 50.0
260          },
261          "gen_ai.cost.total_tokens": {
262            "type": "double",
263            "value": 75.0
264          },
265          "gen_ai.operation.type": {
266            "type": "string",
267            "value": "ai_client"
268          },
269          "gen_ai.request.model": {
270            "type": "string",
271            "value": "claude-2.1"
272          },
273          "gen_ai.response.model": {
274            "type": "string",
275            "value": "claude-2.1"
276          },
277          "gen_ai.response.tokens_per_second": {
278            "type": "double",
279            "value": 2000.0
280          },
281          "gen_ai.usage.input_tokens": {
282            "type": "integer",
283            "value": 1000
284          },
285          "gen_ai.usage.input_tokens.cached": {
286            "type": "integer",
287            "value": 500
288          },
289          "gen_ai.usage.output_tokens": {
290            "type": "integer",
291            "value": 2000
292          },
293          "gen_ai.usage.output_tokens.reasoning": {
294            "type": "integer",
295            "value": 1000
296          },
297          "gen_ai.usage.total_tokens": {
298            "type": "double",
299            "value": 3000.0
300          }
301        }
302        "#);
303    }
304
305    #[test]
306    fn test_normalize_ai_basic_tokens() {
307        let mut attributes = Annotated::new(attributes! {
308            "gen_ai.operation.type" => "ai_client".to_owned(),
309            "gen_ai.usage.input_tokens" => 1000,
310            "gen_ai.usage.output_tokens" => 2000,
311            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
312        });
313
314        normalize_ai(
315            &mut attributes,
316            Some(Duration::from_millis(500)),
317            Some(&model_costs()),
318        );
319
320        assert_annotated_snapshot!(attributes, @r#"
321        {
322          "gen_ai.cost.input_tokens": {
323            "type": "double",
324            "value": 90.0
325          },
326          "gen_ai.cost.output_tokens": {
327            "type": "double",
328            "value": 100.0
329          },
330          "gen_ai.cost.total_tokens": {
331            "type": "double",
332            "value": 190.0
333          },
334          "gen_ai.operation.type": {
335            "type": "string",
336            "value": "ai_client"
337          },
338          "gen_ai.request.model": {
339            "type": "string",
340            "value": "gpt4-21-04"
341          },
342          "gen_ai.response.model": {
343            "type": "string",
344            "value": "gpt4-21-04"
345          },
346          "gen_ai.response.tokens_per_second": {
347            "type": "double",
348            "value": 4000.0
349          },
350          "gen_ai.usage.input_tokens": {
351            "type": "integer",
352            "value": 1000
353          },
354          "gen_ai.usage.output_tokens": {
355            "type": "integer",
356            "value": 2000
357          },
358          "gen_ai.usage.total_tokens": {
359            "type": "double",
360            "value": 3000.0
361          }
362        }
363        "#);
364    }
365
366    #[test]
367    fn test_normalize_ai_basic_tokens_no_duration_no_cost() {
368        let mut attributes = Annotated::new(attributes! {
369            "gen_ai.operation.type" => "ai_client".to_owned(),
370            "gen_ai.usage.input_tokens" => 1000,
371            "gen_ai.usage.output_tokens" => 2000,
372            "gen_ai.request.model" => "unknown".to_owned(),
373        });
374
375        normalize_ai(&mut attributes, Some(Duration::ZERO), Some(&model_costs()));
376
377        assert_annotated_snapshot!(attributes, @r#"
378        {
379          "gen_ai.operation.type": {
380            "type": "string",
381            "value": "ai_client"
382          },
383          "gen_ai.request.model": {
384            "type": "string",
385            "value": "unknown"
386          },
387          "gen_ai.response.model": {
388            "type": "string",
389            "value": "unknown"
390          },
391          "gen_ai.usage.input_tokens": {
392            "type": "integer",
393            "value": 1000
394          },
395          "gen_ai.usage.output_tokens": {
396            "type": "integer",
397            "value": 2000
398          },
399          "gen_ai.usage.total_tokens": {
400            "type": "double",
401            "value": 3000.0
402          }
403        }
404        "#);
405    }
406
407    #[test]
408    fn test_normalize_ai_does_not_overwrite() {
409        let mut attributes = Annotated::new(attributes! {
410            "gen_ai.operation.type" => "ai_client".to_owned(),
411            "gen_ai.usage.input_tokens" => 1000,
412            "gen_ai.usage.output_tokens" => 2000,
413            "gen_ai.request.model" => "gpt4".to_owned(),
414            "gen_ai.response.model" => "gpt4-21-04".to_owned(),
415
416            "gen_ai.cost.input_tokens" => 999.0,
417        });
418
419        normalize_ai(
420            &mut attributes,
421            Some(Duration::from_millis(500)),
422            Some(&model_costs()),
423        );
424
425        assert_annotated_snapshot!(attributes, @r#"
426        {
427          "gen_ai.cost.input_tokens": {
428            "type": "double",
429            "value": 90.0
430          },
431          "gen_ai.cost.output_tokens": {
432            "type": "double",
433            "value": 100.0
434          },
435          "gen_ai.cost.total_tokens": {
436            "type": "double",
437            "value": 190.0
438          },
439          "gen_ai.operation.type": {
440            "type": "string",
441            "value": "ai_client"
442          },
443          "gen_ai.request.model": {
444            "type": "string",
445            "value": "gpt4"
446          },
447          "gen_ai.response.model": {
448            "type": "string",
449            "value": "gpt4-21-04"
450          },
451          "gen_ai.response.tokens_per_second": {
452            "type": "double",
453            "value": 4000.0
454          },
455          "gen_ai.usage.input_tokens": {
456            "type": "integer",
457            "value": 1000
458          },
459          "gen_ai.usage.output_tokens": {
460            "type": "integer",
461            "value": 2000
462          },
463          "gen_ai.usage.total_tokens": {
464            "type": "double",
465            "value": 3000.0
466          }
467        }
468        "#);
469    }
470
471    #[test]
472    fn test_normalize_ai_overwrite_costs() {
473        let mut attributes = Annotated::new(attributes! {
474            "gen_ai.operation.type" => "ai_client".to_owned(),
475            "gen_ai.usage.input_tokens" => 1000,
476            "gen_ai.usage.output_tokens" => 2000,
477            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
478
479            "gen_ai.usage.total_tokens" => 1337,
480
481            "gen_ai.cost.input_tokens" => 99.0,
482            "gen_ai.cost.output_tokens" => 99.0,
483            "gen_ai.cost.total_tokens" => 123.0,
484
485            "gen_ai.response.tokens_per_second" => 42.0,
486        });
487
488        normalize_ai(
489            &mut attributes,
490            Some(Duration::from_millis(500)),
491            Some(&model_costs()),
492        );
493
494        assert_annotated_snapshot!(attributes, @r#"
495        {
496          "gen_ai.cost.input_tokens": {
497            "type": "double",
498            "value": 90.0
499          },
500          "gen_ai.cost.output_tokens": {
501            "type": "double",
502            "value": 100.0
503          },
504          "gen_ai.cost.total_tokens": {
505            "type": "double",
506            "value": 190.0
507          },
508          "gen_ai.operation.type": {
509            "type": "string",
510            "value": "ai_client"
511          },
512          "gen_ai.request.model": {
513            "type": "string",
514            "value": "gpt4-21-04"
515          },
516          "gen_ai.response.model": {
517            "type": "string",
518            "value": "gpt4-21-04"
519          },
520          "gen_ai.response.tokens_per_second": {
521            "type": "double",
522            "value": 4000.0
523          },
524          "gen_ai.usage.input_tokens": {
525            "type": "integer",
526            "value": 1000
527          },
528          "gen_ai.usage.output_tokens": {
529            "type": "integer",
530            "value": 2000
531          },
532          "gen_ai.usage.total_tokens": {
533            "type": "double",
534            "value": 3000.0
535          }
536        }
537        "#);
538    }
539
540    #[test]
541    fn test_normalize_ai_no_ai_attributes() {
542        let mut attributes = Annotated::new(attributes! {
543            "gen_ai.usage.input_tokens" => 1000,
544            "gen_ai.usage.output_tokens" => 2000,
545        });
546
547        normalize_ai(
548            &mut attributes,
549            Some(Duration::from_millis(500)),
550            Some(&model_costs()),
551        );
552
553        assert_annotated_snapshot!(&mut attributes, @r#"
554        {
555          "gen_ai.usage.input_tokens": {
556            "type": "integer",
557            "value": 1000
558          },
559          "gen_ai.usage.output_tokens": {
560            "type": "integer",
561            "value": 2000
562          }
563        }
564        "#);
565    }
566
567    #[test]
568    fn test_normalize_ai_no_ai_indicator_attribute() {
569        let mut attributes = Annotated::new(attributes! {
570            "foo" => 123,
571        });
572
573        normalize_ai(
574            &mut attributes,
575            Some(Duration::from_millis(500)),
576            Some(&model_costs()),
577        );
578
579        assert_annotated_snapshot!(&mut attributes, @r#"
580        {
581          "foo": {
582            "type": "integer",
583            "value": 123
584          }
585        }
586        "#);
587    }
588
589    #[test]
590    fn test_normalize_ai_empty() {
591        let mut attributes = Annotated::empty();
592
593        normalize_ai(
594            &mut attributes,
595            Some(Duration::from_millis(500)),
596            Some(&model_costs()),
597        );
598
599        assert!(attributes.is_empty());
600    }
601}
relay_event_normalization/eap/ai.rs

relay_event_normalization/eap/
ai.rs
