fix(llm-monitoring): Fix alerts, remove hacky ai_total_tokens_used fu…

…nction (getsentry#71766) Make alerts go to the spansMetrics dataset, and remove the old hack where `sum` didn't work on counters. Also add an alias to avoid hardcoding MRIs.
uyefe · May 30, 2024 · 0dedb6c · 0dedb6c
1 parent cb56da9
commit 0dedb6c
Show file tree

Hide file tree

Showing 7 changed files with 37 additions and 35 deletions.
diff --git a/static/app/views/alerts/rules/metric/triggers/chart/index.tsx b/static/app/views/alerts/rules/metric/triggers/chart/index.tsx
@@ -262,6 +262,7 @@ class TriggersChart extends PureComponent<Props, State> {
     });
 
     let queryDataset = queryExtras.dataset as undefined | DiscoverDatasets;
+    const queryOverride = (queryExtras.query as string | undefined) ?? query;
 
     if (shouldUseErrorsDiscoverDataset(query, dataset)) {
       queryDataset = DiscoverDatasets.ERRORS;
@@ -271,7 +272,7 @@ class TriggersChart extends PureComponent<Props, State> {
       const totalCount = await fetchTotalCount(api, organization.slug, {
         field: [],
         project: projects.map(({id}) => id),
-        query,
+        query: queryOverride,
         statsPeriod,
         environment: environment ? [environment] : [],
         dataset: queryDataset,

diff --git a/static/app/views/alerts/rules/metric/utils/getMetricDatasetQueryExtras.tsx b/static/app/views/alerts/rules/metric/utils/getMetricDatasetQueryExtras.tsx
@@ -26,10 +26,14 @@ export function getMetricDatasetQueryExtras({
   const disableMetricDataset =
     decodeScalar(location?.query?.disableMetricDataset) === 'true';
 
-  const queryExtras: Record<string, string> =
-    hasMetricDataset && !disableMetricDataset
-      ? {dataset: getMEPAlertsDataset(dataset, newAlertOrQuery)}
-      : {};
+  const queryExtras: Record<string, string> = {};
+  if (hasMetricDataset && !disableMetricDataset) {
+    queryExtras.dataset = getMEPAlertsDataset(dataset, newAlertOrQuery);
+  }
+  if (location?.query?.aggregate?.includes('ai.total')) {
+    queryExtras.dataset = 'spansMetrics';
+    queryExtras.query = '';
+  }
 
   if (useOnDemandMetrics) {
     queryExtras.useOnDemandMetrics = 'true';

diff --git a/static/app/views/alerts/wizard/options.tsx b/static/app/views/alerts/wizard/options.tsx
@@ -192,12 +192,12 @@ export const AlertWizardRuleTemplates: Record<
     eventTypes: EventTypes.TRANSACTION,
   },
   llm_tokens: {
-    aggregate: 'sum(c:spans/ai.total_tokens.used@none)',
+    aggregate: 'sum(ai.total_tokens.used)',
     dataset: Dataset.GENERIC_METRICS,
     eventTypes: EventTypes.TRANSACTION,
   },
   llm_cost: {
-    aggregate: 'sum(c:spans/ai.total_cost@usd)',
+    aggregate: 'sum(ai.total_cost)',
     dataset: Dataset.GENERIC_METRICS,
     eventTypes: EventTypes.TRANSACTION,
   },

diff --git a/static/app/views/llmMonitoring/llmMonitoringCharts.tsx b/static/app/views/llmMonitoring/llmMonitoringCharts.tsx
@@ -10,7 +10,7 @@ interface TotalTokensUsedChartProps {
 }
 
 export function TotalTokensUsedChart({groupId}: TotalTokensUsedChartProps) {
-  const aggregate = 'ai_total_tokens_used()';
+  const aggregate = 'sum(ai.total_tokens.used)';
 
   let query = 'span.category:"ai"';
   if (groupId) {

diff --git a/static/app/views/llmMonitoring/llmMonitoringDetailsPage.tsx b/static/app/views/llmMonitoring/llmMonitoringDetailsPage.tsx
@@ -76,10 +76,7 @@ export function LLMMonitoringPage({params}: Props) {
         'span.category': 'ai',
         'span.ai.pipeline.group': groupId,
       }),
-      fields: [
-        'ai_total_tokens_used()',
-        'ai_total_tokens_used(c:spans/ai.total_cost@usd)',
-      ],
+      fields: ['sum(ai.total_tokens.used)', 'sum(ai.total_cost)'],
       enabled: Boolean(groupId),
     },
     'api.ai-pipelines.view'
@@ -117,16 +114,14 @@ export function LLMMonitoringPage({params}: Props) {
                   <MetricsRibbon>
                     <MetricReadout
                       title={t('Total Tokens Used')}
-                      value={tokenUsedMetric['ai_total_tokens_used()']}
+                      value={tokenUsedMetric['sum(ai.total_tokens.used)']}
                       unit={'count'}
                       isLoading={isTotalTokenDataLoading}
                     />
 
                     <MetricReadout
                       title={t('Total Cost')}
-                      value={
-                        tokenUsedMetric['ai_total_tokens_used(c:spans/ai.total_cost@usd)']
-                      }
+                      value={tokenUsedMetric['sum(ai.total_cost)']}
                       unit={CurrencyUnit.USD}
                       isLoading={isTotalTokenDataLoading}
                     />

diff --git a/static/app/views/llmMonitoring/pipelinesTable.tsx b/static/app/views/llmMonitoring/pipelinesTable.tsx
@@ -39,16 +39,16 @@ type Row = Pick<
   | 'spm()'
   | 'avg(span.duration)'
   | 'sum(span.duration)'
-  | 'ai_total_tokens_used()'
-  | 'ai_total_tokens_used(c:spans/ai.total_cost@usd)'
+  | 'sum(ai.total_tokens.used)'
+  | 'sum(ai.total_cost)'
 >;
 
 type Column = GridColumnHeader<
   | 'span.description'
   | 'spm()'
   | 'avg(span.duration)'
-  | 'ai_total_tokens_used()'
-  | 'ai_total_tokens_used(c:spans/ai.total_cost@usd)'
+  | 'sum(ai.total_tokens.used)'
+  | 'sum(ai.total_cost)'
 >;
 
 const COLUMN_ORDER: Column[] = [
@@ -58,12 +58,12 @@ const COLUMN_ORDER: Column[] = [
     width: COL_WIDTH_UNDEFINED,
   },
   {
-    key: 'ai_total_tokens_used()',
+    key: 'sum(ai.total_tokens.used)',
     name: t('Total tokens used'),
     width: 180,
   },
   {
-    key: 'ai_total_tokens_used(c:spans/ai.total_cost@usd)',
+    key: 'sum(ai.total_cost)',
     name: t('Total cost'),
     width: 180,
   },
@@ -79,7 +79,7 @@ const COLUMN_ORDER: Column[] = [
   },
 ];
 
-const SORTABLE_FIELDS = ['ai_total_tokens_used()', 'avg(span.duration)', 'spm()'];
+const SORTABLE_FIELDS = ['sum(ai.total_tokens.used)', 'avg(span.duration)', 'spm()'];
 
 type ValidSort = Sort & {
   field: 'spm()' | 'avg(span.duration)';
@@ -131,8 +131,8 @@ export function PipelinesTable() {
       ),
       fields: [
         'span.ai.pipeline.group',
-        'ai_total_tokens_used()',
-        'ai_total_tokens_used(c:spans/ai.total_cost@usd)',
+        'sum(ai.total_tokens.used)',
+        'sum(ai.total_cost)',
       ],
     },
     'api.performance.ai-analytics.token-usage-chart'
@@ -141,17 +141,17 @@ export function PipelinesTable() {
   const rows: Row[] = (data as Row[]).map(baseRow => {
     const row: Row = {
       ...baseRow,
-      'ai_total_tokens_used()': 0,
-      'ai_total_tokens_used(c:spans/ai.total_cost@usd)': 0,
+      'sum(ai.total_tokens.used)': 0,
+      'sum(ai.total_cost)': 0,
     };
     if (!tokensUsedLoading) {
       const tokenUsedDataPoint = tokensUsedData.find(
         tokenRow => tokenRow['span.ai.pipeline.group'] === row['span.group']
       );
       if (tokenUsedDataPoint) {
-        row['ai_total_tokens_used()'] = tokenUsedDataPoint['ai_total_tokens_used()'];
-        row['ai_total_tokens_used(c:spans/ai.total_cost@usd)'] =
-          tokenUsedDataPoint['ai_total_tokens_used(c:spans/ai.total_cost@usd)'];
+        row['sum(ai.total_tokens.used)'] =
+          tokenUsedDataPoint['sum(ai.total_tokens.used)'];
+        row['sum(ai.total_cost)'] = tokenUsedDataPoint['sum(ai.total_cost)'];
       }
     }
     return row;
@@ -246,8 +246,8 @@ function renderBodyCell(
       </Link>
     );
   }
-  if (column.key === 'ai_total_tokens_used(c:spans/ai.total_cost@usd)') {
-    const cost = row['ai_total_tokens_used(c:spans/ai.total_cost@usd)'];
+  if (column.key === 'sum(ai.total_cost)') {
+    const cost = row[column.key];
     if (cost) {
       return <span>US ${cost.toFixed(3)}</span>;
     }

diff --git a/static/app/views/starfish/types.tsx b/static/app/views/starfish/types.tsx
@@ -42,6 +42,8 @@ export enum SpanMetricsField {
   HTTP_DECODED_RESPONSE_CONTENT_LENGTH = 'http.decoded_response_content_length',
   HTTP_RESPONSE_TRANSFER_SIZE = 'http.response_transfer_size',
   FILE_EXTENSION = 'file_extension',
+  AI_TOTAL_TOKENS_USED = 'ai.total_tokens.used',
+  AI_TOTAL_COST = 'ai.total_cost',
   OS_NAME = 'os.name',
   APP_START_TYPE = 'app_start_type',
   DEVICE_CLASS = 'device.class',
@@ -51,6 +53,8 @@ export enum SpanMetricsField {
 }
 
 export type SpanNumberFields =
+  | SpanMetricsField.AI_TOTAL_COST
+  | SpanMetricsField.AI_TOTAL_TOKENS_USED
   | SpanMetricsField.SPAN_SELF_TIME
   | SpanMetricsField.SPAN_DURATION
   | SpanMetricsField.HTTP_DECODED_RESPONSE_CONTENT_LENGTH
@@ -112,7 +116,7 @@ export const SPAN_FUNCTIONS = [
   'http_error_count',
   'cache_hit_rate',
   'cache_miss_rate',
-  'ai_total_tokens_used',
+  'sum',
 ] as const;
 
 const BREAKPOINT_CONDITIONS = ['less', 'greater'] as const;
@@ -142,8 +146,6 @@ export type SpanMetricsResponse = {
   'http_response_rate(3)': number;
   'http_response_rate(4)': number;
   'http_response_rate(5)': number;
-} & {
-  'ai_total_tokens_used(c:spans/ai.total_cost@usd)': number;
 } & {
   ['project']: string;
   ['project.id']: number;