diff --git a/CHANGELOG.md b/CHANGELOG.md index ae688de41df..539c8a468f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,7 +24,8 @@ - Add option gating Snuba publishing to ingest-replay-events for Replays. ([#5088](https://github.com/getsentry/relay/pull/5088), [#5115](https://github.com/getsentry/relay/pull/5115)) - Add gen_ai_cost_total_tokens attribute and double write total tokens cost. ([#5121](https://github.com/getsentry/relay/pull/5121)) - Change mapping of incoming OTLP spans with `ERROR` status to Sentry's `internal_error` status. ([#5127](https://github.com/getsentry/relay/pull/5127)) -- Add `ai_operation_type_map` to global config ([#5125](https://github.com/getsentry/relay/pull/5125)) +- Add `ai_operation_type_map` to global config. ([#5125](https://github.com/getsentry/relay/pull/5125)) +- Infer `gen_ai.operation.type` from `span.op`. ([#5129](https://github.com/getsentry/relay/pull/5129)) ## 25.8.0 diff --git a/relay-cabi/src/processing.rs b/relay-cabi/src/processing.rs index 1bd123ece50..d0845608516 100644 --- a/relay-cabi/src/processing.rs +++ b/relay-cabi/src/processing.rs @@ -267,8 +267,9 @@ pub unsafe extern "C" fn relay_store_normalizer_normalize_event( max_tag_value_length: usize::MAX, span_description_rules: None, performance_score: None, - geoip_lookup: None, // only supported in relay - ai_model_costs: None, // only supported in relay + geoip_lookup: None, // only supported in relay + ai_model_costs: None, // only supported in relay + ai_operation_type_map: None, // only supported in relay enable_trimming: config.enable_trimming.unwrap_or_default(), measurements: None, normalize_spans: config.normalize_spans, diff --git a/relay-event-normalization/src/event.rs b/relay-event-normalization/src/event.rs index b588676d600..d64994a95cc 100644 --- a/relay-event-normalization/src/event.rs +++ b/relay-event-normalization/src/event.rs @@ -27,8 +27,9 @@ use relay_protocol::{ use smallvec::SmallVec; use uuid::Uuid; +use crate::normalize::AiOperationTypeMap; use crate::normalize::request; -use crate::span::ai::enrich_ai_span_data; +use crate::span::ai::enrich_ai_event_data; use crate::span::tag_extraction::extract_span_tags_from_event; use crate::utils::{self, MAX_DURATION_MOBILE_MS, get_event_user_tag}; use crate::{ @@ -140,6 +141,9 @@ pub struct NormalizationConfig<'a> { /// Configuration for calculating the cost of AI model runs pub ai_model_costs: Option<&'a ModelCosts>, + /// Configuration for mapping AI operation types from span.op to gen_ai.operation.type + pub ai_operation_type_map: Option<&'a AiOperationTypeMap>, + /// An initialized GeoIP lookup. pub geoip_lookup: Option<&'a GeoIpLookup>, @@ -194,6 +198,7 @@ impl Default for NormalizationConfig<'_> { performance_score: Default::default(), geoip_lookup: Default::default(), ai_model_costs: Default::default(), + ai_operation_type_map: Default::default(), enable_trimming: false, measurements: None, normalize_spans: true, @@ -323,7 +328,7 @@ fn normalize(event: &mut Event, meta: &mut Meta, config: &NormalizationConfig) { .get_or_default::() .score_profile_version = Annotated::new(version); } - enrich_ai_span_data(event, config.ai_model_costs); + enrich_ai_event_data(event, config.ai_model_costs, config.ai_operation_type_map); normalize_breakdowns(event, config.breakdowns_config); // Breakdowns are part of the metric extraction too normalize_default_attributes(event, meta, config); normalize_trace_context_tags(event); @@ -2847,6 +2852,161 @@ mod tests { ); } + #[test] + fn test_ai_operation_type_mapping() { + let json = r#" + { + "type": "transaction", + "transaction": "test-transaction", + "spans": [ + { + "op": "gen_ai.chat", + "description": "AI chat completion", + "data": {} + }, + { + "op": "gen_ai.handoff", + "description": "AI agent handoff", + "data": {} + }, + { + "op": "gen_ai.unknown", + "description": "Unknown AI operation", + "data": {} + } + ] + } + "#; + + let mut event = Annotated::::from_json(json).unwrap(); + + let operation_type_map = AiOperationTypeMap { + version: 1, + operation_types: HashMap::from([ + (Pattern::new("gen_ai.chat").unwrap(), "chat".to_owned()), + ( + Pattern::new("gen_ai.execute_tool").unwrap(), + "execute_tool".to_owned(), + ), + ( + Pattern::new("gen_ai.handoff").unwrap(), + "handoff".to_owned(), + ), + ( + Pattern::new("gen_ai.invoke_agent").unwrap(), + "invoke_agent".to_owned(), + ), + // fallback to agent + (Pattern::new("gen_ai.*").unwrap(), "agent".to_owned()), + ]), + }; + + normalize_event( + &mut event, + &NormalizationConfig { + ai_operation_type_map: Some(&operation_type_map), + ..NormalizationConfig::default() + }, + ); + + let span_data_0 = get_value!(event.spans[0].data!); + let span_data_1 = get_value!(event.spans[1].data!); + let span_data_2 = get_value!(event.spans[2].data!); + + assert_eq!( + span_data_0.gen_ai_operation_type.value(), + Some(&"chat".to_owned()) + ); + + assert_eq!( + span_data_1.gen_ai_operation_type.value(), + Some(&"handoff".to_owned()) + ); + + // Third span should have operation type mapped to "agent" fallback + assert_eq!( + span_data_2.gen_ai_operation_type.value(), + Some(&"agent".to_owned()) + ); + } + + #[test] + fn test_ai_operation_type_disabled_map() { + let json = r#" + { + "type": "transaction", + "transaction": "test-transaction", + "spans": [ + { + "op": "gen_ai.chat", + "description": "AI chat completion", + "data": {} + } + ] + } + "#; + + let mut event = Annotated::::from_json(json).unwrap(); + + let operation_type_map = AiOperationTypeMap { + version: 0, // Disabled version + operation_types: HashMap::from([( + Pattern::new("gen_ai.chat").unwrap(), + "chat".to_owned(), + )]), + }; + + normalize_event( + &mut event, + &NormalizationConfig { + ai_operation_type_map: Some(&operation_type_map), + ..NormalizationConfig::default() + }, + ); + + let span_data = get_value!(event.spans[0].data!); + + // Should not set operation type when map is disabled + assert_eq!(span_data.gen_ai_operation_type.value(), None); + } + + #[test] + fn test_ai_operation_type_empty_map() { + let json = r#" + { + "type": "transaction", + "transaction": "test-transaction", + "spans": [ + { + "op": "gen_ai.chat", + "description": "AI chat completion", + "data": {} + } + ] + } + "#; + + let mut event = Annotated::::from_json(json).unwrap(); + + let operation_type_map = AiOperationTypeMap { + version: 1, + operation_types: HashMap::new(), + }; + + normalize_event( + &mut event, + &NormalizationConfig { + ai_operation_type_map: Some(&operation_type_map), + ..NormalizationConfig::default() + }, + ); + + let span_data = get_value!(event.spans[0].data!); + + // Should not set operation type when map is empty + assert_eq!(span_data.gen_ai_operation_type.value(), None); + } + #[test] fn test_apple_high_device_class() { let mut event = Event { diff --git a/relay-event-normalization/src/normalize/mod.rs b/relay-event-normalization/src/normalize/mod.rs index fb223c49fa9..a8637234b28 100644 --- a/relay-event-normalization/src/normalize/mod.rs +++ b/relay-event-normalization/src/normalize/mod.rs @@ -300,7 +300,7 @@ pub struct ModelCostV2 { } /// A mapping of AI operation types from span.op to gen_ai.operation.type. /// -/// This struct uses a dictionary-based mapping structure with exact span operation keys +/// This struct uses a dictionary-based mapping structure with pattern-based span operation keys /// and corresponding AI operation type values. /// /// Example JSON: @@ -322,7 +322,7 @@ pub struct AiOperationTypeMap { /// The mappings of span.op => gen_ai.operation.type as a dictionary #[serde(default, skip_serializing_if = "HashMap::is_empty")] - pub operation_types: HashMap, + pub operation_types: HashMap, } impl AiOperationTypeMap { @@ -343,7 +343,22 @@ impl AiOperationTypeMap { if !self.is_enabled() { return None; } - self.operation_types.get(span_op).map(String::as_str) + + // try first direct match with span_op + if let Some(value) = self.operation_types.get(span_op) { + return Some(value.as_str()); + } + + // if there is not a direct match, try to find the match using a pattern + let operation_type = self.operation_types.iter().find_map(|(key, value)| { + if key.is_match(span_op) { + Some(value) + } else { + None + } + }); + + operation_type.map(String::as_str) } } #[cfg(test)] @@ -546,6 +561,61 @@ mod tests { assert!(!deserialized.is_enabled()); } + #[test] + fn test_ai_operation_type_map_serialization() { + // Test serialization and deserialization with patterns + let mut operation_types = HashMap::new(); + operation_types.insert( + Pattern::new("gen_ai.chat*").unwrap(), + "Inference".to_owned(), + ); + operation_types.insert( + Pattern::new("gen_ai.execute_tool").unwrap(), + "Tool".to_owned(), + ); + + let original = AiOperationTypeMap { + version: 1, + operation_types, + }; + + let json = serde_json::to_string(&original).unwrap(); + let deserialized: AiOperationTypeMap = serde_json::from_str(&json).unwrap(); + + assert!(deserialized.is_enabled()); + assert_eq!( + deserialized.get_operation_type("gen_ai.chat.completions"), + Some("Inference") + ); + assert_eq!( + deserialized.get_operation_type("gen_ai.execute_tool"), + Some("Tool") + ); + assert_eq!(deserialized.get_operation_type("unknown_op"), None); + } + + #[test] + fn test_ai_operation_type_map_pattern_matching() { + let mut operation_types = HashMap::new(); + operation_types.insert(Pattern::new("gen_ai.*").unwrap(), "default".to_owned()); + operation_types.insert(Pattern::new("gen_ai.chat").unwrap(), "chat".to_owned()); + + let map = AiOperationTypeMap { + version: 1, + operation_types, + }; + + let result = map.get_operation_type("gen_ai.chat"); + assert!(Some("chat") == result); + + let result = map.get_operation_type("gen_ai.chat.completions"); + assert!(Some("default") == result); + + assert_eq!(map.get_operation_type("gen_ai.other"), Some("default")); + + assert_eq!(map.get_operation_type("other.operation"), None); + } + #[test] fn test_merge_builtin_measurement_keys() { let foo = BuiltinMeasurementKey::new("foo", MetricUnit::Duration(DurationUnit::Hour)); diff --git a/relay-event-normalization/src/normalize/span/ai.rs b/relay-event-normalization/src/normalize/span/ai.rs index 1a2cf44e4fa..7041f7b2bba 100644 --- a/relay-event-normalization/src/normalize/span/ai.rs +++ b/relay-event-normalization/src/normalize/span/ai.rs @@ -1,5 +1,6 @@ //! AI cost calculation. +use crate::normalize::AiOperationTypeMap; use crate::{ModelCostV2, ModelCosts}; use relay_event_schema::protocol::{Event, Span, SpanData}; use relay_protocol::{Annotated, Getter, Value}; @@ -72,11 +73,7 @@ fn extract_ai_model_cost_data(model_cost: Option, data: &mut SpanDa } /// Maps AI-related measurements (legacy) to span data. -pub fn map_ai_measurements_to_data(span: &mut Span) { - if !is_ai_span(span) { - return; - }; - +fn map_ai_measurements_to_data(span: &mut Span) { let measurements = span.measurements.value(); let data = span.data.get_or_insert_with(SpanData::default); @@ -120,11 +117,7 @@ pub fn map_ai_measurements_to_data(span: &mut Span) { } /// Extract the additional data into the span -pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) { - if !is_ai_span(span) { - return; - } - +fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) { let duration = span .get_value("span.duration") .and_then(|v| v.as_f64()) @@ -161,22 +154,57 @@ pub fn extract_ai_data(span: &mut Span, ai_model_costs: &ModelCosts) { } } +/// Enrich the AI span data +pub fn enrich_ai_span_data( + span: &mut Span, + model_costs: Option<&ModelCosts>, + operation_type_map: Option<&AiOperationTypeMap>, +) { + if !is_ai_span(span) { + return; + } + + map_ai_measurements_to_data(span); + if let Some(model_costs) = model_costs { + extract_ai_data(span, model_costs); + } + if let Some(operation_type_map) = operation_type_map { + infer_ai_operation_type(span, operation_type_map); + } +} + /// Extract the ai data from all of an event's spans -pub fn enrich_ai_span_data(event: &mut Event, model_costs: Option<&ModelCosts>) { +pub fn enrich_ai_event_data( + event: &mut Event, + model_costs: Option<&ModelCosts>, + operation_type_map: Option<&AiOperationTypeMap>, +) { let spans = event.spans.value_mut().iter_mut().flatten(); let spans = spans.filter_map(|span| span.value_mut().as_mut()); for span in spans { - map_ai_measurements_to_data(span); - if let Some(model_costs) = model_costs { - extract_ai_data(span, model_costs); - } + enrich_ai_span_data(span, model_costs, operation_type_map); + } +} + +/// Infer AI operation type mapping to a span. +/// +/// This function maps span.op values to gen_ai.operation.type based on the provided +/// operation type map configuration. +fn infer_ai_operation_type(span: &mut Span, operation_type_map: &AiOperationTypeMap) { + let data = span.data.get_or_insert_with(SpanData::default); + + if let Some(op) = span.op.value() + && let Some(operation_type) = operation_type_map.get_operation_type(op) + { + data.gen_ai_operation_type + .set_value(Some(operation_type.to_owned())); } } /// Returns true if the span is an AI span. /// AI spans are spans with op starting with "ai." (legacy) or "gen_ai." (new). -pub fn is_ai_span(span: &Span) -> bool { +fn is_ai_span(span: &Span) -> bool { span.op .value() .is_some_and(|op| op.starts_with("ai.") || op.starts_with("gen_ai.")) diff --git a/relay-event-schema/src/protocol/span.rs b/relay-event-schema/src/protocol/span.rs index b1ead2ef9dd..4c072c50e79 100644 --- a/relay-event-schema/src/protocol/span.rs +++ b/relay-event-schema/src/protocol/span.rs @@ -641,6 +641,10 @@ pub struct SpanData { #[metastructure(field = "gen_ai.operation.name", pii = "maybe")] pub gen_ai_operation_name: Annotated, + /// The type of the operation being performed. + #[metastructure(field = "gen_ai.operation.type", pii = "maybe")] + pub gen_ai_operation_type: Annotated, + /// The client's browser name. #[metastructure(field = "browser.name")] pub browser_name: Annotated, @@ -1443,6 +1447,7 @@ mod tests { gen_ai_system: ~, gen_ai_tool_name: ~, gen_ai_operation_name: ~, + gen_ai_operation_type: ~, browser_name: ~, code_filepath: String( "task.py", diff --git a/relay-event-schema/src/protocol/span/convert.rs b/relay-event-schema/src/protocol/span/convert.rs index fd1bd3ebfff..bf8a9f5b35d 100644 --- a/relay-event-schema/src/protocol/span/convert.rs +++ b/relay-event-schema/src/protocol/span/convert.rs @@ -189,6 +189,7 @@ mod tests { gen_ai_system: ~, gen_ai_tool_name: ~, gen_ai_operation_name: ~, + gen_ai_operation_type: ~, browser_name: "Chrome", code_filepath: ~, code_lineno: ~, diff --git a/relay-pattern/src/lib.rs b/relay-pattern/src/lib.rs index f20e3919c7e..3864747152a 100644 --- a/relay-pattern/src/lib.rs +++ b/relay-pattern/src/lib.rs @@ -37,6 +37,7 @@ //! //! For untrusted user input it is highly recommended to limit the maximum complexity. +use std::borrow::Borrow; use std::fmt; use std::num::NonZeroUsize; @@ -143,6 +144,12 @@ impl PartialEq for Pattern { impl Eq for Pattern {} +impl Borrow for Pattern { + fn borrow(&self) -> &str { + &self.pattern + } +} + impl std::hash::Hash for Pattern { fn hash(&self, state: &mut H) { self.pattern.hash(state); diff --git a/relay-server/src/services/processor.rs b/relay-server/src/services/processor.rs index b3ba2fb2e75..e711dd736ae 100644 --- a/relay-server/src/services/processor.rs +++ b/relay-server/src/services/processor.rs @@ -1518,7 +1518,8 @@ impl EnvelopeProcessorService { .aggregator_config_for(MetricNamespace::Transactions); let global_config = self.inner.global_config.current(); - let ai_model_costs = global_config.ai_model_costs.clone().ok(); + let ai_model_costs = global_config.ai_model_costs.as_ref().ok(); + let ai_operation_type_map = global_config.ai_operation_type_map.as_ref().ok(); let http_span_allowed_hosts = global_config.options.http_span_allowed_hosts.as_slice(); let retention_days: i64 = project_info @@ -1590,7 +1591,8 @@ impl EnvelopeProcessorService { emit_event_errors: full_normalization, span_description_rules: project_info.config.span_description_rules.as_ref(), geoip_lookup: Some(&self.inner.geoip_lookup), - ai_model_costs: ai_model_costs.as_ref(), + ai_model_costs, + ai_operation_type_map, enable_trimming: true, measurements: Some(CombinedMeasurementsConfig::new( project_info.config().measurements.as_ref(), diff --git a/relay-server/src/services/processor/span/processing.rs b/relay-server/src/services/processor/span/processing.rs index a3eacae0385..98298a4c191 100644 --- a/relay-server/src/services/processor/span/processing.rs +++ b/relay-server/src/services/processor/span/processing.rs @@ -21,7 +21,8 @@ use relay_config::Config; use relay_dynamic_config::{ CombinedMetricExtractionConfig, ErrorBoundary, Feature, GlobalConfig, ProjectConfig, }; -use relay_event_normalization::span::ai::{extract_ai_data, map_ai_measurements_to_data}; +use relay_event_normalization::AiOperationTypeMap; +use relay_event_normalization::span::ai::enrich_ai_span_data; use relay_event_normalization::{ BorrowedSpanOpDefaults, ClientHints, CombinedMeasurementsConfig, FromUserAgentInfo, GeoIpLookup, MeasurementsConfig, ModelCosts, PerformanceScoreConfig, RawUserAgentInfo, @@ -435,6 +436,8 @@ struct NormalizeSpanConfig<'a> { measurements: Option>, /// Configuration for AI model cost calculation ai_model_costs: Option<&'a ModelCosts>, + /// Configuration to derive the `gen_ai.operation.type` field from other fields + ai_operation_type_map: Option<&'a AiOperationTypeMap>, /// The maximum length for names of custom measurements. /// /// Measurements with longer names are removed from the transaction event and replaced with a @@ -478,10 +481,8 @@ impl<'a> NormalizeSpanConfig<'a> { project_config.measurements.as_ref(), global_config.measurements.as_ref(), )), - ai_model_costs: match &global_config.ai_model_costs { - ErrorBoundary::Err(_) => None, - ErrorBoundary::Ok(costs) => Some(costs), - }, + ai_model_costs: global_config.ai_model_costs.as_ref().ok(), + ai_operation_type_map: global_config.ai_operation_type_map.as_ref().ok(), max_name_and_unit_len: aggregator_config .max_name_length .saturating_sub(MeasurementsConfig::MEASUREMENT_MRI_OVERHEAD), @@ -544,6 +545,7 @@ fn normalize( performance_score, measurements, ai_model_costs, + ai_operation_type_map, max_name_and_unit_len, tx_name_rules, user_agent, @@ -651,10 +653,7 @@ fn normalize( normalize_performance_score(span, performance_score); - map_ai_measurements_to_data(span); - if let Some(model_costs_config) = ai_model_costs { - extract_ai_data(span, model_costs_config); - } + enrich_ai_span_data(span, ai_model_costs, ai_operation_type_map); tag_extraction::extract_measurements(span, is_mobile); @@ -1253,6 +1252,7 @@ mod tests { performance_score: None, measurements: None, ai_model_costs: None, + ai_operation_type_map: None, max_name_and_unit_len: 200, tx_name_rules: &[], user_agent: None,