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

relay_event_normalization/eap/
ai.rs
