Skip to content

Commit dc8ee98

Browse files
Dav1ddeDav1dde
committed
ref(ai): Embedd AI operation types into Relay
1 parent 125a8ea commit dc8ee98

File tree

10 files changed

+114
-427
lines changed

10 files changed

+114
-427
lines changed

‎CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
**Internal**:
66

7+
- Embed AI operation type mappings into Relay. ([#5555](https://github.com/getsentry/relay/pull/5555))
78
- Use new processor architecture to process transactions. ([#5379](https://github.com/getsentry/relay/pull/5379))
89

910

‎relay-cabi/src/processing.rs‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -267,9 +267,8 @@ pub unsafe extern "C" fn relay_store_normalizer_normalize_event(
267267
max_tag_value_length: usize::MAX,
268268
span_description_rules: None,
269269
performance_score: None,
270-
geoip_lookup: None, // only supported in relay
271-
ai_model_costs: None, // only supported in relay
272-
ai_operation_type_map: None, // only supported in relay
270+
geoip_lookup: None, // only supported in relay
271+
ai_model_costs: None, // only supported in relay
273272
enable_trimming: config.enable_trimming.unwrap_or_default(),
274273
measurements: None,
275274
normalize_spans: config.normalize_spans,

‎relay-conventions/src/consts.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@ convention_attributes!(
3636
GEN_AI_COST_INPUT_TOKENS => "gen_ai.cost.input_tokens",
3737
GEN_AI_COST_OUTPUT_TOKENS => "gen_ai.cost.output_tokens",
3838
GEN_AI_COST_TOTAL_TOKENS => "gen_ai.cost.total_tokens",
39+
GEN_AI_OPERATION_TYPE => "gen_ai.operation.type",
40+
GEN_AI_OPERATION_NAME => "gen_ai.operation.name",
3941
GEN_AI_REQUEST_MODEL => "gen_ai.request.model",
4042
GEN_AI_RESPONSE_MODEL => "gen_ai.response.model",
4143
GEN_AI_RESPONSE_TPS => "gen_ai.response.tokens_per_second",

‎relay-dynamic-config/src/global.rs‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,7 @@ use std::io::BufReader;
55
use std::path::Path;
66

77
use relay_base_schema::metrics::MetricNamespace;
8-
use relay_event_normalization::{
9-
AiOperationTypeMap, MeasurementsConfig, ModelCosts, SpanOpDefaults,
10-
};
8+
use relay_event_normalization::{MeasurementsConfig, ModelCosts, SpanOpDefaults};
119
use relay_filter::GenericFiltersConfig;
1210
use relay_quotas::Quota;
1311
use serde::{Deserialize, Serialize, de};
@@ -52,10 +50,6 @@ pub struct GlobalConfig {
5250
#[serde(skip_serializing_if = "is_model_costs_empty")]
5351
pub ai_model_costs: ErrorBoundary<ModelCosts>,
5452

55-
/// Configuration to derive the `gen_ai.operation.type` field from other fields
56-
#[serde(skip_serializing_if = "is_ai_operation_type_map_empty")]
57-
pub ai_operation_type_map: ErrorBoundary<AiOperationTypeMap>,
58-
5953
/// Configuration to derive the `span.op` from other span fields.
6054
#[serde(
6155
deserialize_with = "default_on_error",
@@ -351,10 +345,6 @@ fn is_model_costs_empty(value: &ErrorBoundary<ModelCosts>) -> bool {
351345
matches!(value, ErrorBoundary::Ok(model_costs) if model_costs.is_empty())
352346
}
353347

354-
fn is_ai_operation_type_map_empty(value: &ErrorBoundary<AiOperationTypeMap>) -> bool {
355-
matches!(value, ErrorBoundary::Ok(ai_operation_type_map) if ai_operation_type_map.is_empty())
356-
}
357-
358348
#[cfg(test)]
359349
mod tests {
360350
use super::*;

‎relay-event-normalization/src/event.rs‎

Lines changed: 24 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,6 @@ use relay_protocol::{
2727
use smallvec::SmallVec;
2828
use uuid::Uuid;
2929

30-
use crate::normalize::AiOperationTypeMap;
3130
use crate::normalize::request;
3231
use crate::span::ai::enrich_ai_event_data;
3332
use crate::span::tag_extraction::extract_span_tags_from_event;
@@ -141,9 +140,6 @@ pub struct NormalizationConfig<'a> {
141140
/// Configuration for calculating the cost of AI model runs
142141
pub ai_model_costs: Option<&'a ModelCosts>,
143142

144-
/// Configuration for mapping AI operation types from span.op to gen_ai.operation.type
145-
pub ai_operation_type_map: Option<&'a AiOperationTypeMap>,
146-
147143
/// An initialized GeoIP lookup.
148144
pub geoip_lookup: Option<&'a GeoIpLookup>,
149145

@@ -198,7 +194,6 @@ impl Default for NormalizationConfig<'_> {
198194
performance_score: Default::default(),
199195
geoip_lookup: Default::default(),
200196
ai_model_costs: Default::default(),
201-
ai_operation_type_map: Default::default(),
202197
enable_trimming: false,
203198
measurements: None,
204199
normalize_spans: true,
@@ -328,7 +323,7 @@ fn normalize(event: &mut Event, meta: &mut Meta, config: &NormalizationConfig) {
328323
.get_or_default::<PerformanceScoreContext>()
329324
.score_profile_version = Annotated::new(version);
330325
}
331-
enrich_ai_event_data(event, config.ai_model_costs, config.ai_operation_type_map);
326+
enrich_ai_event_data(event, config.ai_model_costs);
332327
normalize_breakdowns(event, config.breakdowns_config); // Breakdowns are part of the metric extraction too
333328
normalize_default_attributes(event, meta, config);
334329
normalize_trace_context_tags(event);
@@ -2334,7 +2329,8 @@ mod tests {
23342329
"gen_ai.cost.total_tokens": 50.0,
23352330
"gen_ai.cost.input_tokens": 10.0,
23362331
"gen_ai.cost.output_tokens": 40.0,
2337-
"gen_ai.response.tokens_per_second": 62500.0
2332+
"gen_ai.response.tokens_per_second": 62500.0,
2333+
"gen_ai.operation.type": "ai_client"
23382334
}
23392335
"#);
23402336
assert_annotated_snapshot!(span2, @r#"
@@ -2346,7 +2342,8 @@ mod tests {
23462342
"gen_ai.cost.total_tokens": 80.0,
23472343
"gen_ai.cost.input_tokens": 20.0,
23482344
"gen_ai.cost.output_tokens": 60.0,
2349-
"gen_ai.response.tokens_per_second": 62500.0
2345+
"gen_ai.response.tokens_per_second": 62500.0,
2346+
"gen_ai.operation.type": "ai_client"
23502347
}
23512348
"#);
23522349
}
@@ -2451,7 +2448,8 @@ mod tests {
24512448
"gen_ai.cost.total_tokens": 75.0,
24522449
"gen_ai.cost.input_tokens": 25.0,
24532450
"gen_ai.cost.output_tokens": 50.0,
2454-
"gen_ai.response.tokens_per_second": 2000.0
2451+
"gen_ai.response.tokens_per_second": 2000.0,
2452+
"gen_ai.operation.type": "ai_client"
24552453
}
24562454
"#);
24572455
assert_annotated_snapshot!(span2, @r#"
@@ -2463,7 +2461,8 @@ mod tests {
24632461
"gen_ai.cost.total_tokens": 190.0,
24642462
"gen_ai.cost.input_tokens": 90.0,
24652463
"gen_ai.cost.output_tokens": 100.0,
2466-
"gen_ai.response.tokens_per_second": 2000.0
2464+
"gen_ai.response.tokens_per_second": 2000.0,
2465+
"gen_ai.operation.type": "ai_client"
24672466
}
24682467
"#);
24692468
assert_annotated_snapshot!(span3, @r#"
@@ -2475,7 +2474,8 @@ mod tests {
24752474
"gen_ai.cost.total_tokens": 190.0,
24762475
"gen_ai.cost.input_tokens": 90.0,
24772476
"gen_ai.cost.output_tokens": 100.0,
2478-
"gen_ai.response.tokens_per_second": 2000.0
2477+
"gen_ai.response.tokens_per_second": 2000.0,
2478+
"gen_ai.operation.type": "ai_client"
24792479
}
24802480
"#);
24812481
}
@@ -2527,7 +2527,8 @@ mod tests {
25272527

25282528
assert_annotated_snapshot!(span, @r#"
25292529
{
2530-
"gen_ai.request.model": "claude-2.1"
2530+
"gen_ai.request.model": "claude-2.1",
2531+
"gen_ai.operation.type": "agent"
25312532
}
25322533
"#);
25332534
}
@@ -2618,7 +2619,8 @@ mod tests {
26182619
"gen_ai.cost.total_tokens": 65.0,
26192620
"gen_ai.cost.input_tokens": 25.0,
26202621
"gen_ai.cost.output_tokens": 40.0,
2621-
"gen_ai.response.tokens_per_second": 62500.0
2622+
"gen_ai.response.tokens_per_second": 62500.0,
2623+
"gen_ai.operation.type": "ai_client"
26222624
}
26232625
"#);
26242626
assert_annotated_snapshot!(span2, @r#"
@@ -2630,7 +2632,8 @@ mod tests {
26302632
"gen_ai.cost.total_tokens": 190.0,
26312633
"gen_ai.cost.input_tokens": 90.0,
26322634
"gen_ai.cost.output_tokens": 100.0,
2633-
"gen_ai.response.tokens_per_second": 62500.0
2635+
"gen_ai.response.tokens_per_second": 62500.0,
2636+
"gen_ai.operation.type": "ai_client"
26342637
}
26352638
"#);
26362639
}
@@ -2673,7 +2676,8 @@ mod tests {
26732676
assert_annotated_snapshot!(span, @r#"
26742677
{
26752678
"gen_ai.usage.total_tokens": 500.0,
2676-
"gen_ai.usage.input_tokens": 500
2679+
"gen_ai.usage.input_tokens": 500,
2680+
"gen_ai.operation.type": "ai_client"
26772681
}
26782682
"#);
26792683
}
@@ -2716,7 +2720,8 @@ mod tests {
27162720
assert_annotated_snapshot!(span, @r#"
27172721
{
27182722
"gen_ai.usage.total_tokens": 1000.0,
2719-
"gen_ai.usage.output_tokens": 1000
2723+
"gen_ai.usage.output_tokens": 1000,
2724+
"gen_ai.operation.type": "ai_client"
27202725
}
27212726
"#);
27222727
}
@@ -2749,40 +2754,13 @@ mod tests {
27492754

27502755
let mut event = Annotated::<Event>::from_json(json).unwrap();
27512756

2752-
let operation_type_map = AiOperationTypeMap {
2753-
version: 1,
2754-
operation_types: HashMap::from([
2755-
(Pattern::new("gen_ai.chat").unwrap(), "chat".to_owned()),
2756-
(
2757-
Pattern::new("gen_ai.execute_tool").unwrap(),
2758-
"execute_tool".to_owned(),
2759-
),
2760-
(
2761-
Pattern::new("gen_ai.handoff").unwrap(),
2762-
"handoff".to_owned(),
2763-
),
2764-
(
2765-
Pattern::new("gen_ai.invoke_agent").unwrap(),
2766-
"invoke_agent".to_owned(),
2767-
),
2768-
// fallback to agent
2769-
(Pattern::new("gen_ai.*").unwrap(), "agent".to_owned()),
2770-
]),
2771-
};
2772-
2773-
normalize_event(
2774-
&mut event,
2775-
&NormalizationConfig {
2776-
ai_operation_type_map: Some(&operation_type_map),
2777-
..NormalizationConfig::default()
2778-
},
2779-
);
2757+
normalize_event(&mut event, &NormalizationConfig::default());
27802758

27812759
let [span1, span2, span3] = collect_span_data(event);
27822760

27832761
assert_annotated_snapshot!(span1, @r#"
27842762
{
2785-
"gen_ai.operation.type": "chat"
2763+
"gen_ai.operation.type": "ai_client"
27862764
}
27872765
"#);
27882766
assert_annotated_snapshot!(span2, @r#"
@@ -2792,88 +2770,11 @@ mod tests {
27922770
"#);
27932771
assert_annotated_snapshot!(span3, @r#"
27942772
{
2795-
"gen_ai.operation.type": "agent"
2773+
"gen_ai.operation.type": "ai_client"
27962774
}
27972775
"#);
27982776
}
27992777

2800-
#[test]
2801-
fn test_ai_operation_type_disabled_map() {
2802-
let json = r#"
2803-
{
2804-
"type": "transaction",
2805-
"transaction": "test-transaction",
2806-
"spans": [
2807-
{
2808-
"op": "gen_ai.chat",
2809-
"description": "AI chat completion",
2810-
"data": {}
2811-
}
2812-
]
2813-
}
2814-
"#;
2815-
2816-
let mut event = Annotated::<Event>::from_json(json).unwrap();
2817-
2818-
let operation_type_map = AiOperationTypeMap {
2819-
version: 0, // Disabled version
2820-
operation_types: HashMap::from([(
2821-
Pattern::new("gen_ai.chat").unwrap(),
2822-
"chat".to_owned(),
2823-
)]),
2824-
};
2825-
2826-
normalize_event(
2827-
&mut event,
2828-
&NormalizationConfig {
2829-
ai_operation_type_map: Some(&operation_type_map),
2830-
..NormalizationConfig::default()
2831-
},
2832-
);
2833-
2834-
let [span] = collect_span_data(event);
2835-
2836-
// Should not set operation type when map is disabled
2837-
assert_annotated_snapshot!(span, @"{}");
2838-
}
2839-
2840-
#[test]
2841-
fn test_ai_operation_type_empty_map() {
2842-
let json = r#"
2843-
{
2844-
"type": "transaction",
2845-
"transaction": "test-transaction",
2846-
"spans": [
2847-
{
2848-
"op": "gen_ai.chat",
2849-
"description": "AI chat completion",
2850-
"data": {}
2851-
}
2852-
]
2853-
}
2854-
"#;
2855-
2856-
let mut event = Annotated::<Event>::from_json(json).unwrap();
2857-
2858-
let operation_type_map = AiOperationTypeMap {
2859-
version: 1,
2860-
operation_types: HashMap::new(),
2861-
};
2862-
2863-
normalize_event(
2864-
&mut event,
2865-
&NormalizationConfig {
2866-
ai_operation_type_map: Some(&operation_type_map),
2867-
..NormalizationConfig::default()
2868-
},
2869-
);
2870-
2871-
let [span] = collect_span_data(event);
2872-
2873-
// Should not set operation type when map is empty
2874-
assert_annotated_snapshot!(span, @"{}");
2875-
}
2876-
28772778
#[test]
28782779
fn test_apple_high_device_class() {
28792780
let mut event = Event {

0 commit comments

Comments
 (0)