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_protocol::{Empty, assert_annotated_snapshot};
238
239    use crate::{ModelCostV2, ModelMetadataEntry};
240
241    use super::*;
242
243    macro_rules! attributes {
244        ($($key:expr => $value:expr),* $(,)?) => {
245            Attributes::from([
246                $(($key.into(), Annotated::new($value.into())),)*
247            ])
248        };
249    }
250
251    fn model_metadata() -> ModelMetadata {
252        ModelMetadata {
253            version: 1,
254            models: HashMap::from([
255                (
256                    "claude-2.1".parse().unwrap(),
257                    ModelMetadataEntry {
258                        costs: Some(ModelCostV2 {
259                            input_per_token: 0.01,
260                            output_per_token: 0.02,
261                            output_reasoning_per_token: 0.03,
262                            input_cached_per_token: 0.04,
263                            input_cache_write_per_token: 0.0,
264                        }),
265                        context_size: None,
266                    },
267                ),
268                (
269                    "gpt4-21-04".parse().unwrap(),
270                    ModelMetadataEntry {
271                        costs: Some(ModelCostV2 {
272                            input_per_token: 0.09,
273                            output_per_token: 0.05,
274                            output_reasoning_per_token: 0.0,
275                            input_cached_per_token: 0.0,
276                            input_cache_write_per_token: 0.0,
277                        }),
278                        context_size: None,
279                    },
280                ),
281            ]),
282        }
283    }
284
285    fn model_metadata_with_context_size() -> ModelMetadata {
286        ModelMetadata {
287            version: 1,
288            models: HashMap::from([(
289                "claude-2.1".parse().unwrap(),
290                ModelMetadataEntry {
291                    costs: Some(ModelCostV2 {
292                        input_per_token: 0.01,
293                        output_per_token: 0.02,
294                        output_reasoning_per_token: 0.03,
295                        input_cached_per_token: 0.04,
296                        input_cache_write_per_token: 0.0,
297                    }),
298                    context_size: Some(100_000),
299                },
300            )]),
301        }
302    }
303
304    #[test]
305    fn test_normalize_ai_all_tokens() {
306        let mut attributes = Annotated::new(attributes! {
307            "gen_ai.operation.type" => "ai_client".to_owned(),
308            "gen_ai.usage.input_tokens" => 1000,
309            "gen_ai.usage.output_tokens" => 2000,
310            "gen_ai.usage.reasoning.output_tokens" => 1000,
311            "gen_ai.usage.cache_read.input_tokens" => 500,
312            "gen_ai.request.model" => "claude-2.1".to_owned(),
313        });
314
315        normalize_ai(
316            &mut attributes,
317            Some(Duration::from_secs(1)),
318            Some(&model_metadata()),
319        );
320
321        assert_annotated_snapshot!(attributes, @r#"
322        {
323          "gen_ai.cost.cache_creation.input_tokens": {
324            "type": "double",
325            "value": 0.0
326          },
327          "gen_ai.cost.cache_read.input_tokens": {
328            "type": "double",
329            "value": 20.0
330          },
331          "gen_ai.cost.input_tokens": {
332            "type": "double",
333            "value": 25.0
334          },
335          "gen_ai.cost.output_tokens": {
336            "type": "double",
337            "value": 50.0
338          },
339          "gen_ai.cost.reasoning.output_tokens": {
340            "type": "double",
341            "value": 30.0
342          },
343          "gen_ai.cost.total_tokens": {
344            "type": "double",
345            "value": 75.0
346          },
347          "gen_ai.operation.type": {
348            "type": "string",
349            "value": "ai_client"
350          },
351          "gen_ai.request.model": {
352            "type": "string",
353            "value": "claude-2.1"
354          },
355          "gen_ai.response.model": {
356            "type": "string",
357            "value": "claude-2.1"
358          },
359          "gen_ai.response.tokens_per_second": {
360            "type": "double",
361            "value": 2000.0
362          },
363          "gen_ai.usage.cache_read.input_tokens": {
364            "type": "integer",
365            "value": 500
366          },
367          "gen_ai.usage.input_tokens": {
368            "type": "integer",
369            "value": 1000
370          },
371          "gen_ai.usage.output_tokens": {
372            "type": "integer",
373            "value": 2000
374          },
375          "gen_ai.usage.reasoning.output_tokens": {
376            "type": "integer",
377            "value": 1000
378          },
379          "gen_ai.usage.total_tokens": {
380            "type": "double",
381            "value": 3000.0
382          }
383        }
384        "#);
385    }
386
387    #[test]
388    fn test_normalize_ai_basic_tokens() {
389        let mut attributes = Annotated::new(attributes! {
390            "gen_ai.operation.type" => "ai_client".to_owned(),
391            "gen_ai.usage.input_tokens" => 1000,
392            "gen_ai.usage.output_tokens" => 2000,
393            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
394        });
395
396        normalize_ai(
397            &mut attributes,
398            Some(Duration::from_millis(500)),
399            Some(&model_metadata()),
400        );
401
402        assert_annotated_snapshot!(attributes, @r#"
403        {
404          "gen_ai.cost.cache_creation.input_tokens": {
405            "type": "double",
406            "value": 0.0
407          },
408          "gen_ai.cost.cache_read.input_tokens": {
409            "type": "double",
410            "value": 0.0
411          },
412          "gen_ai.cost.input_tokens": {
413            "type": "double",
414            "value": 90.0
415          },
416          "gen_ai.cost.output_tokens": {
417            "type": "double",
418            "value": 100.0
419          },
420          "gen_ai.cost.reasoning.output_tokens": {
421            "type": "double",
422            "value": 0.0
423          },
424          "gen_ai.cost.total_tokens": {
425            "type": "double",
426            "value": 190.0
427          },
428          "gen_ai.operation.type": {
429            "type": "string",
430            "value": "ai_client"
431          },
432          "gen_ai.request.model": {
433            "type": "string",
434            "value": "gpt4-21-04"
435          },
436          "gen_ai.response.model": {
437            "type": "string",
438            "value": "gpt4-21-04"
439          },
440          "gen_ai.response.tokens_per_second": {
441            "type": "double",
442            "value": 4000.0
443          },
444          "gen_ai.usage.input_tokens": {
445            "type": "integer",
446            "value": 1000
447          },
448          "gen_ai.usage.output_tokens": {
449            "type": "integer",
450            "value": 2000
451          },
452          "gen_ai.usage.total_tokens": {
453            "type": "double",
454            "value": 3000.0
455          }
456        }
457        "#);
458    }
459
460    #[test]
461    fn test_normalize_ai_basic_tokens_no_duration_no_cost() {
462        let mut attributes = Annotated::new(attributes! {
463            "gen_ai.operation.type" => "ai_client".to_owned(),
464            "gen_ai.usage.input_tokens" => 1000,
465            "gen_ai.usage.output_tokens" => 2000,
466            "gen_ai.request.model" => "unknown".to_owned(),
467        });
468
469        normalize_ai(
470            &mut attributes,
471            Some(Duration::ZERO),
472            Some(&model_metadata()),
473        );
474
475        assert_annotated_snapshot!(attributes, @r#"
476        {
477          "gen_ai.operation.type": {
478            "type": "string",
479            "value": "ai_client"
480          },
481          "gen_ai.request.model": {
482            "type": "string",
483            "value": "unknown"
484          },
485          "gen_ai.response.model": {
486            "type": "string",
487            "value": "unknown"
488          },
489          "gen_ai.usage.input_tokens": {
490            "type": "integer",
491            "value": 1000
492          },
493          "gen_ai.usage.output_tokens": {
494            "type": "integer",
495            "value": 2000
496          },
497          "gen_ai.usage.total_tokens": {
498            "type": "double",
499            "value": 3000.0
500          }
501        }
502        "#);
503    }
504
505    #[test]
506    fn test_normalize_ai_does_not_overwrite() {
507        let mut attributes = Annotated::new(attributes! {
508            "gen_ai.operation.type" => "ai_client".to_owned(),
509            "gen_ai.usage.input_tokens" => 1000,
510            "gen_ai.usage.output_tokens" => 2000,
511            "gen_ai.request.model" => "gpt4".to_owned(),
512            "gen_ai.response.model" => "gpt4-21-04".to_owned(),
513
514            "gen_ai.cost.input_tokens" => 999.0,
515        });
516
517        normalize_ai(
518            &mut attributes,
519            Some(Duration::from_millis(500)),
520            Some(&model_metadata()),
521        );
522
523        assert_annotated_snapshot!(attributes, @r#"
524        {
525          "gen_ai.cost.cache_creation.input_tokens": {
526            "type": "double",
527            "value": 0.0
528          },
529          "gen_ai.cost.cache_read.input_tokens": {
530            "type": "double",
531            "value": 0.0
532          },
533          "gen_ai.cost.input_tokens": {
534            "type": "double",
535            "value": 90.0
536          },
537          "gen_ai.cost.output_tokens": {
538            "type": "double",
539            "value": 100.0
540          },
541          "gen_ai.cost.reasoning.output_tokens": {
542            "type": "double",
543            "value": 0.0
544          },
545          "gen_ai.cost.total_tokens": {
546            "type": "double",
547            "value": 190.0
548          },
549          "gen_ai.operation.type": {
550            "type": "string",
551            "value": "ai_client"
552          },
553          "gen_ai.request.model": {
554            "type": "string",
555            "value": "gpt4"
556          },
557          "gen_ai.response.model": {
558            "type": "string",
559            "value": "gpt4-21-04"
560          },
561          "gen_ai.response.tokens_per_second": {
562            "type": "double",
563            "value": 4000.0
564          },
565          "gen_ai.usage.input_tokens": {
566            "type": "integer",
567            "value": 1000
568          },
569          "gen_ai.usage.output_tokens": {
570            "type": "integer",
571            "value": 2000
572          },
573          "gen_ai.usage.total_tokens": {
574            "type": "double",
575            "value": 3000.0
576          }
577        }
578        "#);
579    }
580
581    #[test]
582    fn test_normalize_ai_preserves_costs() {
583        let mut attributes = Annotated::new(attributes! {
584            "gen_ai.operation.type" => "ai_client".to_owned(),
585            "gen_ai.usage.input_tokens" => 1000,
586            "gen_ai.usage.output_tokens" => 2000,
587            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
588
589            "gen_ai.usage.total_tokens" => 1337,
590
591            "gen_ai.cost.input_tokens" => 99.0,
592            "gen_ai.cost.output_tokens" => 99.0,
593            "gen_ai.cost.total_tokens" => 123.0,
594
595            "gen_ai.response.tokens_per_second" => 42.0,
596        });
597
598        normalize_ai(
599            &mut attributes,
600            Some(Duration::from_millis(500)),
601            Some(&model_metadata()),
602        );
603
604        assert_annotated_snapshot!(attributes, @r#"
605        {
606          "gen_ai.cost.input_tokens": {
607            "type": "double",
608            "value": 99.0
609          },
610          "gen_ai.cost.output_tokens": {
611            "type": "double",
612            "value": 99.0
613          },
614          "gen_ai.cost.total_tokens": {
615            "type": "double",
616            "value": 123.0
617          },
618          "gen_ai.operation.type": {
619            "type": "string",
620            "value": "ai_client"
621          },
622          "gen_ai.request.model": {
623            "type": "string",
624            "value": "gpt4-21-04"
625          },
626          "gen_ai.response.model": {
627            "type": "string",
628            "value": "gpt4-21-04"
629          },
630          "gen_ai.response.tokens_per_second": {
631            "type": "double",
632            "value": 4000.0
633          },
634          "gen_ai.usage.input_tokens": {
635            "type": "integer",
636            "value": 1000
637          },
638          "gen_ai.usage.output_tokens": {
639            "type": "integer",
640            "value": 2000
641          },
642          "gen_ai.usage.total_tokens": {
643            "type": "double",
644            "value": 3000.0
645          }
646        }
647        "#);
648    }
649
650    #[test]
651    fn test_normalize_ai_no_ai_attributes() {
652        let mut attributes = Annotated::new(attributes! {
653            "gen_ai.usage.input_tokens" => 1000,
654            "gen_ai.usage.output_tokens" => 2000,
655        });
656
657        normalize_ai(
658            &mut attributes,
659            Some(Duration::from_millis(500)),
660            Some(&model_metadata()),
661        );
662
663        assert_annotated_snapshot!(&mut attributes, @r#"
664        {
665          "gen_ai.usage.input_tokens": {
666            "type": "integer",
667            "value": 1000
668          },
669          "gen_ai.usage.output_tokens": {
670            "type": "integer",
671            "value": 2000
672          }
673        }
674        "#);
675    }
676
677    #[test]
678    fn test_normalize_ai_no_ai_indicator_attribute() {
679        let mut attributes = Annotated::new(attributes! {
680            "foo" => 123,
681        });
682
683        normalize_ai(
684            &mut attributes,
685            Some(Duration::from_millis(500)),
686            Some(&model_metadata()),
687        );
688
689        assert_annotated_snapshot!(&mut attributes, @r#"
690        {
691          "foo": {
692            "type": "integer",
693            "value": 123
694          }
695        }
696        "#);
697    }
698
699    #[test]
700    fn test_normalize_ai_empty() {
701        let mut attributes = Annotated::empty();
702
703        normalize_ai(
704            &mut attributes,
705            Some(Duration::from_millis(500)),
706            Some(&model_metadata()),
707        );
708
709        assert!(attributes.is_empty());
710    }
711
712    #[test]
713    fn test_context_utilization_with_total_tokens() {
714        let mut attributes = Annotated::new(attributes! {
715            "gen_ai.operation.type" => "ai_client".to_owned(),
716            "gen_ai.usage.input_tokens" => 30000,
717            "gen_ai.usage.output_tokens" => 12000,
718            "gen_ai.request.model" => "claude-2.1".to_owned(),
719        });
720
721        normalize_ai(
722            &mut attributes,
723            Some(Duration::from_secs(1)),
724            Some(&model_metadata_with_context_size()),
725        );
726
727        assert_annotated_snapshot!(attributes, @r#"
728        {
729          "gen_ai.context.utilization": {
730            "type": "double",
731            "value": 0.42
732          },
733          "gen_ai.context.window_size": {
734            "type": "integer",
735            "value": 100000
736          },
737          "gen_ai.cost.cache_creation.input_tokens": {
738            "type": "double",
739            "value": 0.0
740          },
741          "gen_ai.cost.cache_read.input_tokens": {
742            "type": "double",
743            "value": 0.0
744          },
745          "gen_ai.cost.input_tokens": {
746            "type": "double",
747            "value": 300.0
748          },
749          "gen_ai.cost.output_tokens": {
750            "type": "double",
751            "value": 240.0
752          },
753          "gen_ai.cost.reasoning.output_tokens": {
754            "type": "double",
755            "value": 0.0
756          },
757          "gen_ai.cost.total_tokens": {
758            "type": "double",
759            "value": 540.0
760          },
761          "gen_ai.operation.type": {
762            "type": "string",
763            "value": "ai_client"
764          },
765          "gen_ai.request.model": {
766            "type": "string",
767            "value": "claude-2.1"
768          },
769          "gen_ai.response.model": {
770            "type": "string",
771            "value": "claude-2.1"
772          },
773          "gen_ai.response.tokens_per_second": {
774            "type": "double",
775            "value": 12000.0
776          },
777          "gen_ai.usage.input_tokens": {
778            "type": "integer",
779            "value": 30000
780          },
781          "gen_ai.usage.output_tokens": {
782            "type": "integer",
783            "value": 12000
784          },
785          "gen_ai.usage.total_tokens": {
786            "type": "double",
787            "value": 42000.0
788          }
789        }
790        "#);
791    }
792
793    #[test]
794    fn test_context_utilization_no_context_size() {
795        let mut attributes = Annotated::new(attributes! {
796            "gen_ai.operation.type" => "ai_client".to_owned(),
797            "gen_ai.usage.input_tokens" => 1000,
798            "gen_ai.usage.output_tokens" => 2000,
799            "gen_ai.request.model" => "claude-2.1".to_owned(),
800        });
801
802        // model_metadata() has no context_size set.
803        normalize_ai(
804            &mut attributes,
805            Some(Duration::from_secs(1)),
806            Some(&model_metadata()),
807        );
808
809        let attrs = attributes.value().unwrap();
810        assert!(attrs.get_value("gen_ai.context.window_size").is_none());
811        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
812    }
813
814    #[test]
815    fn test_context_utilization_no_total_tokens() {
816        // Only context_size is available, but no token counts at all.
817        let mut attributes = Annotated::new(attributes! {
818            "gen_ai.operation.type" => "ai_client".to_owned(),
819            "gen_ai.request.model" => "claude-2.1".to_owned(),
820        });
821
822        normalize_ai(
823            &mut attributes,
824            Some(Duration::from_secs(1)),
825            Some(&model_metadata_with_context_size()),
826        );
827
828        let attrs = attributes.value().unwrap();
829        // window_size should still be set even without tokens.
830        assert_eq!(
831            attrs
832                .get_value("gen_ai.context.window_size")
833                .unwrap()
834                .as_f64(),
835            Some(100_000.0)
836        );
837        // But utilization cannot be computed without total_tokens.
838        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
839    }
840
841    #[test]
842    fn test_context_utilization_unknown_model() {
843        let mut attributes = Annotated::new(attributes! {
844            "gen_ai.operation.type" => "ai_client".to_owned(),
845            "gen_ai.usage.input_tokens" => 1000,
846            "gen_ai.usage.output_tokens" => 2000,
847            "gen_ai.request.model" => "unknown-model".to_owned(),
848        });
849
850        normalize_ai(
851            &mut attributes,
852            Some(Duration::from_secs(1)),
853            Some(&model_metadata_with_context_size()),
854        );
855
856        let attrs = attributes.value().unwrap();
857        assert!(attrs.get_value("gen_ai.context.window_size").is_none());
858        assert!(attrs.get_value("gen_ai.context.utilization").is_none());
859    }
860}
relay_event_normalization/eap/ai.rs

relay_event_normalization/eap/
ai.rs
