relay_event_normalization/normalize/span/
ai.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
//! AI cost calculation.

use crate::ModelCosts;
use relay_base_schema::metrics::MetricUnit;
use relay_event_schema::protocol::{Event, Measurement, Span};

/// Calculated cost is in US dollars.
fn calculate_ai_model_cost(
    model_id: &str,
    prompt_tokens_used: Option<f64>,
    completion_tokens_used: Option<f64>,
    total_tokens_used: Option<f64>,
    ai_model_costs: &ModelCosts,
) -> Option<f64> {
    if let Some(prompt_tokens) = prompt_tokens_used {
        if let Some(completion_tokens) = completion_tokens_used {
            let mut result = 0.0;
            if let Some(cost_per_1k) = ai_model_costs.cost_per_1k_tokens(model_id, false) {
                result += cost_per_1k * (prompt_tokens / 1000.0)
            }
            if let Some(cost_per_1k) = ai_model_costs.cost_per_1k_tokens(model_id, true) {
                result += cost_per_1k * (completion_tokens / 1000.0)
            }
            return Some(result);
        }
    }
    if let Some(total_tokens) = total_tokens_used {
        ai_model_costs
            .cost_per_1k_tokens(model_id, false)
            .map(|cost| cost * (total_tokens / 1000.0))
    } else {
        None
    }
}

/// Extract the ai_total_cost measurement into the span.
pub fn extract_ai_measurements(span: &mut Span, ai_model_costs: &ModelCosts) {
    let Some(span_op) = span.op.value() else {
        return;
    };

    if !span_op.starts_with("ai.") {
        return;
    }

    let Some(measurements) = span.measurements.value() else {
        return;
    };

    let total_tokens_used = measurements.get_value("ai_total_tokens_used");
    let prompt_tokens_used = measurements.get_value("ai_prompt_tokens_used");
    let completion_tokens_used = measurements.get_value("ai_completion_tokens_used");
    if let Some(model_id) = span
        .data
        .value()
        .and_then(|d| d.ai_model_id.value())
        .and_then(|val| val.as_str())
    {
        if let Some(total_cost) = calculate_ai_model_cost(
            model_id,
            prompt_tokens_used,
            completion_tokens_used,
            total_tokens_used,
            ai_model_costs,
        ) {
            span.measurements
                .get_or_insert_with(Default::default)
                .insert(
                    "ai_total_cost".to_owned(),
                    Measurement {
                        value: total_cost.into(),
                        unit: MetricUnit::None.into(),
                    }
                    .into(),
                );
        }
    }
}

/// Extract the ai_total_cost measurements from all of an event's spans
pub fn normalize_ai_measurements(event: &mut Event, model_costs: Option<&ModelCosts>) {
    if let Some(model_costs) = model_costs {
        if let Some(spans) = event.spans.value_mut() {
            for span in spans {
                if let Some(mut_span) = span.value_mut() {
                    extract_ai_measurements(mut_span, model_costs);
                }
            }
        }
    }
}