relay_event_normalization/eap/
ai.rs

1use std::time::Duration;
2
3use relay_conventions::consts::*;
4use relay_event_schema::protocol::Attributes;
5use relay_protocol::Annotated;
6
7use crate::ModelCosts;
8use crate::span::ai;
9
10/// Normalizes AI attributes.
11pub fn normalize_ai(
12    attributes: &mut Annotated<Attributes>,
13    duration: Option<Duration>,
14    costs: Option<&ModelCosts>,
15) {
16    let Some(attributes) = attributes.value_mut() else {
17        return;
18    };
19
20    normalize_total_tokens(attributes);
21    normalize_tokens_per_second(attributes, duration);
22    normalize_ai_costs(attributes, costs);
23}
24
25/// Calculates the [`GEN_AI_USAGE_TOTAL_TOKENS`] attribute.
26fn normalize_total_tokens(attributes: &mut Attributes) {
27    if attributes.contains_key(GEN_AI_USAGE_TOTAL_TOKENS) {
28        return;
29    }
30
31    let input_tokens = attributes
32        .get_value(GEN_AI_USAGE_INPUT_TOKENS)
33        .and_then(|v| v.as_f64());
34
35    let output_tokens = attributes
36        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
37        .and_then(|v| v.as_f64());
38
39    if input_tokens.is_none() && output_tokens.is_none() {
40        return;
41    }
42
43    let total_tokens = input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0);
44    attributes.insert(GEN_AI_USAGE_TOTAL_TOKENS, total_tokens);
45}
46
47/// Calculates the [`GEN_AI_RESPONSE_TPS`] attribute.
48fn normalize_tokens_per_second(attributes: &mut Attributes, duration: Option<Duration>) {
49    let Some(duration) = duration.filter(|d| !d.is_zero()) else {
50        return;
51    };
52
53    if attributes.contains_key(GEN_AI_RESPONSE_TPS) {
54        return;
55    }
56
57    let output_tokens = attributes
58        .get_value(GEN_AI_USAGE_OUTPUT_TOKENS)
59        .and_then(|v| v.as_f64())
60        .filter(|v| *v > 0.0);
61
62    if let Some(output_tokens) = output_tokens {
63        let tps = output_tokens / duration.as_secs_f64();
64        attributes.insert(GEN_AI_RESPONSE_TPS, tps);
65    }
66}
67
68/// Calculates model costs and serializes them into attributes.
69fn normalize_ai_costs(attributes: &mut Attributes, model_costs: Option<&ModelCosts>) {
70    if attributes.contains_key(GEN_AI_COST_TOTAL_TOKENS) {
71        return;
72    }
73
74    let model_cost = attributes
75        .get_value(GEN_AI_REQUEST_MODEL)
76        .or_else(|| attributes.get_value(GEN_AI_RESPONSE_MODEL))
77        .and_then(|v| v.as_str())
78        .and_then(|model| model_costs?.cost_per_token(model));
79
80    let Some(model_cost) = model_cost else { return };
81
82    let get_tokens = |key| {
83        attributes
84            .get_value(key)
85            .and_then(|v| v.as_f64())
86            .unwrap_or(0.0)
87    };
88
89    let tokens = ai::UsedTokens {
90        input_tokens: get_tokens(GEN_AI_USAGE_INPUT_TOKENS),
91        input_cached_tokens: get_tokens(GEN_AI_USAGE_INPUT_CACHED_TOKENS),
92        output_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_TOKENS),
93        output_reasoning_tokens: get_tokens(GEN_AI_USAGE_OUTPUT_REASONING_TOKENS),
94    };
95
96    let Some(costs) = ai::calculate_costs(model_cost, tokens) else {
97        return;
98    };
99
100    // Overwrite all values, the attributes should reflect the values we used to calculate the total.
101    attributes.insert(GEN_AI_COST_INPUT_TOKENS, costs.input);
102    attributes.insert(GEN_AI_COST_OUTPUT_TOKENS, costs.output);
103    attributes.insert(GEN_AI_COST_TOTAL_TOKENS, costs.total());
104}
105
106#[cfg(test)]
107mod tests {
108    use std::collections::HashMap;
109
110    use relay_pattern::Pattern;
111    use relay_protocol::{Empty, assert_annotated_snapshot};
112
113    use crate::ModelCostV2;
114
115    use super::*;
116
117    macro_rules! attributes {
118        ($($key:expr => $value:expr),* $(,)?) => {
119            Attributes::from([
120                $(($key.into(), Annotated::new($value.into())),)*
121            ])
122        };
123    }
124
125    fn model_costs() -> ModelCosts {
126        ModelCosts {
127            version: 2,
128            models: HashMap::from([
129                (
130                    Pattern::new("claude-2.1").unwrap(),
131                    ModelCostV2 {
132                        input_per_token: 0.01,
133                        output_per_token: 0.02,
134                        output_reasoning_per_token: 0.03,
135                        input_cached_per_token: 0.04,
136                    },
137                ),
138                (
139                    Pattern::new("gpt4-21-04").unwrap(),
140                    ModelCostV2 {
141                        input_per_token: 0.09,
142                        output_per_token: 0.05,
143                        output_reasoning_per_token: 0.0,
144                        input_cached_per_token: 0.0,
145                    },
146                ),
147            ]),
148        }
149    }
150
151    #[test]
152    fn test_normalize_ai_all_tokens() {
153        let mut attributes = Annotated::new(attributes! {
154            "gen_ai.usage.input_tokens" => 1000,
155            "gen_ai.usage.output_tokens" => 2000,
156            "gen_ai.usage.output_tokens.reasoning" => 1000,
157            "gen_ai.usage.input_tokens.cached" => 500,
158            "gen_ai.request.model" => "claude-2.1".to_owned(),
159        });
160
161        normalize_ai(
162            &mut attributes,
163            Some(Duration::from_secs(1)),
164            Some(&model_costs()),
165        );
166
167        assert_annotated_snapshot!(attributes, @r#"
168        {
169          "gen_ai.cost.input_tokens": {
170            "type": "double",
171            "value": 25.0
172          },
173          "gen_ai.cost.output_tokens": {
174            "type": "double",
175            "value": 50.0
176          },
177          "gen_ai.cost.total_tokens": {
178            "type": "double",
179            "value": 75.0
180          },
181          "gen_ai.request.model": {
182            "type": "string",
183            "value": "claude-2.1"
184          },
185          "gen_ai.response.tokens_per_second": {
186            "type": "double",
187            "value": 2000.0
188          },
189          "gen_ai.usage.input_tokens": {
190            "type": "integer",
191            "value": 1000
192          },
193          "gen_ai.usage.input_tokens.cached": {
194            "type": "integer",
195            "value": 500
196          },
197          "gen_ai.usage.output_tokens": {
198            "type": "integer",
199            "value": 2000
200          },
201          "gen_ai.usage.output_tokens.reasoning": {
202            "type": "integer",
203            "value": 1000
204          },
205          "gen_ai.usage.total_tokens": {
206            "type": "double",
207            "value": 3000.0
208          }
209        }
210        "#);
211    }
212
213    #[test]
214    fn test_normalize_ai_basic_tokens() {
215        let mut attributes = Annotated::new(attributes! {
216            "gen_ai.usage.input_tokens" => 1000,
217            "gen_ai.usage.output_tokens" => 2000,
218            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
219        });
220
221        normalize_ai(
222            &mut attributes,
223            Some(Duration::from_millis(500)),
224            Some(&model_costs()),
225        );
226
227        assert_annotated_snapshot!(attributes, @r#"
228        {
229          "gen_ai.cost.input_tokens": {
230            "type": "double",
231            "value": 90.0
232          },
233          "gen_ai.cost.output_tokens": {
234            "type": "double",
235            "value": 100.0
236          },
237          "gen_ai.cost.total_tokens": {
238            "type": "double",
239            "value": 190.0
240          },
241          "gen_ai.request.model": {
242            "type": "string",
243            "value": "gpt4-21-04"
244          },
245          "gen_ai.response.tokens_per_second": {
246            "type": "double",
247            "value": 4000.0
248          },
249          "gen_ai.usage.input_tokens": {
250            "type": "integer",
251            "value": 1000
252          },
253          "gen_ai.usage.output_tokens": {
254            "type": "integer",
255            "value": 2000
256          },
257          "gen_ai.usage.total_tokens": {
258            "type": "double",
259            "value": 3000.0
260          }
261        }
262        "#);
263    }
264
265    #[test]
266    fn test_normalize_ai_basic_tokens_no_duration_no_cost() {
267        let mut attributes = Annotated::new(attributes! {
268            "gen_ai.usage.input_tokens" => 1000,
269            "gen_ai.usage.output_tokens" => 2000,
270            "gen_ai.request.model" => "unknown".to_owned(),
271        });
272
273        normalize_ai(&mut attributes, Some(Duration::ZERO), Some(&model_costs()));
274
275        assert_annotated_snapshot!(attributes, @r#"
276        {
277          "gen_ai.request.model": {
278            "type": "string",
279            "value": "unknown"
280          },
281          "gen_ai.usage.input_tokens": {
282            "type": "integer",
283            "value": 1000
284          },
285          "gen_ai.usage.output_tokens": {
286            "type": "integer",
287            "value": 2000
288          },
289          "gen_ai.usage.total_tokens": {
290            "type": "double",
291            "value": 3000.0
292          }
293        }
294        "#);
295    }
296
297    #[test]
298    fn test_normalize_ai_does_not_overwrite() {
299        let mut attributes = Annotated::new(attributes! {
300            "gen_ai.usage.input_tokens" => 1000,
301            "gen_ai.usage.output_tokens" => 2000,
302            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
303
304            "gen_ai.cost.input_tokens" => 999.0,
305        });
306
307        normalize_ai(
308            &mut attributes,
309            Some(Duration::from_millis(500)),
310            Some(&model_costs()),
311        );
312
313        assert_annotated_snapshot!(attributes, @r#"
314        {
315          "gen_ai.cost.input_tokens": {
316            "type": "double",
317            "value": 90.0
318          },
319          "gen_ai.cost.output_tokens": {
320            "type": "double",
321            "value": 100.0
322          },
323          "gen_ai.cost.total_tokens": {
324            "type": "double",
325            "value": 190.0
326          },
327          "gen_ai.request.model": {
328            "type": "string",
329            "value": "gpt4-21-04"
330          },
331          "gen_ai.response.tokens_per_second": {
332            "type": "double",
333            "value": 4000.0
334          },
335          "gen_ai.usage.input_tokens": {
336            "type": "integer",
337            "value": 1000
338          },
339          "gen_ai.usage.output_tokens": {
340            "type": "integer",
341            "value": 2000
342          },
343          "gen_ai.usage.total_tokens": {
344            "type": "double",
345            "value": 3000.0
346          }
347        }
348        "#);
349    }
350
351    #[test]
352    fn test_normalize_ai_overwrite_individual_cost_if_not_total() {
353        let mut attributes = Annotated::new(attributes! {
354            "gen_ai.usage.input_tokens" => 1000,
355            "gen_ai.usage.output_tokens" => 2000,
356            "gen_ai.request.model" => "gpt4-21-04".to_owned(),
357
358            "gen_ai.usage.total_tokens" => 1337,
359
360            "gen_ai.cost.input_tokens" => 99.0,
361            "gen_ai.cost.output_tokens" => 99.0,
362            "gen_ai.cost.total_tokens" => 123.0,
363
364            "gen_ai.response.tokens_per_second" => 42.0,
365        });
366
367        normalize_ai(
368            &mut attributes,
369            Some(Duration::from_millis(500)),
370            Some(&model_costs()),
371        );
372
373        assert_annotated_snapshot!(attributes, @r#"
374        {
375          "gen_ai.cost.input_tokens": {
376            "type": "double",
377            "value": 99.0
378          },
379          "gen_ai.cost.output_tokens": {
380            "type": "double",
381            "value": 99.0
382          },
383          "gen_ai.cost.total_tokens": {
384            "type": "double",
385            "value": 123.0
386          },
387          "gen_ai.request.model": {
388            "type": "string",
389            "value": "gpt4-21-04"
390          },
391          "gen_ai.response.tokens_per_second": {
392            "type": "double",
393            "value": 42.0
394          },
395          "gen_ai.usage.input_tokens": {
396            "type": "integer",
397            "value": 1000
398          },
399          "gen_ai.usage.output_tokens": {
400            "type": "integer",
401            "value": 2000
402          },
403          "gen_ai.usage.total_tokens": {
404            "type": "integer",
405            "value": 1337
406          }
407        }
408        "#);
409    }
410
411    #[test]
412    fn test_normalize_ai_no_ai_attributes() {
413        let mut attributes = Annotated::new(attributes! {
414            "foo" => 123,
415        });
416
417        normalize_ai(
418            &mut attributes,
419            Some(Duration::from_millis(500)),
420            Some(&model_costs()),
421        );
422
423        assert_annotated_snapshot!(&mut attributes, @r#"
424        {
425          "foo": {
426            "type": "integer",
427            "value": 123
428          }
429        }
430        "#);
431    }
432
433    #[test]
434    fn test_normalize_ai_empty() {
435        let mut attributes = Annotated::empty();
436
437        normalize_ai(
438            &mut attributes,
439            Some(Duration::from_millis(500)),
440            Some(&model_costs()),
441        );
442
443        assert!(attributes.is_empty());
444    }
445}
relay_event_normalization/eap/ai.rs

relay_event_normalization/eap/
ai.rs
