relay_event_normalization/normalize/span/
ai.rs

1//! AI cost calculation.
2
3use crate::{ModelCostV2, ModelCosts};
4use relay_event_schema::protocol::{Event, Span, SpanData};
5use relay_protocol::{Annotated, Getter, Value};
6
7/// Calculates the cost of an AI model based on the model cost and the tokens used.
8/// Calculated cost is in US dollars.
9fn calculate_ai_model_cost(model_cost: Option<ModelCostV2>, data: &SpanData) -> Option<f64> {
10    let cost_per_token = model_cost?;
11    let input_tokens_used = data
12        .gen_ai_usage_input_tokens
13        .value()
14        .and_then(Value::as_f64);
15
16    let output_tokens_used = data
17        .gen_ai_usage_output_tokens
18        .value()
19        .and_then(Value::as_f64);
20    let output_reasoning_tokens_used = data
21        .gen_ai_usage_output_tokens_reasoning
22        .value()
23        .and_then(Value::as_f64);
24    let input_cached_tokens_used = data
25        .gen_ai_usage_input_tokens_cached
26        .value()
27        .and_then(Value::as_f64);
28
29    if input_tokens_used.is_none() && output_tokens_used.is_none() {
30        return None;
31    }
32
33    let mut result = 0.0;
34
35    result += cost_per_token.input_per_token * input_tokens_used.unwrap_or(0.0);
36    result += cost_per_token.output_per_token * output_tokens_used.unwrap_or(0.0);
37    result +=
38        cost_per_token.output_reasoning_per_token * output_reasoning_tokens_used.unwrap_or(0.0);
39    result += cost_per_token.input_cached_per_token * input_cached_tokens_used.unwrap_or(0.0);
40
41    Some(result)
42}
43
44/// Maps AI-related measurements (legacy) to span data.
45pub fn map_ai_measurements_to_data(span: &mut Span) {
46    if !is_ai_span(span) {
47        return;
48    };
49
50    let measurements = span.measurements.value();
51    let data = span.data.get_or_insert_with(SpanData::default);
52
53    let set_field_from_measurement = |target_field: &mut Annotated<Value>,
54                                      measurement_key: &str| {
55        if let Some(measurements) = measurements {
56            if target_field.value().is_none() {
57                if let Some(value) = measurements.get_value(measurement_key) {
58                    target_field.set_value(Value::F64(value.to_f64()).into());
59                }
60            }
61        }
62    };
63
64    set_field_from_measurement(&mut data.gen_ai_usage_total_tokens, "ai_total_tokens_used");
65    set_field_from_measurement(&mut data.gen_ai_usage_input_tokens, "ai_prompt_tokens_used");
66    set_field_from_measurement(
67        &mut data.gen_ai_usage_output_tokens,
68        "ai_completion_tokens_used",
69    );
70
71    // It might be that 'total_tokens' is not set in which case we need to calculate it
72    if data.gen_ai_usage_total_tokens.value().is_none() {
73        let input_tokens = data
74            .gen_ai_usage_input_tokens
75            .value()
76            .and_then(Value::as_f64);
77        let output_tokens = data
78            .gen_ai_usage_output_tokens
79            .value()
80            .and_then(Value::as_f64);
81
82        if input_tokens.is_none() && output_tokens.is_none() {
83            // don't set total_tokens if there are no input nor output tokens
84            return;
85        }
86
87        data.gen_ai_usage_total_tokens.set_value(
88            Value::F64(input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0)).into(),
89        );
90    }
91}
92
93/// Extract the additional data into the span
94pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) {
95    if !is_ai_span(span) {
96        return;
97    }
98
99    let duration = span
100        .get_value("span.duration")
101        .and_then(|v| v.as_f64())
102        .unwrap_or(0.0);
103
104    let Some(data) = span.data.value_mut() else {
105        return;
106    };
107
108    // Extracts the response tokens per second
109    if data.gen_ai_response_tokens_per_second.value().is_none() && duration > 0.0 {
110        if let Some(output_tokens) = data
111            .gen_ai_usage_output_tokens
112            .value()
113            .and_then(Value::as_f64)
114        {
115            data.gen_ai_response_tokens_per_second
116                .set_value(Value::F64(output_tokens / (duration / 1000.0)).into());
117        }
118    }
119
120    // Extracts the total cost of the AI model used
121    if let Some(model_id) = data
122        .gen_ai_request_model
123        .value()
124        .and_then(|val| val.as_str())
125        // xxx (vgrozdanic): temporal fallback to legacy field, until we fix
126        // sentry conventions and standardize what SDKs send
127        .or_else(|| data.ai_model_id.value().and_then(|val| val.as_str()))
128    {
129        if let Some(total_cost) =
130            calculate_ai_model_cost(ai_model_costs.cost_per_token(model_id), data)
131        {
132            data.gen_ai_usage_total_cost
133                .set_value(Value::F64(total_cost).into());
134        }
135    }
136}
137
138/// Extract the ai data from all of an event's spans
139pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) {
140    let spans = event.spans.value_mut().iter_mut().flatten();
141    let spans = spans.filter_map(|span| span.value_mut().as_mut());
142
143    for span in spans {
144        map_ai_measurements_to_data(span);
145        if let Some(model_costs) = model_costs {
146            extract_ai_data(span, model_costs);
147        }
148    }
149}
150
151/// Returns true if the span is an AI span.
152/// AI spans are spans with op starting with "ai." (legacy) or "gen_ai." (new).
153pub fn is_ai_span(span: &Span) -> bool {
154    span.op
155        .value()
156        .is_some_and(|op| op.starts_with("ai.") || op.starts_with("gen_ai."))
157}
relay_event_normalization/normalize/span/ai.rs

relay_event_normalization/normalize/span/
ai.rs
