feat(ai-monitoring): Add tokens used to pipeline table, fix number of pipelines graph (#69601)

static/app/views/aiMonitoring/PipelinesTable.tsx

@@ -1,8 +1,10 @@
 import {browserHistory} from 'react-router';
 import type {Location} from 'history';
 
-import type {GridColumnHeader} from 'sentry/components/gridEditable';
-import GridEditable, {COL_WIDTH_UNDEFINED} from 'sentry/components/gridEditable';
+import GridEditable, {
+  COL_WIDTH_UNDEFINED,
+  type GridColumnHeader,
+} from 'sentry/components/gridEditable';
 import Link from 'sentry/components/links/link';
 import type {CursorHandler} from 'sentry/components/pagination';
 import Pagination from 'sentry/components/pagination';
@@ -32,17 +34,17 @@ type Row = Pick<
   | 'spm()'
   | 'avg(span.duration)'
   | 'sum(span.duration)'
-  | 'time_spent_percentage()'
+  | 'ai_total_tokens_used()'
 >;
 
 type Column = GridColumnHeader<
-  'span.description' | 'spm()' | 'avg(span.duration)' | 'time_spent_percentage()'
+  'span.description' | 'spm()' | 'avg(span.duration)' | 'ai_total_tokens_used()'
 >;
 
 const COLUMN_ORDER: Column[] = [
   {
     key: 'span.description',
-    name: t('AI Pipeline name'),
+    name: t('AI Pipeline Name'),
     width: COL_WIDTH_UNDEFINED,
   },
   {
@@ -51,21 +53,21 @@ const COLUMN_ORDER: Column[] = [
     width: COL_WIDTH_UNDEFINED,
   },
   {
-    key: `avg(span.duration)`,
-    name: DataTitles.avg,
-    width: COL_WIDTH_UNDEFINED,
+    key: 'ai_total_tokens_used()',
+    name: t('Total tokens used'),
+    width: 180,
   },
   {
-    key: 'time_spent_percentage()',
-    name: DataTitles.timeSpent,
+    key: `avg(span.duration)`,
+    name: DataTitles.avg,
     width: COL_WIDTH_UNDEFINED,
   },
 ];
 
-const SORTABLE_FIELDS = ['avg(span.duration)', 'spm()', 'time_spent_percentage()'];
+const SORTABLE_FIELDS = ['avg(span.duration)', 'spm()'];
 
 type ValidSort = Sort & {
-  field: 'spm()' | 'avg(span.duration)' | 'time_spent_percentage()';
+  field: 'spm()' | 'avg(span.duration)';
 };
 
 export function isAValidSort(sort: Sort): sort is ValidSort {
@@ -80,7 +82,7 @@ export function PipelinesTable() {
 
   let sort = decodeSorts(sortField).filter(isAValidSort)[0];
   if (!sort) {
-    sort = {field: 'time_spent_percentage()', kind: 'desc'};
+    sort = {field: 'spm()', kind: 'desc'};
   }
   const {data, isLoading, meta, pageLinks, error} = useSpanMetrics({
     search: new MutableSearch('span.category:ai.pipeline'),
@@ -91,7 +93,7 @@ export function PipelinesTable() {
       'spm()',
       'avg(span.duration)',
       'sum(span.duration)',
-      'time_spent_percentage()',
+      'ai_total_tokens_used()', // this is zero initially and overwritten below.
     ],
     sorts: [sort],
     limit: 25,
@@ -99,6 +101,26 @@ export function PipelinesTable() {
     referrer: 'api.ai-pipelines.view',
   });
 
+  const {
+    data: tokensUsedData,
+    isLoading: tokensUsedLoading,
+    error: tokensUsedError,
+  } = useSpanMetrics({
+    search: new MutableSearch(
+      `span.ai.pipeline.group:[${(data as Row[])?.map(x => x['span.group']).join(',')}] span.category:ai`
+    ),
+    fields: ['span.ai.pipeline.group', 'ai_total_tokens_used()'],
+  });
+  if (!tokensUsedLoading) {
+    for (const tokenUsedRow of tokensUsedData) {
+      const groupId = tokenUsedRow['span.ai.pipeline.group'];
+      const tokensUsed = tokenUsedRow['ai_total_tokens_used()'];
+      data
+        .filter(x => x['span.group'] === groupId)
+        .forEach(x => (x['ai_total_tokens_used()'] = tokensUsed));
+    }
+  }
+
   const handleCursor: CursorHandler = (newCursor, pathname, query) => {
     browserHistory.push({
       pathname,
@@ -114,7 +136,7 @@ export function PipelinesTable() {
     >
       <GridEditable
         isLoading={isLoading}
-        error={error}
+        error={error ?? tokensUsedError}
         data={data}
         columnOrder={COLUMN_ORDER}
         columnSortBy={[

static/app/views/aiMonitoring/aiMonitoringCharts.tsx

@@ -8,8 +8,15 @@ import {useMetricsQuery} from 'sentry/utils/metrics/useMetricsQuery';
 import usePageFilters from 'sentry/utils/usePageFilters';
 import {MetricChartContainer} from 'sentry/views/dashboards/metrics/chart';
 
-export function TotalTokensUsedChart() {
+interface TotalTokensUsedChartProps {
+  groupId?: string;
+}
+export function TotalTokensUsedChart({groupId}: TotalTokensUsedChartProps) {
   const {selection, isReady: isGlobalSelectionReady} = usePageFilters();
+  let query = 'span.category:"ai"';
+  if (groupId) {
+    query = `${query} span.ai.pipeline.group:"${groupId}"`;
+  }
   const {
     data: timeseriesData,
     isLoading,
@@ -21,7 +28,7 @@ export function TotalTokensUsedChart() {
         name: 'total',
         mri: `c:spans/ai.total_tokens.used@none`,
         op: 'sum',
-        // TODO this double counts the (e.g.) langchain and openai token usage
+        query,
       },
     ],
     selection,
@@ -75,7 +82,7 @@ export function NumberOfPipelinesChart({groupId}: NumberOfPipelinesChartProps) {
     [
       {
         name: 'number',
-        mri: `d:spans/exclusive_time@millisecond`,
+        mri: `d:spans/exclusive_time_light@millisecond`,
         op: 'count',
         query,
       },
@@ -130,7 +137,7 @@ export function PipelineDurationChart({groupId}: PipelineDurationChartProps) {
   } = useMetricsQuery(
     [
       {
-        name: 'a',
+        name: 'duration',
         mri: `d:spans/duration@millisecond`,
         op: 'avg',
         query,
@@ -164,7 +171,7 @@ export function PipelineDurationChart({groupId}: PipelineDurationChartProps) {
         metricQueries={[
           {
             name: 'mql',
-            formula: '$a',
+            formula: '$duration',
           },
         ]}
         displayType={MetricDisplayType.AREA}

static/app/views/aiMonitoring/aiMonitoringDetailsPage.tsx

@@ -19,6 +19,7 @@ import useOrganization from 'sentry/utils/useOrganization';
 import {
   NumberOfPipelinesChart,
   PipelineDurationChart,
+  TotalTokensUsedChart,
 } from 'sentry/views/aiMonitoring/aiMonitoringCharts';
 import {PipelineSpansTable} from 'sentry/views/aiMonitoring/pipelineSpansTable';
 import {MetricReadout} from 'sentry/views/performance/metricReadout';
@@ -146,12 +147,15 @@ export default function AiMonitoringPage({params}: Props) {
                         </MetricsRibbon>
                       </SpaceBetweenWrap>
                     </ModuleLayout.Full>
-                    <ModuleLayout.Half>
+                    <ModuleLayout.Third>
+                      <TotalTokensUsedChart groupId={groupId} />
+                    </ModuleLayout.Third>
+                    <ModuleLayout.Third>
                       <NumberOfPipelinesChart groupId={groupId} />
-                    </ModuleLayout.Half>
-                    <ModuleLayout.Half>
+                    </ModuleLayout.Third>
+                    <ModuleLayout.Third>
                       <PipelineDurationChart groupId={groupId} />
-                    </ModuleLayout.Half>
+                    </ModuleLayout.Third>
                     <ModuleLayout.Full>
                       <PipelineSpansTable groupId={groupId} />
                     </ModuleLayout.Full>

static/app/views/starfish/types.tsx

@@ -58,7 +58,8 @@ export type SpanStringFields =
   | 'transaction.method'
   | 'release'
   | 'os.name'
-  | 'span.status_code';
+  | 'span.status_code'
+  | 'span.ai.pipeline.group';
 
 export type SpanMetricsQueryFilters = {
   [Field in SpanStringFields]?: string;