feat(llm-monitoring): Backend changes to alias the AI MRIs (#71770)

Also allows `sum` on counters, and the AI metrics counters in
particular.

src/sentry/search/events/constants.py

@@ -335,6 +335,8 @@ SPAN_METRICS_MAP = {
     "user": "s:spans/user@none",
     "span.self_time": "d:spans/exclusive_time@millisecond",
     "span.duration": "d:spans/duration@millisecond",
+    "ai.total_tokens.used": "c:spans/ai.total_tokens.used@none",
+    "ai.total_cost": "c:spans/ai.total_cost@usd",
     "http.response_content_length": "d:spans/http.response_content_length@byte",
     "http.decoded_response_content_length": "d:spans/http.decoded_response_content_length@byte",
     "http.response_transfer_size": "d:spans/http.response_transfer_size@byte",
@@ -369,6 +371,9 @@ SPAN_METRIC_DURATION_COLUMNS = {
     for key, value in SPAN_METRICS_MAP.items()
     if value.endswith("@millisecond") or value.endswith("@second")
 }
+SPAN_METRIC_SUMMABLE_COLUMNS = SPAN_METRIC_DURATION_COLUMNS.union(
+    {"ai.total_tokens.used", "ai.total_cost"}
+)
 SPAN_METRIC_COUNT_COLUMNS = {
     key
     for key, value in SPAN_METRICS_MAP.items()

src/sentry/search/events/datasets/spans_metrics.py

@@ -138,7 +138,7 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                             "span.self_time",
                             fields.MetricArg(
                                 "column",
-                                allowed_columns=constants.SPAN_METRIC_DURATION_COLUMNS,
+                                allowed_columns=constants.SPAN_METRIC_SUMMABLE_COLUMNS,
                                 allow_custom_measurements=False,
                             ),
                         ),
@@ -152,6 +152,14 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                         ],
                         alias,
                     ),
+                    snql_counter=lambda args, alias: Function(
+                        "sumIf",
+                        [
+                            Column("value"),
+                            Function("equals", [Column("metric_id"), args["metric_id"]]),
+                        ],
+                        alias,
+                    ),
                     default_result_type="duration",
                 ),
                 fields.MetricsFunction(
@@ -324,29 +332,6 @@ class SpansMetricsDatasetConfig(DatasetConfig):
                     snql_distribution=self._resolve_time_spent_percentage,
                     default_result_type="percentage",
                 ),
-                fields.MetricsFunction(
-                    "ai_total_tokens_used",
-                    optional_args=[
-                        fields.with_default(
-                            "c:spans/ai.total_tokens.used@none",
-                            fields.MetricArg(
-                                "column",
-                                allowed_columns=["c:spans/ai.total_tokens.used@none"],
-                                allow_custom_measurements=False,
-                                allow_mri=True,
-                            ),
-                        ),
-                    ],
-                    calculated_args=[resolve_metric_id],
-                    snql_counter=lambda args, alias: Function(
-                        "sumIf",
-                        [
-                            Column("value"),
-                            Function("equals", [Column("metric_id"), args["metric_id"]]),
-                        ],
-                        alias,
-                    ),
-                ),
                 fields.MetricsFunction(
                     "http_response_rate",
                     required_args=[
@@ -1353,7 +1338,7 @@ class SpansMetricsLayerDatasetConfig(DatasetConfig):
                             "span.self_time",
                             fields.MetricArg(
                                 "column",
-                                allowed_columns=constants.SPAN_METRIC_DURATION_COLUMNS,
+                                allowed_columns=constants.SPAN_METRIC_SUMMABLE_COLUMNS,
                                 allow_custom_measurements=False,
                             ),
                         ),

src/sentry/sentry_metrics/querying/samples_list.py

@@ -827,6 +827,7 @@ class SpansMeasurementsSamplesListExecutor(SpansSamplesListExecutor):
         SpanMRI.DECODED_RESPONSE_CONTENT_LENGTH.value: "http.decoded_response_content_length",
         SpanMRI.RESPONSE_TRANSFER_SIZE.value: "http.response_transfer_size",
         SpanMRI.AI_TOTAL_TOKENS.value: "ai_total_tokens_used",
+        SpanMRI.AI_TOTAL_COST.value: "ai_total_cost",
         SpanMRI.CACHE_ITEM_SIZE.value: "cache.item_size",
         SpanMRI.MOBILE_SLOW_FRAMES.value: "frames.slow",
         SpanMRI.MOBILE_FROZEN_FRAMES.value: "frames.frozen",

src/sentry/snuba/metrics/naming_layer/mri.py

@@ -167,7 +167,8 @@ class SpanMRI(Enum):
     SELF_TIME_LIGHT = "d:spans/exclusive_time_light@millisecond"
 
     # Measurement-based metrics
-    AI_TOTAL_TOKENS = "c:spans/ai.total_tokens@none"
+    AI_TOTAL_TOKENS = "c:spans/ai.total_tokens.used@none"
+    AI_TOTAL_COST = "c:spans/ai.total_cost@usd"
     CACHE_ITEM_SIZE = "d:spans/cache.item_size@byte"
     DECODED_RESPONSE_CONTENT_LENGTH = "d:spans/http.decoded_response_content_length@byte"
     MESSAGE_RECEIVE_LATENCY = "g:spans/messaging.message.receive.latency@millisecond"