relay_event_normalization/eap/
ai.rs

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

relay_event_normalization/eap/
ai.rs
