feat(ai): Add missing gen_ai fields (#5246)

Adds a couple of missing gen_ai fields, sentry convetions PR will be
done as a follow up.

Also marked all of the fields that report token counts as `pii = maybe`
just in case we change scrubbing algorithm, we don't want to filter out
those values.

Closes [TET-1272: Add fields containing `token` to the scrubbing
exceptions in
relay](https://linear.app/getsentry/issue/TET-1272/add-fields-containing-token-to-the-scrubbing-exceptions-in-relay)

Author: vgrozdanic

relay-event-schema/src/protocol/span.rs

@@ -468,7 +468,7 @@ pub struct SpanData {
     pub app_start_type: Annotated<Value>,
 
     /// The maximum number of tokens that should be used by an LLM call.
-    #[metastructure(field = "gen_ai.request.max_tokens")]
+    #[metastructure(field = "gen_ai.request.max_tokens", pii = "maybe")]
     pub gen_ai_request_max_tokens: Annotated<Value>,
 
     /// Name of the AI pipeline or chain being executed.
@@ -478,36 +478,61 @@ pub struct SpanData {
     /// The total tokens that were used by an LLM call
     #[metastructure(
         field = "gen_ai.usage.total_tokens",
-        legacy_alias = "ai.total_tokens.used"
+        legacy_alias = "ai.total_tokens.used",
+        pii = "maybe"
     )]
     pub gen_ai_usage_total_tokens: Annotated<Value>,
 
     /// The input tokens used by an LLM call (usually cheaper than output tokens)
     #[metastructure(
         field = "gen_ai.usage.input_tokens",
         legacy_alias = "ai.prompt_tokens.used",
-        legacy_alias = "gen_ai.usage.prompt_tokens"
+        legacy_alias = "gen_ai.usage.prompt_tokens",
+        pii = "maybe"
     )]
     pub gen_ai_usage_input_tokens: Annotated<Value>,
 
     /// The input tokens used by an LLM call that were cached
     /// (cheaper and faster than non-cached input tokens)
-    #[metastructure(field = "gen_ai.usage.input_tokens.cached")]
+    #[metastructure(field = "gen_ai.usage.input_tokens.cached", pii = "maybe")]
     pub gen_ai_usage_input_tokens_cached: Annotated<Value>,
 
+    /// The input tokens written to cache during an LLM call
+    #[metastructure(field = "gen_ai.usage.input_tokens.cache_write", pii = "maybe")]
+    pub gen_ai_usage_input_tokens_cache_write: Annotated<Value>,
+
+    /// The input tokens that missed the cache (DeepSeek provider)
+    #[metastructure(field = "gen_ai.usage.input_tokens.cache_miss", pii = "maybe")]
+    pub gen_ai_usage_input_tokens_cache_miss: Annotated<Value>,
+
     /// The output tokens used by an LLM call (the ones the LLM actually generated)
     #[metastructure(
         field = "gen_ai.usage.output_tokens",
         legacy_alias = "ai.completion_tokens.used",
-        legacy_alias = "gen_ai.usage.completion_tokens"
+        legacy_alias = "gen_ai.usage.completion_tokens",
+        pii = "maybe"
     )]
     pub gen_ai_usage_output_tokens: Annotated<Value>,
 
     /// The output tokens used to represent the model's internal thought
     /// process while generating a response
-    #[metastructure(field = "gen_ai.usage.output_tokens.reasoning")]
+    #[metastructure(field = "gen_ai.usage.output_tokens.reasoning", pii = "maybe")]
     pub gen_ai_usage_output_tokens_reasoning: Annotated<Value>,
 
+    /// The output tokens for accepted predictions (OpenAI provider)
+    #[metastructure(
+        field = "gen_ai.usage.output_tokens.prediction_accepted",
+        pii = "maybe"
+    )]
+    pub gen_ai_usage_output_tokens_prediction_accepted: Annotated<Value>,
+
+    /// The output tokens for rejected predictions (OpenAI provider)
+    #[metastructure(
+        field = "gen_ai.usage.output_tokens.prediction_rejected",
+        pii = "maybe"
+    )]
+    pub gen_ai_usage_output_tokens_prediction_rejected: Annotated<Value>,
+
     // Exact model used to generate the response (e.g. gpt-4o-mini-2024-07-18)
     #[metastructure(field = "gen_ai.response.model")]
     pub gen_ai_response_model: Annotated<Value>,
@@ -521,15 +546,15 @@ pub struct SpanData {
     pub gen_ai_usage_total_cost: Annotated<Value>,
 
     /// The total cost for the tokens used (duplicate field for migration)
-    #[metastructure(field = "gen_ai.cost.total_tokens")]
+    #[metastructure(field = "gen_ai.cost.total_tokens", pii = "maybe")]
     pub gen_ai_cost_total_tokens: Annotated<Value>,
 
     /// The cost for input tokens used
-    #[metastructure(field = "gen_ai.cost.input_tokens")]
+    #[metastructure(field = "gen_ai.cost.input_tokens", pii = "maybe")]
     pub gen_ai_cost_input_tokens: Annotated<Value>,
 
     /// The cost for output tokens used
-    #[metastructure(field = "gen_ai.cost.output_tokens")]
+    #[metastructure(field = "gen_ai.cost.output_tokens", pii = "maybe")]
     pub gen_ai_cost_output_tokens: Annotated<Value>,
 
     /// Prompt passed to LLM (Vercel AI SDK)
@@ -587,7 +612,7 @@ pub struct SpanData {
     pub gen_ai_response_streaming: Annotated<Value>,
 
     ///  Total output tokens per seconds throughput
-    #[metastructure(field = "gen_ai.response.tokens_per_second")]
+    #[metastructure(field = "gen_ai.response.tokens_per_second", pii = "maybe")]
     pub gen_ai_response_tokens_per_second: Annotated<Value>,
 
     /// The available tools for a request to an LLM
@@ -1444,16 +1469,20 @@ mod tests {
             .unwrap()
             .into_value()
             .unwrap();
-        insta::assert_debug_snapshot!(data, @r###"
+        insta::assert_debug_snapshot!(data, @r#"
         SpanData {
             app_start_type: ~,
             gen_ai_request_max_tokens: ~,
             gen_ai_pipeline_name: ~,
             gen_ai_usage_total_tokens: ~,
             gen_ai_usage_input_tokens: ~,
             gen_ai_usage_input_tokens_cached: ~,
+            gen_ai_usage_input_tokens_cache_write: ~,
+            gen_ai_usage_input_tokens_cache_miss: ~,
             gen_ai_usage_output_tokens: ~,
             gen_ai_usage_output_tokens_reasoning: ~,
+            gen_ai_usage_output_tokens_prediction_accepted: ~,
+            gen_ai_usage_output_tokens_prediction_rejected: ~,
             gen_ai_response_model: ~,
             gen_ai_request_model: ~,
             gen_ai_usage_total_cost: ~,
@@ -1581,7 +1610,7 @@ mod tests {
                 ),
             },
         }
-        "###);
+        "#);
 
         assert_eq!(data.get_value("foo"), Some(Val::U64(2)));
         assert_eq!(data.get_value("bar"), Some(Val::String("3")));

relay-event-schema/src/protocol/span/convert.rs

@@ -160,8 +160,12 @@ mod tests {
                 gen_ai_usage_total_tokens: ~,
                 gen_ai_usage_input_tokens: ~,
                 gen_ai_usage_input_tokens_cached: ~,
+                gen_ai_usage_input_tokens_cache_write: ~,
+                gen_ai_usage_input_tokens_cache_miss: ~,
                 gen_ai_usage_output_tokens: ~,
                 gen_ai_usage_output_tokens_reasoning: ~,
+                gen_ai_usage_output_tokens_prediction_accepted: ~,
+                gen_ai_usage_output_tokens_prediction_rejected: ~,
                 gen_ai_response_model: ~,
                 gen_ai_request_model: ~,
                 gen_ai_usage_total_cost: ~,