From 2124b3635945df2e749dd0a2a3adc278b21a7137 Mon Sep 17 00:00:00 2001
From: colin-sentry <161344340+colin-sentry@users.noreply.github.com>
Date: Thu, 4 Apr 2024 16:32:49 -0400
Subject: [PATCH] Graph logic for AI analytics page (#68296)

Adds time series graphs for prompt tokens and completion tokens
---
 static/app/routes.tsx                         |  4 +-
 .../{ai-analytics => aiAnalytics}/filters.tsx |  0
 .../{ai-analytics => aiAnalytics}/index.tsx   | 12 +++-
 .../{ai-analytics => aiAnalytics}/landing.tsx |  0
 .../app/views/aiAnalytics/tokenUsageChart.tsx | 64 +++++++++++++++++++
 5 files changed, 77 insertions(+), 3 deletions(-)
 rename static/app/views/{ai-analytics => aiAnalytics}/filters.tsx (100%)
 rename static/app/views/{ai-analytics => aiAnalytics}/index.tsx (81%)
 rename static/app/views/{ai-analytics => aiAnalytics}/landing.tsx (100%)
 create mode 100644 static/app/views/aiAnalytics/tokenUsageChart.tsx

diff --git a/static/app/routes.tsx b/static/app/routes.tsx
index f3dd5dc4fc4b61..a5aa6ab2ef6b4d 100644
--- a/static/app/routes.tsx
+++ b/static/app/routes.tsx
@@ -1493,10 +1493,10 @@ function buildRoutes() {
   const aiAnalyticsRoutes = (
     <Route
       path="/ai-analytics/"
-      component={make(() => import('sentry/views/ai-analytics'))}
+      component={make(() => import('sentry/views/aiAnalytics'))}
       withOrgPath
     >
-      <IndexRoute component={make(() => import('sentry/views/ai-analytics/landing'))} />
+      <IndexRoute component={make(() => import('sentry/views/aiAnalytics/landing'))} />
     </Route>
   );
 
diff --git a/static/app/views/ai-analytics/filters.tsx b/static/app/views/aiAnalytics/filters.tsx
similarity index 100%
rename from static/app/views/ai-analytics/filters.tsx
rename to static/app/views/aiAnalytics/filters.tsx
diff --git a/static/app/views/ai-analytics/index.tsx b/static/app/views/aiAnalytics/index.tsx
similarity index 81%
rename from static/app/views/ai-analytics/index.tsx
rename to static/app/views/aiAnalytics/index.tsx
index 114d8259755336..dce6c51ad05921 100644
--- a/static/app/views/ai-analytics/index.tsx
+++ b/static/app/views/aiAnalytics/index.tsx
@@ -9,7 +9,8 @@ import PageFiltersContainer from 'sentry/components/organizations/pageFilters/co
 import {t} from 'sentry/locale';
 import {space} from 'sentry/styles/space';
 import useOrganization from 'sentry/utils/useOrganization';
-import AiAnalyticsFilters from 'sentry/views/ai-analytics/filters';
+import AiAnalyticsFilters from 'sentry/views/aiAnalytics/filters';
+import TokenUsageChart from 'sentry/views/aiAnalytics/tokenUsageChart';
 
 function NoAccessComponent() {
   return (
@@ -35,6 +36,10 @@ function AiAnalyticsContainer() {
               <PageFiltersContainer>
                 <AiAnalyticsFilters onSearch={x => setSearch(x)} query={search} />
               </PageFiltersContainer>
+              <TwoColumns>
+                <TokenUsageChart metric="ai.prompt_tokens.used" />
+                <TokenUsageChart metric="ai.completion_tokens.used" />
+              </TwoColumns>
             </StyledMain>
           </StyledBody>
         </Layout.Page>
@@ -56,6 +61,11 @@ const StyledBody = styled('div')`
   grid-template-areas: 'content saved-searches';
 `;
 
+const TwoColumns = styled('div')`
+  display: flex;
+  gap: ${space(2)};
+`;
+
 const StyledMain = styled('section')`
   grid-area: content;
   padding: ${space(2)};
diff --git a/static/app/views/ai-analytics/landing.tsx b/static/app/views/aiAnalytics/landing.tsx
similarity index 100%
rename from static/app/views/ai-analytics/landing.tsx
rename to static/app/views/aiAnalytics/landing.tsx
diff --git a/static/app/views/aiAnalytics/tokenUsageChart.tsx b/static/app/views/aiAnalytics/tokenUsageChart.tsx
new file mode 100644
index 00000000000000..13602ba78ea041
--- /dev/null
+++ b/static/app/views/aiAnalytics/tokenUsageChart.tsx
@@ -0,0 +1,64 @@
+import styled from '@emotion/styled';
+
+import type {PageFilters} from 'sentry/types';
+import {MetricDisplayType} from 'sentry/utils/metrics/types';
+import {useMetricsQuery} from 'sentry/utils/metrics/useMetricsQuery';
+import withPageFilters from 'sentry/utils/withPageFilters';
+import {MetricChartContainer} from 'sentry/views/dashboards/metrics/chart';
+
+interface TokenUsageChartProps {
+  metric: 'ai.total_tokens.used' | 'ai.prompt_tokens.used' | 'ai.completion_tokens.used';
+  selection: PageFilters;
+  isGlobalSelectionReady?: boolean;
+}
+
+function TokenUsageChart({selection, metric}: TokenUsageChartProps) {
+  const {
+    data: timeseriesData,
+    isLoading,
+    isError,
+    error,
+  } = useMetricsQuery(
+    [
+      {
+        name: 'a',
+        mri: `c:custom/${metric}@tokens`,
+        op: 'sum',
+      },
+    ],
+    selection,
+    {
+      intervalLadder: 'dashboard',
+    }
+  );
+
+  if (isError) {
+    return <div>{'' + error}</div>;
+  }
+
+  return (
+    <TokenChartContainer>
+      <MetricChartContainer
+        timeseriesData={timeseriesData}
+        isLoading={isLoading}
+        metricQueries={[
+          {
+            name: 'mql',
+            formula: '$a',
+          },
+        ]}
+        displayType={MetricDisplayType.AREA}
+        chartHeight={200}
+      />
+    </TokenChartContainer>
+  );
+}
+
+const TokenChartContainer = styled('div')`
+  overflow: hidden;
+  border: 1px solid ${p => p.theme.border};
+  border-radius: ${p => p.theme.borderRadius};
+  flex: 1 1 content;
+`;
+
+export default withPageFilters(TokenUsageChart);