Flame Graph: Analyze with Grafana Assistant (#108684)

* Bare-bones mocked integration

* Create correct context based on the query

* Add data source name

* Do not bundle grafana/assistant with flame graph

* Rename component

* Add tests

* Mock grafana/assistant

* Update feature toggle and allow hiding the button

* Update deps

* Update types

* Update yarn.lock

* Fix typo in feature toggle description

* Enable grafanaAssistantInProfilesDrilldown by default

* Enable grafanaAssistantInProfilesDrilldown by default

* Show Analyze Flame Graph button only if there's context for the assistant

packages/grafana-flamegraph/src/AnalyzeFlameGraphButton.tsx

@@ -0,0 +1,34 @@
+import { ChatContextItem, useAssistant } from '@grafana/assistant';
+import { Button } from '@grafana/ui';
+
+type Props = {
+  assistantContext: ChatContextItem[];
+  className?: string;
+};
+
+export function AnalyzeFlameGraphButton(props: Props) {
+  const { assistantContext, className } = props;
+  const [isAvailable, openAssistant] = useAssistant();
+
+  if (!isAvailable || !openAssistant || assistantContext.length === 0) {
+    return null;
+  }
+
+  return (
+    <Button
+      className={className}
+      onClick={() =>
+        openAssistant({
+          prompt: 'Analyze Flame Graph',
+          context: assistantContext,
+        })
+      }
+      variant="secondary"
+      fill="outline"
+      icon="ai-sparkle"
+      size="sm"
+    >
+      Analyze Flame Graph
+    </Button>
+  );
+}

packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx

@@ -9,6 +9,11 @@ import { data } from './FlameGraph/testData/dataNestedSet';
 import FlameGraphContainer, { labelSearch } from './FlameGraphContainer';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
 
+jest.mock('@grafana/assistant', () => ({
+  useAssistant: jest.fn(() => [false, null]), // [isAvailable, openAssistant]
+  createContext: jest.fn(),
+}));
+
 jest.mock('react-use', () => ({
   ...jest.requireActual('react-use'),
   useMeasure: () => {

packages/grafana-flamegraph/src/FlameGraphContainer.tsx

@@ -14,6 +14,7 @@ import FlameGraphHeader from './FlameGraphHeader';
 import FlameGraphTopTableContainer from './TopTable/FlameGraphTopTableContainer';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
 import { ClickedItemData, ColorScheme, ColorSchemeDiff, SelectedView, TextAlign } from './types';
+import { getAssistantContextFromDataFrame } from './utils';
 
 const ufuzzy = new uFuzzy();
 
@@ -77,6 +78,14 @@ export type Props = {
    * Whether or not to keep any focused item when the profile data changes.
    */
   keepFocusOnDataChange?: boolean;
+
+  /**
+   * If true, the assistant button will be shown in the header if available.
+   * This is needed mainly for Profiles Drilldown where in some cases we need to hide the button to show alternative
+   * option to use AI.
+   * @default true
+   */
+  showAnalyzeWithAssistant?: boolean;
 };
 
 const FlameGraphContainer = ({
@@ -93,6 +102,7 @@ const FlameGraphContainer = ({
   disableCollapsing,
   keepFocusOnDataChange,
   getExtraContextMenuButtons,
+  showAnalyzeWithAssistant = true,
 }: Props) => {
   const [focusedItemData, setFocusedItemData] = useState<ClickedItemData>();
 
@@ -293,6 +303,7 @@ const FlameGraphContainer = ({
             isDiffMode={dataContainer.isDiffFlamegraph()}
             setCollapsedMap={setCollapsedMap}
             collapsedMap={collapsedMap}
+            assistantContext={data && showAnalyzeWithAssistant ? getAssistantContextFromDataFrame(data) : undefined}
           />
         )}
 

packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx

@@ -7,6 +7,11 @@ import { CollapsedMap } from './FlameGraph/dataTransform';
 import FlameGraphHeader from './FlameGraphHeader';
 import { ColorScheme, SelectedView } from './types';
 
+jest.mock('@grafana/assistant', () => ({
+  useAssistant: jest.fn(() => [false, null]), // [isAvailable, openAssistant]
+  createContext: jest.fn(),
+}));
+
 describe('FlameGraphHeader', () => {
   function setup(props: Partial<React.ComponentProps<typeof FlameGraphHeader>> = {}) {
     const setSearch = jest.fn();

packages/grafana-flamegraph/src/FlameGraphHeader.tsx

@@ -3,9 +3,11 @@ import { useEffect, useState } from 'react';
 import * as React from 'react';
 import { useDebounce, usePrevious } from 'react-use';
 
+import { ChatContextItem } from '@grafana/assistant';
 import { GrafanaTheme2, SelectableValue } from '@grafana/data';
 import { Button, ButtonGroup, Dropdown, Input, Menu, RadioButtonGroup, useStyles2 } from '@grafana/ui';
 
+import { AnalyzeFlameGraphButton } from './AnalyzeFlameGraphButton';
 import { byPackageGradient, byValueGradient, diffColorBlindGradient, diffDefaultGradient } from './FlameGraph/colors';
 import { CollapsedMap } from './FlameGraph/dataTransform';
 import { MIN_WIDTH_TO_SHOW_BOTH_TOPTABLE_AND_FLAMEGRAPH } from './constants';
@@ -30,6 +32,8 @@ type Props = {
   collapsedMap: CollapsedMap;
 
   extraHeaderElements?: React.ReactNode;
+
+  assistantContext?: ChatContextItem[];
 };
 
 const FlameGraphHeader = ({
@@ -50,6 +54,7 @@ const FlameGraphHeader = ({
   isDiffMode,
   setCollapsedMap,
   collapsedMap,
+  assistantContext,
 }: Props) => {
   const styles = useStyles2(getStyles);
   const [localSearch, setLocalSearch] = useSearchInput(search, setSearch);
@@ -84,6 +89,9 @@ const FlameGraphHeader = ({
       </div>
 
       <div className={styles.rightContainer}>
+        {assistantContext && (
+          <AnalyzeFlameGraphButton className={styles.buttonSpacing} assistantContext={assistantContext} />
+        )}
         {showResetButton && (
           <Button
             variant={'secondary'}

packages/grafana-flamegraph/src/utils.ts

@@ -0,0 +1,6 @@
+import { ChatContextItem } from '@grafana/assistant';
+import { DataFrame } from '@grafana/data';
+
+export function getAssistantContextFromDataFrame(data: DataFrame): ChatContextItem[] {
+  return data.meta?.custom?.assistantContext || [];
+}