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    // Cached tokens are subset of the input tokens, so we need to subtract them
36    // from the input tokens
37    result += cost_per_token.input_per_token
38        * (input_tokens_used.unwrap_or(0.0) - input_cached_tokens_used.unwrap_or(0.0));
39    result += cost_per_token.input_cached_per_token * input_cached_tokens_used.unwrap_or(0.0);
40    // Reasoning tokens are subset of the output tokens, so we need to subtract
41    // them from the output tokens
42    result += cost_per_token.output_per_token
43        * (output_tokens_used.unwrap_or(0.0) - output_reasoning_tokens_used.unwrap_or(0.0));
44
45    if cost_per_token.output_reasoning_per_token > 0.0 {
46        // for now most of the models do not differentiate between reasoning and output token cost,
47        // it costs the same
48        result +=
49            cost_per_token.output_reasoning_per_token * output_reasoning_tokens_used.unwrap_or(0.0);
50    } else {
51        result += cost_per_token.output_per_token * output_reasoning_tokens_used.unwrap_or(0.0);
52    }
53
54    Some(result)
55}
56
57/// Maps AI-related measurements (legacy) to span data.
58pub fn map_ai_measurements_to_data(span: &mut Span) {
59    if !is_ai_span(span) {
60        return;
61    };
62
63    let measurements = span.measurements.value();
64    let data = span.data.get_or_insert_with(SpanData::default);
65
66    let set_field_from_measurement = |target_field: &mut Annotated<Value>,
67                                      measurement_key: &str| {
68        if let Some(measurements) = measurements
69            && target_field.value().is_none()
70            && let Some(value) = measurements.get_value(measurement_key)
71        {
72            target_field.set_value(Value::F64(value.to_f64()).into());
73        }
74    };
75
76    set_field_from_measurement(&mut data.gen_ai_usage_total_tokens, "ai_total_tokens_used");
77    set_field_from_measurement(&mut data.gen_ai_usage_input_tokens, "ai_prompt_tokens_used");
78    set_field_from_measurement(
79        &mut data.gen_ai_usage_output_tokens,
80        "ai_completion_tokens_used",
81    );
82
83    // It might be that 'total_tokens' is not set in which case we need to calculate it
84    if data.gen_ai_usage_total_tokens.value().is_none() {
85        let input_tokens = data
86            .gen_ai_usage_input_tokens
87            .value()
88            .and_then(Value::as_f64);
89        let output_tokens = data
90            .gen_ai_usage_output_tokens
91            .value()
92            .and_then(Value::as_f64);
93
94        if input_tokens.is_none() && output_tokens.is_none() {
95            // don't set total_tokens if there are no input nor output tokens
96            return;
97        }
98
99        data.gen_ai_usage_total_tokens.set_value(
100            Value::F64(input_tokens.unwrap_or(0.0) + output_tokens.unwrap_or(0.0)).into(),
101        );
102    }
103}
104
105/// Extract the additional data into the span
106pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) {
107    if !is_ai_span(span) {
108        return;
109    }
110
111    let duration = span
112        .get_value("span.duration")
113        .and_then(|v| v.as_f64())
114        .unwrap_or(0.0);
115
116    let Some(data) = span.data.value_mut() else {
117        return;
118    };
119
120    // Extracts the response tokens per second
121    if data.gen_ai_response_tokens_per_second.value().is_none()
122        && duration > 0.0
123        && let Some(output_tokens) = data
124            .gen_ai_usage_output_tokens
125            .value()
126            .and_then(Value::as_f64)
127    {
128        data.gen_ai_response_tokens_per_second
129            .set_value(Value::F64(output_tokens / (duration / 1000.0)).into());
130    }
131
132    // Extracts the total cost of the AI model used
133    if let Some(model_id) = data
134        .gen_ai_request_model
135        .value()
136        .and_then(|val| val.as_str())
137        .or_else(|| {
138            data.gen_ai_response_model
139                .value()
140                .and_then(|val| val.as_str())
141        })
142        && let Some(total_cost) =
143            calculate_ai_model_cost(ai_model_costs.cost_per_token(model_id), data)
144    {
145        data.gen_ai_usage_total_cost
146            .set_value(Value::F64(total_cost).into());
147    }
148}
149
150/// Extract the ai data from all of an event's spans
151pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) {
152    let spans = event.spans.value_mut().iter_mut().flatten();
153    let spans = spans.filter_map(|span| span.value_mut().as_mut());
154
155    for span in spans {
156        map_ai_measurements_to_data(span);
157        if let Some(model_costs) = model_costs {
158            extract_ai_data(span, model_costs);
159        }
160    }
161}
162
163/// Returns true if the span is an AI span.
164/// AI spans are spans with op starting with "ai." (legacy) or "gen_ai." (new).
165pub fn is_ai_span(span: &Span) -> bool {
166    span.op
167        .value()
168        .is_some_and(|op| op.starts_with("ai.") || op.starts_with("gen_ai."))
169}
relay_event_normalization/normalize/span/ai.rs

relay_event_normalization/normalize/span/
ai.rs
