relay_event_normalization/eap/
ai.rs

1use std::time::Duration;
2
3use relay_conventions::attributes::*;
4use relay_event_schema::protocol::{Attribute, AttributeType, Attributes};
5use relay_protocol::{Annotated, Value};
6
7use crate::ModelMetadata;
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, except for existing model costs, in order
14/// to guarantee a consistent data set for the AI product module.
15///
16/// Callers may choose to only run this normalization in processing mode to not have the
17/// normalization run multiple times.
18pub fn normalize_ai(
19    attributes: &mut Annotated<Attributes>,
20    duration: Option<Duration>,
21    model_metadata: Option<&ModelMetadata>,
22) {
23    let Some(attributes) = attributes.value_mut() else {
24        return;
25    };
26
27    // Specifically only apply normalizations if the item is recognized as an AI item by the
28    // product.
29    if !is_ai_item(attributes) {
30        return;
31    }
32
33    normalize_model(attributes);
34    normalize_ai_type(attributes);
35    normalize_total_tokens(attributes);
36    normalize_tokens_per_second(attributes, duration);
37    normalize_context_utilization(attributes, model_metadata);
38    normalize_ai_costs(attributes, model_metadata);
39    normalize_tool_call_result_size(attributes);
40}
41
42/// Returns whether the item is should have AI normalizations applied.
43fn is_ai_item(attributes: &mut Attributes) -> bool {
44    // The product indicator whether we consider an item to be an EAP item.
45    if attributes.get_value(GEN_AI__OPERATION__TYPE).is_some() {
46        return true;
47    }
48
49    // We use the operation name to infer the operation type.
50    if attributes.get_value(GEN_AI__OPERATION__NAME).is_some() {
51        return true;
52    }
53
54    // Older SDKs may only send a (span) op which we also use to infer the operation type.
55    let op = attributes.get_value(SENTRY__OP).and_then(|op| op.as_str());
56    op.is_some_and(|op| op.starts_with("gen_ai.") || op.starts_with("ai."))
57}
58
59/// Normalizes the [`GEN_AI__RESPONSE__MODEL`] attribute by defaulting to the [`GEN_AI__REQUEST__MODEL`] if it is missing.
60fn normalize_model(attributes: &mut Attributes) {
61    if attributes.contains_key(GEN_AI__RESPONSE__MODEL) {
62        return;
63    }
64    let Some(model) = attributes
65        .get_value(GEN_AI__REQUEST__MODEL)
66        .and_then(|v| v.as_str())
67    else {
68        return;
69    };
70    attributes.insert(GEN_AI__RESPONSE__MODEL, model.to_owned());
71}
72
73/// Normalizes the [`GEN_AI__OPERATION__TYPE`] and infers it from the AI operation if it is missing.
74fn normalize_ai_type(attributes: &mut Attributes) {
75    let op_name = attributes
76        .get_value(GEN_AI__OPERATION__NAME)
77        .or_else(|| attributes.get_value(SENTRY__OP))
78        .and_then(|op| op.as_str())
79        .and_then(|op| ai::infer_ai_operation_type(op))
80        .unwrap_or_else(|| ai::default_ai_operation_type(attributes));
81
82    attributes.insert(GEN_AI__OPERATION__TYPE, op_name.to_owned());
83}
84
85/// Calculates the [`GEN_AI__USAGE__TOTAL_TOKENS`] attribute.
86fn normalize_total_tokens(attributes: &mut Attributes) {
87    let input_tokens = attributes
88        .get_value(GEN_AI__USAGE__INPUT_TOKENS)
89        .and_then(|v| v.as_f64());
90
91    let output_tokens = attributes
92        .get_value(GEN_AI__USAGE__OUTPUT_TOKENS)
93        .and_then(|v| v.as_f64());
94
95    if input_tokens.is_none() && output_tokens.is_none() {
96        return;
97    }
98
99    let total_tokens = input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0);
100    attributes.insert(GEN_AI__USAGE__TOTAL_TOKENS, total_tokens);
101}
102
103/// Calculates the [`GEN_AI__RESPONSE__TOKENS_PER_SECOND`] attribute.
104fn normalize_tokens_per_second(attributes: &mut Attributes, duration: Option<Duration>) {
105    let Some(duration) = duration.filter(|d| !d.is_zero()) else {
106        return;
107    };
108
109    let output_tokens = attributes
110        .get_value(GEN_AI__USAGE__OUTPUT_TOKENS)
111        .and_then(|v| v.as_f64())
112        .filter(|v| *v > 0.0);
113
114    if let Some(output_tokens) = output_tokens {
115        let tps = output_tokens / duration.as_secs_f64();
116        attributes.insert(GEN_AI__RESPONSE__TOKENS_PER_SECOND, tps);
117    }
118}
119
120/// Sets the context window size and utilization for the model.
121fn normalize_context_utilization(
122    attributes: &mut Attributes,
123    model_metadata: Option<&ModelMetadata>,
124) {
125    let model_id = attributes
126        .get_value(GEN_AI__RESPONSE__MODEL)
127        .and_then(|v| v.as_str());
128
129    let context_size = model_id.and_then(|id| model_metadata.and_then(|m| m.context_size(id)));
130
131    let Some(context_size) = context_size else {
132        return;
133    };
134
135    attributes.insert(GEN_AI__CONTEXT__WINDOW_SIZE, context_size as i64);
136
137    let total_tokens = attributes
138        .get_value(GEN_AI__USAGE__TOTAL_TOKENS)
139        .and_then(|v| v.as_f64());
140
141    if let Some(total_tokens) = total_tokens {
142        attributes.insert(
143            GEN_AI__CONTEXT__UTILIZATION,
144            total_tokens / context_size as f64,
145        );
146    }
147}
148
149/// Calculates model costs and serializes them into attributes.
150fn normalize_ai_costs(attributes: &mut Attributes, model_metadata: Option<&ModelMetadata>) {
151    // Preserve a valid attached total cost because Relay cannot calculate costs for self-hosted or
152    // otherwise unknown models.
153    if ai::has_valid_total_cost(attributes) {
154        return;
155    }
156
157    let origin = extract_string_value(attributes, SENTRY__ORIGIN);
158    let platform = extract_string_value(attributes, SENTRY__PLATFORM);
159
160    let integration = map_origin_to_integration(origin);
161    let platform_tag = platform_tag(platform);
162
163    let Some(model_id) = attributes
164        .get_value(GEN_AI__RESPONSE__MODEL)
165        .and_then(|v| v.as_str())
166    else {
167        relay_statsd::metric!(
168            counter(Counters::GenAiCostCalculationResult) += 1,
169            result = "calculation_no_model_id_available",
170            integration = integration,
171            platform = platform_tag,
172        );
173        return;
174    };
175
176    let Some(model_cost) = model_metadata.and_then(|m| m.cost_per_token(model_id)) else {
177        relay_statsd::metric!(
178            counter(Counters::GenAiCostCalculationResult) += 1,
179            result = "calculation_no_model_cost_available",
180            integration = integration,
181            platform = platform_tag,
182        );
183        return;
184    };
185
186    let get_tokens = |key| {
187        attributes
188            .get_value(key)
189            .and_then(|v| v.as_f64())
190            .unwrap_or(0.0)
191    };
192
193    let tokens = ai::UsedTokens {
194        input_tokens: get_tokens(GEN_AI__USAGE__INPUT_TOKENS),
195        input_cached_tokens: get_tokens(GEN_AI__USAGE__CACHE_READ__INPUT_TOKENS),
196        input_cache_write_tokens: get_tokens(GEN_AI__USAGE__CACHE_CREATION__INPUT_TOKENS),
197        output_tokens: get_tokens(GEN_AI__USAGE__OUTPUT_TOKENS),
198        output_reasoning_tokens: get_tokens(GEN_AI__USAGE__REASONING__OUTPUT_TOKENS),
199    };
200
201    let Some(costs) = ai::calculate_costs(model_cost, tokens, integration, platform_tag) else {
202        return;
203    };
204
205    // Overwrite all values, the attributes should reflect the values we used to calculate the total.
206    attributes.insert(GEN_AI__COST__INPUT_TOKENS, costs.input);
207    attributes.insert(
208        GEN_AI__COST__CACHE_READ__INPUT_TOKENS,
209        costs.cache_read_input,
210    );
211    attributes.insert(
212        GEN_AI__COST__CACHE_CREATION__INPUT_TOKENS,
213        costs.cache_creation_input,
214    );
215
216    attributes.insert(GEN_AI__COST__OUTPUT_TOKENS, costs.output);
217    attributes.insert(
218        GEN_AI__COST__REASONING__OUTPUT_TOKENS,
219        costs.reasoning_output,
220    );
221
222    attributes.insert(GEN_AI__COST__TOTAL_TOKENS, costs.total());
223}
224
225/// Records a tool result's UTF-8 byte length, preserving SDK-provided sizes.
226fn normalize_tool_call_result_size(attributes: &mut Attributes) {
227    if extract_string_value(attributes, GEN_AI__OPERATION__TYPE) != Some("tool") {
228        return;
229    }
230
231    let has_valid_size = matches!(
232        attributes.get_value(GEN_AI__TOOL__CALL__RESULT__SIZE),
233        Some(Value::U64(_)) | Some(Value::I64(0..))
234    );
235
236    if has_valid_size {
237        return;
238    }
239
240    let Some(result) = extract_string_value(attributes, GEN_AI__TOOL__CALL__RESULT) else {
241        return;
242    };
243
244    let size = Attribute::new(AttributeType::Integer, Value::U64(result.len() as u64));
245    attributes.0.insert(
246        GEN_AI__TOOL__CALL__RESULT__SIZE.to_owned(),
247        Annotated::new(size),
248    );
249}
250
251fn extract_string_value<'a>(attributes: &'a Attributes, key: &str) -> Option<&'a str> {
252    attributes.get_value(key).and_then(|v| v.as_str())
253}
254
255#[cfg(test)]
256mod tests {
257    use std::collections::HashMap;
258
259    use relay_protocol::{Empty, assert_annotated_snapshot};
260
261    use crate::{ModelCostV2, ModelMetadataEntry};
262
263    use super::*;
264
265    macro_rules! attributes {
266        ($($key:expr => $value:expr),* $(,)?) => {
267            Attributes::from([
268                $(($key.into(), Annotated::new($value.into())),)*
269            ])
270        };
271    }
272
273    fn model_metadata() -> ModelMetadata {
274        ModelMetadata {
275            version: 1,
276            models: HashMap::from([
277                (
278                    "claude-2.1".parse().unwrap(),
279                    ModelMetadataEntry {
280                        costs: Some(ModelCostV2 {
281                            input_per_token: 0.01,
282                            output_per_token: 0.02,
283                            output_reasoning_per_token: 0.03,
284                            input_cached_per_token: 0.04,
285                            input_cache_write_per_token: 0.0,
286                        }),
287                        context_size: None,
288                    },
289                ),
290                (
291                    "gpt4-21-04".parse().unwrap(),
292                    ModelMetadataEntry {
293                        costs: Some(ModelCostV2 {
294                            input_per_token: 0.09,
295                            output_per_token: 0.05,
296                            output_reasoning_per_token: 0.0,
297                            input_cached_per_token: 0.0,
298                            input_cache_write_per_token: 0.0,
299                        }),
300                        context_size: None,
301                    },
302                ),
303            ]),
304        }
305    }
306
307    fn model_metadata_with_context_size() -> ModelMetadata {
308        ModelMetadata {
309            version: 1,
310            models: HashMap::from([(
311                "claude-2.1".parse().unwrap(),
312                ModelMetadataEntry {
313                    costs: Some(ModelCostV2 {
314                        input_per_token: 0.01,
315                        output_per_token: 0.02,
316                        output_reasoning_per_token: 0.03,
317                        input_cached_per_token: 0.04,
318                        input_cache_write_per_token: 0.0,
319                    }),
320                    context_size: Some(100_000),
321                },
322            )]),
323        }
324    }
325
326    #[test]
327    fn test_normalize_ai_type_fallback() {
328        for model_key in [
329            None,
330            Some(GEN_AI__REQUEST__MODEL),
331            Some(GEN_AI__RESPONSE__MODEL),
332        ] {
333            let mut attributes = attributes! {
334                "gen_ai.operation.name" => "embeddings".to_owned(),
335            };
336            if let Some(key) = model_key {
337                attributes.insert(key, "gpt-4".to_owned());
338            }
339            let mut attributes = Annotated::new(attributes);
340
341            normalize_ai(&mut attributes, None, None);
342
343            assert_eq!(
344                attributes
345                    .value()
346                    .unwrap()
347                    .get_value(GEN_AI__OPERATION__TYPE)
348                    .and_then(|v| v.as_str()),
349                Some(if model_key.is_some() {
350                    "ai_client"
351                } else {
352                    "other"
353                }),
354            );
355        }
356    }
357
358    #[test]
359    fn test_normalize_ai_all_tokens() {
360        let mut attributes = Annotated::new(attributes! {
361            "gen_ai.operation.type" => "ai_client".to_owned(),
362            "gen_ai.usage.input_tokens" => 1000,
363            "gen_ai.usage.output_tokens" => 2000,
364            "gen_ai.usage.reasoning.output_tokens" => 1000,
365            "gen_ai.usage.cache_read.input_tokens" => 500,
366            "gen_ai.request.model" => "claude-2.1".to_owned(),
367        });
368
369        normalize_ai(
370            &mut attributes,
371            Some(Duration::from_secs(1)),
372            Some(&model_metadata()),
373        );
374
375        assert_annotated_snapshot!(attributes, @r#"
376        {
377          "gen_ai.cost.cache_creation.input_tokens": {
378            "type": "double",
379            "value": 0.0
380          },
381          "gen_ai.cost.cache_read.input_tokens": {
382            "type": "double",
383            "value": 20.0
384          },
385          "gen_ai.cost.input_tokens": {
386            "type": "double",
387            "value": 25.0
388          },
389          "gen_ai.cost.output_tokens": {
390            "type": "double",
391            "value": 50.0
392          },
393          "gen_ai.cost.reasoning.output_tokens": {
394            "type": "double",
395            "value": 30.0
396          },
397          "gen_ai.cost.total_tokens": {
398            "type": "double",
399            "value": 75.0
400          },
401          "gen_ai.operation.type": {
402            "type": "string",
403            "value": "ai_client"
404          },
405          "gen_ai.request.model": {
406            "type": "string",
407            "value": "claude-2.1"
408          },
409          "gen_ai.response.model": {
410            "type": "string",
411            "value": "claude-2.1"
412          },
413          "gen_ai.response.tokens_per_second": {
414            "type": "double",
415            "value": 2000.0
416          },
417          "gen_ai.usage.cache_read.input_tokens": {
418            "type": "integer",
419            "value": 500
420          },
421          "gen_ai.usage.input_tokens": {
422            "type": "integer",
423            "value": 1000
424          },
425          "gen_ai.usage.output_tokens": {
426            "type": "integer",
427            "value": 2000
428          },
429          "gen_ai.usage.reasoning.output_tokens": {
430            "type": "integer",
431            "value": 1000
432          },
433          "gen_ai.usage.total_tokens": {
434            "type": "double",
435            "value": 3000.0
436          }
437        }
438        "#);
439    }
440
441    #[test]
442    fn test_normalize_ai_basic_tokens() {
443        let mut attributes = Annotated::new(attributes! {
444            "gen_ai.operation.type" => "ai_client".to_owned(),
445            "gen_ai.usage.input_tokens" => 1000,
446            "gen_ai.usage.output_tokens" => 2000,
447            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
448        });
449
450        normalize_ai(
451            &mut attributes,
452            Some(Duration::from_millis(500)),
453            Some(&model_metadata()),
454        );
455
456        assert_annotated_snapshot!(attributes, @r#"
457        {
458          "gen_ai.cost.cache_creation.input_tokens": {
459            "type": "double",
460            "value": 0.0
461          },
462          "gen_ai.cost.cache_read.input_tokens": {
463            "type": "double",
464            "value": 0.0
465          },
466          "gen_ai.cost.input_tokens": {
467            "type": "double",
468            "value": 90.0
469          },
470          "gen_ai.cost.output_tokens": {
471            "type": "double",
472            "value": 100.0
473          },
474          "gen_ai.cost.reasoning.output_tokens": {
475            "type": "double",
476            "value": 0.0
477          },
478          "gen_ai.cost.total_tokens": {
479            "type": "double",
480            "value": 190.0
481          },
482          "gen_ai.operation.type": {
483            "type": "string",
484            "value": "ai_client"
485          },
486          "gen_ai.request.model": {
487            "type": "string",
488            "value": "gpt4-21-04"
489          },
490          "gen_ai.response.model": {
491            "type": "string",
492            "value": "gpt4-21-04"
493          },
494          "gen_ai.response.tokens_per_second": {
495            "type": "double",
496            "value": 4000.0
497          },
498          "gen_ai.usage.input_tokens": {
499            "type": "integer",
500            "value": 1000
501          },
502          "gen_ai.usage.output_tokens": {
503            "type": "integer",
504            "value": 2000
505          },
506          "gen_ai.usage.total_tokens": {
507            "type": "double",
508            "value": 3000.0
509          }
510        }
511        "#);
512    }
513
514    #[test]
515    fn test_normalize_ai_basic_tokens_no_duration_no_cost() {
516        let mut attributes = Annotated::new(attributes! {
517            "gen_ai.operation.type" => "ai_client".to_owned(),
518            "gen_ai.usage.input_tokens" => 1000,
519            "gen_ai.usage.output_tokens" => 2000,
520            "gen_ai.request.model" => "unknown".to_owned(),
521        });
522
523        normalize_ai(
524            &mut attributes,
525            Some(Duration::ZERO),
526            Some(&model_metadata()),
527        );
528
529        assert_annotated_snapshot!(attributes, @r#"
530        {
531          "gen_ai.operation.type": {
532            "type": "string",
533            "value": "ai_client"
534          },
535          "gen_ai.request.model": {
536            "type": "string",
537            "value": "unknown"
538          },
539          "gen_ai.response.model": {
540            "type": "string",
541            "value": "unknown"
542          },
543          "gen_ai.usage.input_tokens": {
544            "type": "integer",
545            "value": 1000
546          },
547          "gen_ai.usage.output_tokens": {
548            "type": "integer",
549            "value": 2000
550          },
551          "gen_ai.usage.total_tokens": {
552            "type": "double",
553            "value": 3000.0
554          }
555        }
556        "#);
557    }
558
559    #[test]
560    fn test_normalize_ai_does_not_overwrite() {
561        let mut attributes = Annotated::new(attributes! {
562            "gen_ai.operation.type" => "ai_client".to_owned(),
563            "gen_ai.usage.input_tokens" => 1000,
564            "gen_ai.usage.output_tokens" => 2000,
565            "gen_ai.request.model" => "gpt4".to_owned(),
566            "gen_ai.response.model" => "gpt4-21-04".to_owned(),
567
568            "gen_ai.cost.input_tokens" => 999.0,
569        });
570
571        normalize_ai(
572            &mut attributes,
573            Some(Duration::from_millis(500)),
574            Some(&model_metadata()),
575        );
576
577        assert_annotated_snapshot!(attributes, @r#"
578        {
579          "gen_ai.cost.cache_creation.input_tokens": {
580            "type": "double",
581            "value": 0.0
582          },
583          "gen_ai.cost.cache_read.input_tokens": {
584            "type": "double",
585            "value": 0.0
586          },
587          "gen_ai.cost.input_tokens": {
588            "type": "double",
589            "value": 90.0
590          },
591          "gen_ai.cost.output_tokens": {
592            "type": "double",
593            "value": 100.0
594          },
595          "gen_ai.cost.reasoning.output_tokens": {
596            "type": "double",
597            "value": 0.0
598          },
599          "gen_ai.cost.total_tokens": {
600            "type": "double",
601            "value": 190.0
602          },
603          "gen_ai.operation.type": {
604            "type": "string",
605            "value": "ai_client"
606          },
607          "gen_ai.request.model": {
608            "type": "string",
609            "value": "gpt4"
610          },
611          "gen_ai.response.model": {
612            "type": "string",
613            "value": "gpt4-21-04"
614          },
615          "gen_ai.response.tokens_per_second": {
616            "type": "double",
617            "value": 4000.0
618          },
619          "gen_ai.usage.input_tokens": {
620            "type": "integer",
621            "value": 1000
622          },
623          "gen_ai.usage.output_tokens": {
624            "type": "integer",
625            "value": 2000
626          },
627          "gen_ai.usage.total_tokens": {
628            "type": "double",
629            "value": 3000.0
630          }
631        }
632        "#);
633    }
634
635    #[test]
636    fn test_normalize_ai_preserves_costs() {
637        let mut attributes = Annotated::new(attributes! {
638            "gen_ai.operation.type" => "ai_client".to_owned(),
639            "gen_ai.usage.input_tokens" => 1000,
640            "gen_ai.usage.output_tokens" => 2000,
641            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
642
643            "gen_ai.usage.total_tokens" => 1337,
644
645            "gen_ai.cost.input_tokens" => 99.0,
646            "gen_ai.cost.output_tokens" => 99.0,
647            "gen_ai.cost.total_tokens" => 123.0,
648
649            "gen_ai.response.tokens_per_second" => 42.0,
650        });
651
652        normalize_ai(
653            &mut attributes,
654            Some(Duration::from_millis(500)),
655            Some(&model_metadata()),
656        );
657
658        assert_annotated_snapshot!(attributes, @r#"
659        {
660          "gen_ai.cost.input_tokens": {
661            "type": "double",
662            "value": 99.0
663          },
664          "gen_ai.cost.output_tokens": {
665            "type": "double",
666            "value": 99.0
667          },
668          "gen_ai.cost.total_tokens": {
669            "type": "double",
670            "value": 123.0
671          },
672          "gen_ai.operation.type": {
673            "type": "string",
674            "value": "ai_client"
675          },
676          "gen_ai.request.model": {
677            "type": "string",
678            "value": "gpt4-21-04"
679          },
680          "gen_ai.response.model": {
681            "type": "string",
682            "value": "gpt4-21-04"
683          },
684          "gen_ai.response.tokens_per_second": {
685            "type": "double",
686            "value": 4000.0
687          },
688          "gen_ai.usage.input_tokens": {
689            "type": "integer",
690            "value": 1000
691          },
692          "gen_ai.usage.output_tokens": {
693            "type": "integer",
694            "value": 2000
695          },
696          "gen_ai.usage.total_tokens": {
697            "type": "double",
698            "value": 3000.0
699          }
700        }
701        "#);
702    }
703
704    #[test]
705    fn test_normalize_ai_no_ai_attributes() {
706        let mut attributes = Annotated::new(attributes! {
707            "gen_ai.usage.input_tokens" => 1000,
708            "gen_ai.usage.output_tokens" => 2000,
709        });
710
711        normalize_ai(
712            &mut attributes,
713            Some(Duration::from_millis(500)),
714            Some(&model_metadata()),
715        );
716
717        assert_annotated_snapshot!(&mut attributes, @r#"
718        {
719          "gen_ai.usage.input_tokens": {
720            "type": "integer",
721            "value": 1000
722          },
723          "gen_ai.usage.output_tokens": {
724            "type": "integer",
725            "value": 2000
726          }
727        }
728        "#);
729    }
730
731    #[test]
732    fn test_normalize_ai_no_ai_indicator_attribute() {
733        let mut attributes = Annotated::new(attributes! {
734            "foo" => 123,
735        });
736
737        normalize_ai(
738            &mut attributes,
739            Some(Duration::from_millis(500)),
740            Some(&model_metadata()),
741        );
742
743        assert_annotated_snapshot!(&mut attributes, @r#"
744        {
745          "foo": {
746            "type": "integer",
747            "value": 123
748          }
749        }
750        "#);
751    }
752
753    #[test]
754    fn test_normalize_ai_empty() {
755        let mut attributes = Annotated::empty();
756
757        normalize_ai(
758            &mut attributes,
759            Some(Duration::from_millis(500)),
760            Some(&model_metadata()),
761        );
762
763        assert!(attributes.is_empty());
764    }
765
766    #[test]
767    fn test_context_utilization_with_total_tokens() {
768        let mut attributes = Annotated::new(attributes! {
769            "gen_ai.operation.type" => "ai_client".to_owned(),
770            "gen_ai.usage.input_tokens" => 30000,
771            "gen_ai.usage.output_tokens" => 12000,
772            "gen_ai.request.model" => "claude-2.1".to_owned(),
773        });
774
775        normalize_ai(
776            &mut attributes,
777            Some(Duration::from_secs(1)),
778            Some(&model_metadata_with_context_size()),
779        );
780
781        assert_annotated_snapshot!(attributes, @r#"
782        {
783          "gen_ai.context.utilization": {
784            "type": "double",
785            "value": 0.42
786          },
787          "gen_ai.context.window_size": {
788            "type": "integer",
789            "value": 100000
790          },
791          "gen_ai.cost.cache_creation.input_tokens": {
792            "type": "double",
793            "value": 0.0
794          },
795          "gen_ai.cost.cache_read.input_tokens": {
796            "type": "double",
797            "value": 0.0
798          },
799          "gen_ai.cost.input_tokens": {
800            "type": "double",
801            "value": 300.0
802          },
803          "gen_ai.cost.output_tokens": {
804            "type": "double",
805            "value": 240.0
806          },
807          "gen_ai.cost.reasoning.output_tokens": {
808            "type": "double",
809            "value": 0.0
810          },
811          "gen_ai.cost.total_tokens": {
812            "type": "double",
813            "value": 540.0
814          },
815          "gen_ai.operation.type": {
816            "type": "string",
817            "value": "ai_client"
818          },
819          "gen_ai.request.model": {
820            "type": "string",
821            "value": "claude-2.1"
822          },
823          "gen_ai.response.model": {
824            "type": "string",
825            "value": "claude-2.1"
826          },
827          "gen_ai.response.tokens_per_second": {
828            "type": "double",
829            "value": 12000.0
830          },
831          "gen_ai.usage.input_tokens": {
832            "type": "integer",
833            "value": 30000
834          },
835          "gen_ai.usage.output_tokens": {
836            "type": "integer",
837            "value": 12000
838          },
839          "gen_ai.usage.total_tokens": {
840            "type": "double",
841            "value": 42000.0
842          }
843        }
844        "#);
845    }
846
847    #[test]
848    fn test_context_utilization_no_context_size() {
849        let mut attributes = Annotated::new(attributes! {
850            "gen_ai.operation.type" => "ai_client".to_owned(),
851            "gen_ai.usage.input_tokens" => 1000,
852            "gen_ai.usage.output_tokens" => 2000,
853            "gen_ai.request.model" => "claude-2.1".to_owned(),
854        });
855
856        // model_metadata() has no context_size set.
857        normalize_ai(
858            &mut attributes,
859            Some(Duration::from_secs(1)),
860            Some(&model_metadata()),
861        );
862
863        let attrs = attributes.value().unwrap();
864        assert!(attrs.get_value("gen_ai.context.window_size").is_none());
865        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
866    }
867
868    #[test]
869    fn test_context_utilization_no_total_tokens() {
870        // Only context_size is available, but no token counts at all.
871        let mut attributes = Annotated::new(attributes! {
872            "gen_ai.operation.type" => "ai_client".to_owned(),
873            "gen_ai.request.model" => "claude-2.1".to_owned(),
874        });
875
876        normalize_ai(
877            &mut attributes,
878            Some(Duration::from_secs(1)),
879            Some(&model_metadata_with_context_size()),
880        );
881
882        let attrs = attributes.value().unwrap();
883        // window_size should still be set even without tokens.
884        assert_eq!(
885            attrs
886                .get_value("gen_ai.context.window_size")
887                .unwrap()
888                .as_f64(),
889            Some(100_000.0)
890        );
891        // But utilization cannot be computed without total_tokens.
892        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
893    }
894
895    #[test]
896    fn test_context_utilization_unknown_model() {
897        let mut attributes = Annotated::new(attributes! {
898            "gen_ai.operation.type" => "ai_client".to_owned(),
899            "gen_ai.usage.input_tokens" => 1000,
900            "gen_ai.usage.output_tokens" => 2000,
901            "gen_ai.request.model" => "unknown-model".to_owned(),
902        });
903
904        normalize_ai(
905            &mut attributes,
906            Some(Duration::from_secs(1)),
907            Some(&model_metadata_with_context_size()),
908        );
909
910        let attrs = attributes.value().unwrap();
911        assert!(attrs.get_value("gen_ai.context.window_size").is_none());
912        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
913    }
914}
relay_event_normalization/eap/ai.rs

relay_event_normalization/eap/
ai.rs
