relay_event_normalization/normalize/span/
ai.rs

1//! AI cost calculation.
2
3use crate::ModelCosts;
4use relay_event_schema::protocol::{Event, Span, SpanData};
5use relay_protocol::{Annotated, Value};
6
7/// Calculated cost is in US dollars.
8fn calculate_ai_model_cost(
9    model_id: &str,
10    prompt_tokens_used: Option<f64>,
11    completion_tokens_used: Option<f64>,
12    total_tokens_used: Option<f64>,
13    ai_model_costs: &ModelCosts,
14) -> Option<f64> {
15    if let Some(prompt_tokens) = prompt_tokens_used {
16        if let Some(completion_tokens) = completion_tokens_used {
17            let mut result = 0.0;
18            if let Some(cost_per_1k) = ai_model_costs.cost_per_1k_tokens(model_id, false) {
19                result += cost_per_1k * (prompt_tokens / 1000.0)
20            }
21            if let Some(cost_per_1k) = ai_model_costs.cost_per_1k_tokens(model_id, true) {
22                result += cost_per_1k * (completion_tokens / 1000.0)
23            }
24            return Some(result);
25        }
26    }
27    if let Some(total_tokens) = total_tokens_used {
28        ai_model_costs
29            .cost_per_1k_tokens(model_id, false)
30            .map(|cost| cost * (total_tokens / 1000.0))
31    } else {
32        None
33    }
34}
35
36/// Maps AI-related measurements (legacy) to span data.
37pub fn map_ai_measurements_to_data(span: &mut Span) {
38    if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
39        return;
40    };
41
42    let measurements = span.measurements.value();
43    let data = span.data.get_or_insert_with(SpanData::default);
44
45    let set_field_from_measurement = |target_field: &mut Annotated<Value>,
46                                      measurement_key: &str| {
47        if let Some(measurements) = measurements {
48            if target_field.value().is_none() {
49                if let Some(value) = measurements.get_value(measurement_key) {
50                    target_field.set_value(Value::F64(value).into());
51                }
52            }
53        }
54    };
55
56    set_field_from_measurement(&mut data.gen_ai_usage_total_tokens, "ai_total_tokens_used");
57    set_field_from_measurement(&mut data.gen_ai_usage_input_tokens, "ai_prompt_tokens_used");
58    set_field_from_measurement(
59        &mut data.gen_ai_usage_output_tokens,
60        "ai_completion_tokens_used",
61    );
62
63    // It might be that 'total_tokens' is not set in which case we need to calculate it
64    if data.gen_ai_usage_total_tokens.value().is_none() {
65        let input_tokens = data
66            .gen_ai_usage_input_tokens
67            .value()
68            .and_then(Value::as_f64)
69            .unwrap_or(0.0);
70        let output_tokens = data
71            .gen_ai_usage_output_tokens
72            .value()
73            .and_then(Value::as_f64)
74            .unwrap_or(0.0);
75        data.gen_ai_usage_total_tokens
76            .set_value(Value::F64(input_tokens + output_tokens).into());
77    }
78}
79
80/// Extract the gen_ai_usage_total_cost data into the span
81pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) {
82    if !span.op.value().is_some_and(|op| op.starts_with("ai.")) {
83        return;
84    }
85
86    let Some(data) = span.data.value_mut() else {
87        return;
88    };
89
90    let total_tokens_used = data
91        .gen_ai_usage_total_tokens
92        .value()
93        .and_then(Value::as_f64);
94    let prompt_tokens_used = data
95        .gen_ai_usage_input_tokens
96        .value()
97        .and_then(Value::as_f64);
98    let completion_tokens_used = data
99        .gen_ai_usage_output_tokens
100        .value()
101        .and_then(Value::as_f64);
102
103    if let Some(model_id) = data.ai_model_id.value().and_then(|val| val.as_str()) {
104        if let Some(total_cost) = calculate_ai_model_cost(
105            model_id,
106            prompt_tokens_used,
107            completion_tokens_used,
108            total_tokens_used,
109            ai_model_costs,
110        ) {
111            data.gen_ai_usage_total_cost
112                .set_value(Value::F64(total_cost).into());
113        }
114    }
115}
116
117/// Extract the ai data from all of an event's spans
118pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) {
119    let spans = event.spans.value_mut().iter_mut().flatten();
120    let spans = spans.filter_map(|span| span.value_mut().as_mut());
121
122    for span in spans {
123        map_ai_measurements_to_data(span);
124        if let Some(model_costs) = model_costs {
125            extract_ai_data(span, model_costs);
126        }
127    }
128}
relay_event_normalization/normalize/span/ai.rs

relay_event_normalization/normalize/span/
ai.rs
