feat: add ability to power LLM monitoring with EAP (#77750)

Requires #77749 first, this implements the LLM monitoring frontend using EAP instead of metrics.
getsentry · Sep 20, 2024 · b08a96a · b08a96a
1 parent a3859f0
commit b08a96a
Show file tree

Hide file tree

Showing 12 changed files with 523 additions and 61 deletions.
diff --git a/static/app/components/events/interfaces/llm-monitoring/llmMonitoringSection.tsx b/static/app/components/events/interfaces/llm-monitoring/llmMonitoringSection.tsx
@@ -1,4 +1,3 @@
-import Alert from 'sentry/components/alert';
 import {LinkButton} from 'sentry/components/button';
 import ButtonBar from 'sentry/components/buttonBar';
 import {IconOpen} from 'sentry/icons';
@@ -7,13 +6,22 @@ import type {Event} from 'sentry/types/event';
 import type {Organization} from 'sentry/types/organization';
 import {MutableSearch} from 'sentry/utils/tokenizeSearch';
 import * as ModuleLayout from 'sentry/views/insights/common/components/moduleLayout';
-import {useSpansIndexed} from 'sentry/views/insights/common/queries/useDiscover';
+import {
+  useEAPSpans,
+  useSpansIndexed,
+} from 'sentry/views/insights/common/queries/useDiscover';
 import {useModuleURL} from 'sentry/views/insights/common/utils/useModuleURL';
 import {
+  EAPNumberOfPipelinesChart,
+  EAPTotalTokensUsedChart,
   NumberOfPipelinesChart,
   TotalTokensUsedChart,
 } from 'sentry/views/insights/llmMonitoring/components/charts/llmMonitoringCharts';
-import {SpanIndexedField, type SpanIndexedResponse} from 'sentry/views/insights/types';
+import {
+  type EAPSpanResponse,
+  SpanIndexedField,
+  type SpanIndexedResponse,
+} from 'sentry/views/insights/types';
 import {SectionKey} from 'sentry/views/issueDetails/streamline/context';
 import {InterimSection} from 'sentry/views/issueDetails/streamline/interimSection';
 
@@ -22,20 +30,53 @@ interface Props {
   organization: Organization;
 }
 
-export default function LLMMonitoringSection({event}: Props) {
-  const traceId = event.contexts.trace?.trace_id;
-  const spanId = event.contexts.trace?.span_id;
-  const {data, error, isPending} = useSpansIndexed(
+function useAIPipelineGroup({
+  useEAP,
+  traceId,
+  spanId,
+}: {
+  useEAP: boolean;
+  spanId?: string;
+  traceId?: string;
+}): string | null {
+  const {data: indexedData} = useSpansIndexed(
     {
       limit: 1,
       fields: [SpanIndexedField.SPAN_AI_PIPELINE_GROUP],
       search: new MutableSearch(`trace:${traceId} id:"${spanId}"`),
+      enabled: !useEAP,
     },
     'api.ai-pipelines.view'
   );
+  const {data: eapData} = useEAPSpans(
+    {
+      limit: 1,
+      fields: [SpanIndexedField.SPAN_AI_PIPELINE_GROUP_TAG],
+      search: new MutableSearch(`trace:${traceId} id:"${spanId}"`),
+      enabled: useEAP,
+    },
+    'api.ai-pipelines-eap.view'
+  );
+
+  if (useEAP) {
+    return (
+      eapData &&
+      (eapData[0] as EAPSpanResponse)?.[SpanIndexedField.SPAN_AI_PIPELINE_GROUP_TAG]
+    );
+  }
+  return (
+    indexedData &&
+    (indexedData[0] as SpanIndexedResponse)?.[SpanIndexedField.SPAN_AI_PIPELINE_GROUP]
+  );
+}
+
+export default function LLMMonitoringSection({event, organization}: Props) {
   const moduleUrl = useModuleURL('ai');
-  const aiPipelineGroup =
-    data && (data[0] as SpanIndexedResponse)?.[SpanIndexedField.SPAN_AI_PIPELINE_GROUP];
+  const aiPipelineGroup = useAIPipelineGroup({
+    useEAP: organization.features.includes('insights-use-eap'),
+    traceId: event.contexts.trace?.trace_id,
+    spanId: event.contexts.trace?.span_id,
+  });
 
   const actions = (
     <ButtonBar gap={1}>
@@ -44,6 +85,7 @@ export default function LLMMonitoringSection({event}: Props) {
       </LinkButton>
     </ButtonBar>
   );
+  const useEAP = organization.features.includes('insights-use-eap');
 
   return (
     <InterimSection
@@ -52,19 +94,23 @@ export default function LLMMonitoringSection({event}: Props) {
       help={t('Charts showing how many tokens are being used')}
       actions={actions}
     >
-      {error ? (
-        <Alert type="error" showIcon>
-          {'' + error}
-        </Alert>
-      ) : isPending ? (
+      {!aiPipelineGroup ? (
         'loading'
       ) : (
         <ModuleLayout.Layout>
           <ModuleLayout.Half>
-            <TotalTokensUsedChart groupId={aiPipelineGroup} />
+            {useEAP ? (
+              <EAPTotalTokensUsedChart groupId={aiPipelineGroup} />
+            ) : (
+              <TotalTokensUsedChart groupId={aiPipelineGroup} />
+            )}
           </ModuleLayout.Half>
           <ModuleLayout.Half>
-            <NumberOfPipelinesChart groupId={aiPipelineGroup} />
+            {useEAP ? (
+              <EAPNumberOfPipelinesChart groupId={aiPipelineGroup} />
+            ) : (
+              <NumberOfPipelinesChart groupId={aiPipelineGroup} />
+            )}
           </ModuleLayout.Half>
         </ModuleLayout.Layout>
       )}

diff --git a/static/app/views/insights/common/queries/useDiscover.spec.tsx b/static/app/views/insights/common/queries/useDiscover.spec.tsx
@@ -13,7 +13,11 @@ import {
   useSpanMetrics,
   useSpansIndexed,
 } from 'sentry/views/insights/common/queries/useDiscover';
-import {SpanIndexedField, type SpanMetricsProperty} from 'sentry/views/insights/types';
+import {
+  SpanIndexedField,
+  type SpanIndexedProperty,
+  type SpanMetricsProperty,
+} from 'sentry/views/insights/types';
 import {OrganizationContext} from 'sentry/views/organizationContext';
 
 jest.mock('sentry/utils/useLocation');
@@ -196,7 +200,7 @@ describe('useDiscover', () => {
         {
           wrapper: Wrapper,
           initialProps: {
-            fields: [SpanIndexedField.SPAN_DESCRIPTION],
+            fields: [SpanIndexedField.SPAN_DESCRIPTION] as SpanIndexedProperty[],
             enabled: false,
           },
         }
@@ -253,7 +257,7 @@ describe('useDiscover', () => {
               SpanIndexedField.SPAN_OP,
               SpanIndexedField.SPAN_GROUP,
               SpanIndexedField.SPAN_DESCRIPTION,
-            ],
+            ] as SpanIndexedProperty[],
             sorts: [{field: 'span.group', kind: 'desc' as const}],
             limit: 10,
             referrer: 'api-spec',

diff --git a/static/app/views/insights/common/queries/useDiscover.ts b/static/app/views/insights/common/queries/useDiscover.ts
@@ -6,9 +6,11 @@ import type {MutableSearch} from 'sentry/utils/tokenizeSearch';
 import usePageFilters from 'sentry/utils/usePageFilters';
 import {useWrappedDiscoverQuery} from 'sentry/views/insights/common/queries/useSpansQuery';
 import type {
+  EAPSpanProperty,
+  EAPSpanResponse,
   MetricsProperty,
   MetricsResponse,
-  SpanIndexedField,
+  SpanIndexedProperty,
   SpanIndexedResponse,
   SpanMetricsProperty,
   SpanMetricsResponse,
@@ -25,7 +27,7 @@ interface UseMetricsOptions<Fields> {
   sorts?: Sort[];
 }
 
-export const useSpansIndexed = <Fields extends SpanIndexedField[]>(
+export const useSpansIndexed = <Fields extends SpanIndexedProperty[]>(
   options: UseMetricsOptions<Fields> = {},
   referrer: string
 ) => {
@@ -36,6 +38,17 @@ export const useSpansIndexed = <Fields extends SpanIndexedField[]>(
   );
 };
 
+export const useEAPSpans = <Fields extends EAPSpanProperty[]>(
+  options: UseMetricsOptions<Fields> = {},
+  referrer: string
+) => {
+  return useDiscover<Fields, EAPSpanResponse>(
+    options,
+    DiscoverDatasets.SPANS_EAP,
+    referrer
+  );
+};
+
 export const useSpanMetrics = <Fields extends SpanMetricsProperty[]>(
   options: UseMetricsOptions<Fields> = {},
   referrer: string

diff --git a/static/app/views/insights/common/queries/useFullSpanFromTrace.tsx b/static/app/views/insights/common/queries/useFullSpanFromTrace.tsx
@@ -4,7 +4,7 @@ import type {Sort} from 'sentry/utils/discover/fields';
 import {MutableSearch} from 'sentry/utils/tokenizeSearch';
 import {useSpansIndexed} from 'sentry/views/insights/common/queries/useDiscover';
 import {useEventDetails} from 'sentry/views/insights/common/queries/useEventDetails';
-import {SpanIndexedField} from 'sentry/views/insights/types';
+import {SpanIndexedField, type SpanIndexedProperty} from 'sentry/views/insights/types';
 
 const DEFAULT_SORT: Sort[] = [{field: 'timestamp', kind: 'desc'}];
 
@@ -34,7 +34,7 @@ export function useFullSpanFromTrace(
         SpanIndexedField.TRANSACTION_ID,
         SpanIndexedField.PROJECT,
         SpanIndexedField.ID,
-        ...(sorts?.map(sort => sort.field as SpanIndexedField) || []),
+        ...(sorts?.map(sort => sort.field as SpanIndexedProperty) || []),
       ],
     },
     'api.starfish.full-span-from-trace'

diff --git a/static/app/views/insights/http/queries/useSpanSamples.spec.tsx b/static/app/views/insights/http/queries/useSpanSamples.spec.tsx
@@ -8,7 +8,7 @@ import {QueryClientProvider} from 'sentry/utils/queryClient';
 import {MutableSearch} from 'sentry/utils/tokenizeSearch';
 import usePageFilters from 'sentry/utils/usePageFilters';
 import {useSpanSamples} from 'sentry/views/insights/http/queries/useSpanSamples';
-import {SpanIndexedField} from 'sentry/views/insights/types';
+import {SpanIndexedField, type SpanIndexedProperty} from 'sentry/views/insights/types';
 import {OrganizationContext} from 'sentry/views/organizationContext';
 
 jest.mock('sentry/utils/usePageFilters');
@@ -59,7 +59,10 @@ describe('useSpanSamples', () => {
       {
         wrapper: Wrapper,
         initialProps: {
-          fields: [SpanIndexedField.TRANSACTION_ID, SpanIndexedField.ID],
+          fields: [
+            SpanIndexedField.TRANSACTION_ID,
+            SpanIndexedField.ID,
+          ] as SpanIndexedProperty[],
           enabled: false,
         },
       }
@@ -100,7 +103,10 @@ describe('useSpanSamples', () => {
             release: '0.0.1',
             environment: undefined,
           },
-          fields: [SpanIndexedField.TRANSACTION_ID, SpanIndexedField.ID],
+          fields: [
+            SpanIndexedField.TRANSACTION_ID,
+            SpanIndexedField.ID,
+          ] as SpanIndexedProperty[],
           referrer: 'api-spec',
         },
       }

diff --git a/static/app/views/insights/http/queries/useSpanSamples.tsx b/static/app/views/insights/http/queries/useSpanSamples.tsx
@@ -6,7 +6,11 @@ import type {MutableSearch} from 'sentry/utils/tokenizeSearch';
 import useOrganization from 'sentry/utils/useOrganization';
 import usePageFilters from 'sentry/utils/usePageFilters';
 import {getDateConditions} from 'sentry/views/insights/common/utils/getDateConditions';
-import type {SpanIndexedField, SpanIndexedResponse} from 'sentry/views/insights/types';
+import type {
+  SpanIndexedField,
+  SpanIndexedProperty,
+  SpanIndexedResponse,
+} from 'sentry/views/insights/types';
 
 interface UseSpanSamplesOptions<Fields> {
   enabled?: boolean;
@@ -17,7 +21,7 @@ interface UseSpanSamplesOptions<Fields> {
   search?: MutableSearch;
 }
 
-export const useSpanSamples = <Fields extends SpanIndexedField[]>(
+export const useSpanSamples = <Fields extends SpanIndexedProperty[]>(
   options: UseSpanSamplesOptions<Fields> = {}
 ) => {
   const {