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
2025-09-15 18:32:51 +08:00 · 2025-08-19 09:54:00 +02:00
parent ea833ddf1f
commit bbf01a6383
19 changed files with 539 additions and 8 deletions
--- a/docs/sources/setup-grafana/configure-grafana/feature-toggles/index.md
+++ b/docs/sources/setup-grafana/configure-grafana/feature-toggles/index.md
@ -76,6 +76,7 @@ Most [generally available](https://grafana.com/docs/release-life-cycle/#general-
 | `alertingMigrationUI`                  | Enables the alerting migration UI, to migrate data source-managed rules to Grafana-managed rules                                                              | Yes                |
 | `alertingImportYAMLUI`                 | Enables a UI feature for importing rules from a Prometheus file to Grafana-managed rules                                                                      | Yes                |
 | `unifiedNavbars`                       | Enables unified navbars                                                                                                                                       |                    |
+| `grafanaAssistantInProfilesDrilldown`  | Enables integration with Grafana Assistant in Profiles Drilldown                                                                                              | Yes                |
 | `tabularNumbers`                       | Use fixed-width numbers globally in the UI                                                                                                                    |                    |

 ## Public preview feature toggles
--- a/packages/grafana-data/src/types/featureToggles.gen.ts
+++ b/packages/grafana-data/src/types/featureToggles.gen.ts
@ -950,6 +950,11 @@ export interface FeatureToggles {
  */
  metricsFromProfiles?: boolean;
  /**
+  * Enables integration with Grafana Assistant in Profiles Drilldown
+  * @default true
+  */
+  grafanaAssistantInProfilesDrilldown?: boolean;
+  /**
  * Enables using PGX instead of libpq for PostgreSQL datasource
  */
  postgresDSUsePGX?: boolean;
--- a/packages/grafana-flamegraph/package.json
+++ b/packages/grafana-flamegraph/package.json
@ -83,6 +83,7 @@
    "typescript": "5.9.2"
  },
  "peerDependencies": {
+    "@grafana/assistant": "^0.0.12",
    "react": "^18.0.0",
    "react-dom": "^18.0.0"
  }
--- a/packages/grafana-flamegraph/src/AnalyzeFlameGraphButton.tsx
+++ b/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>
+  );
+}
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.test.tsx
+++ b/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: () => {
--- a/packages/grafana-flamegraph/src/FlameGraphContainer.tsx
+++ b/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}
          />
        )}

--- a/packages/grafana-flamegraph/src/FlameGraphHeader.test.tsx
+++ b/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();
--- a/packages/grafana-flamegraph/src/FlameGraphHeader.tsx
+++ b/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'}
--- a/packages/grafana-flamegraph/src/utils.ts
+++ b/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 || [];
+}
--- a/pkg/services/featuremgmt/registry.go
+++ b/pkg/services/featuremgmt/registry.go
@ -1638,6 +1638,14 @@ var (
 			Owner:        grafanaObservabilityTracesAndProfilingSquad,
 			FrontendOnly: true,
 		},
+		{
+			Name:         "grafanaAssistantInProfilesDrilldown",
+			Description:  "Enables integration with Grafana Assistant in Profiles Drilldown",
+			Stage:        FeatureStageGeneralAvailability,
+			Owner:        grafanaObservabilityTracesAndProfilingSquad,
+			FrontendOnly: true,
+			Expression:   "true",
+		},
 		{
 			Name:        "postgresDSUsePGX",
 			Description: "Enables using PGX instead of libpq for PostgreSQL datasource",
--- a/pkg/services/featuremgmt/toggles_gen.csv
+++ b/pkg/services/featuremgmt/toggles_gen.csv
@ -213,6 +213,7 @@ localizationForPlugins,experimental,@grafana/plugins-platform-backend,false,fals
 unifiedNavbars,GA,@grafana/plugins-platform-backend,false,false,true
 logsPanelControls,preview,@grafana/observability-logs,false,false,true
 metricsFromProfiles,experimental,@grafana/observability-traces-and-profiling,false,false,true
+grafanaAssistantInProfilesDrilldown,GA,@grafana/observability-traces-and-profiling,false,false,true
 postgresDSUsePGX,experimental,@grafana/oss-big-tent,false,false,false
 tempoAlerting,experimental,@grafana/observability-traces-and-profiling,false,false,true
 pluginsAutoUpdate,experimental,@grafana/plugins-platform-backend,false,false,false
--- a/pkg/services/featuremgmt/toggles_gen.go
+++ b/pkg/services/featuremgmt/toggles_gen.go
@ -863,6 +863,10 @@ const (
 	// Enables creating metrics from profiles and storing them as recording rules
 	FlagMetricsFromProfiles = "metricsFromProfiles"

+	// FlagGrafanaAssistantInProfilesDrilldown
+	// Enables integration with Grafana Assistant in Profiles Drilldown
+	FlagGrafanaAssistantInProfilesDrilldown = "grafanaAssistantInProfilesDrilldown"
+
 	// FlagPostgresDSUsePGX
 	// Enables using PGX instead of libpq for PostgreSQL datasource
 	FlagPostgresDSUsePGX = "postgresDSUsePGX"
--- a/pkg/services/featuremgmt/toggles_gen.json
+++ b/pkg/services/featuremgmt/toggles_gen.json
@ -1472,6 +1472,37 @@
        "codeowner": "@grafana/plugins-platform-backend"
      }
    },
+    {
+      "metadata": {
+        "name": "grafanaAssistantInProfilesDrilldown",
+        "resourceVersion": "1754572610001",
+        "creationTimestamp": "2025-08-01T07:43:17Z",
+        "annotations": {
+          "grafana.app/updatedTimestamp": "2025-08-07 13:16:50.001205 +0000 UTC"
+        }
+      },
+      "spec": {
+        "description": "Enables integration with Grafana Assistant in Profiles Drilldown",
+        "stage": "GA",
+        "codeowner": "@grafana/observability-traces-and-profiling",
+        "frontend": true,
+        "expression": "true"
+      }
+    },
+    {
+      "metadata": {
+        "name": "grafanaAssistantInProfilesDrillfown",
+        "resourceVersion": "1754034112469",
+        "creationTimestamp": "2025-08-01T07:41:52Z",
+        "deletionTimestamp": "2025-08-01T07:43:17Z"
+      },
+      "spec": {
+        "description": "Enables interation with Grafana Assitant in Profiles Drilldown",
+        "stage": "experimental",
+        "codeowner": "@grafana/observability-traces-and-profiling",
+        "frontend": true
+      }
+    },
    {
      "metadata": {
        "name": "grafanaManagedRecordingRules",
--- a/public/app/plugins/datasource/grafana-pyroscope-datasource/datasource.ts
+++ b/public/app/plugins/datasource/grafana-pyroscope-datasource/datasource.ts
@ -1,5 +1,5 @@
 import Prism from 'prismjs';
-import { Observable, of } from 'rxjs';
+import { map, Observable, of } from 'rxjs';

 import {
  AbstractQuery,
@ -18,7 +18,13 @@ import { DataSourceWithBackend, getTemplateSrv, TemplateSrv } from '@grafana/run
 import { VariableSupport } from './VariableSupport';
 import { defaultGrafanaPyroscopeDataQuery, defaultPyroscopeQueryType } from './dataquery.gen';
 import { PyroscopeDataSourceOptions, Query, ProfileTypeMessage } from './types';
-import { addLabelToQuery, extractLabelMatchers, grammar, toPromLikeExpr } from './utils';
+import {
+  addLabelToQuery,
+  extractLabelMatchers,
+  grammar,
+  toPromLikeExpr,
+  enrichDataFrameWithAssistantContentMapper,
+} from './utils';

 export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeDataSourceOptions> {
  constructor(
@ -45,10 +51,12 @@ export class PyroscopeDataSource extends DataSourceWithBackend<Query, PyroscopeD
    if (!validTargets.length) {
      return of({ data: [] });
    }
-    return super.query({
-      ...request,
-      targets: validTargets,
-    });
+    return super
+      .query({
+        ...request,
+        targets: validTargets,
+      })
+      .pipe(map(enrichDataFrameWithAssistantContentMapper(request, this.name)));
  }

  async getProfileTypes(start: number, end: number): Promise<ProfileTypeMessage[]> {
--- a/public/app/plugins/datasource/grafana-pyroscope-datasource/package.json
+++ b/public/app/plugins/datasource/grafana-pyroscope-datasource/package.json
@ -39,6 +39,7 @@
    "webpack": "5.101.0"
  },
  "peerDependencies": {
+    "@grafana/assistant": "^0.0.12",
    "@grafana/runtime": "*"
  },
  "scripts": {
--- a/public/app/plugins/datasource/grafana-pyroscope-datasource/utils.test.ts
+++ b/public/app/plugins/datasource/grafana-pyroscope-datasource/utils.test.ts
@ -0,0 +1,342 @@
+import { createContext, ItemDataType } from '@grafana/assistant';
+import {
+  DataFrame,
+  DataQueryResponse,
+  DataQueryRequest,
+  FieldType,
+  MutableDataFrame,
+  dateTime,
+  PreferredVisualisationType,
+} from '@grafana/data';
+
+import { GrafanaPyroscopeDataQuery } from './dataquery.gen';
+import { enrichDataFrameWithAssistantContentMapper } from './utils';
+
+// Mock the createContext function
+jest.mock('@grafana/assistant', () => ({
+  createContext: jest.fn(),
+  ItemDataType: {
+    Datasource: 'datasource',
+    Structured: 'structured',
+  },
+}));
+
+const mockCreateContext = createContext as jest.MockedFunction<typeof createContext>;
+
+describe('enrichDataFrameWithAssistantContentMapper', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+    mockCreateContext.mockImplementation((type, data) => ({
+      type,
+      data,
+      node: { id: 'test-id', type: 'test', name: 'test-node', navigable: true },
+      occurrences: [],
+    }));
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  const createMockRequest = (
+    targets: Array<Partial<GrafanaPyroscopeDataQuery>> = []
+  ): DataQueryRequest<GrafanaPyroscopeDataQuery> => ({
+    requestId: 'test-request',
+    interval: '1s',
+    intervalMs: 1000,
+    maxDataPoints: 1000,
+    range: {
+      from: dateTime('2023-01-01T00:00:00Z'),
+      to: dateTime('2023-01-01T01:00:00Z'),
+      raw: {
+        from: 'now-1h',
+        to: 'now',
+      },
+    },
+    scopedVars: {},
+    timezone: 'UTC',
+    app: 'explore',
+    startTime: Date.now(),
+    targets: targets.map((target, index) => ({
+      refId: `A${index}`,
+      profileTypeId: 'cpu',
+      labelSelector: '{service="test"}',
+      groupBy: [],
+      queryType: 'profile' as const,
+      datasource: {
+        uid: 'test-uid',
+        type: 'grafana-pyroscope-datasource',
+      },
+      ...target,
+    })) as GrafanaPyroscopeDataQuery[],
+  });
+
+  const createMockDataFrame = (refId: string, preferredVisualisationType?: PreferredVisualisationType): DataFrame => {
+    const frame = new MutableDataFrame({
+      refId,
+      fields: [
+        { name: 'time', type: FieldType.time, values: [1672531200000] },
+        { name: 'value', type: FieldType.number, values: [100] },
+      ],
+    });
+
+    if (preferredVisualisationType) {
+      frame.meta = {
+        preferredVisualisationType,
+      };
+    }
+
+    return frame;
+  };
+
+  const createMockResponse = (dataFrames: DataFrame[]): DataQueryResponse => ({
+    data: dataFrames,
+  });
+
+  describe('when processing flamegraph data frames', () => {
+    it('should enrich flamegraph data frame with assistant context', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+        },
+      ]);
+      const dataFrame = createMockDataFrame('A0', 'flamegraph');
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'PyroscopeDatasource');
+      const result = mapper(response);
+
+      expect(result.data).toHaveLength(1);
+      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined();
+      expect(mockCreateContext).toHaveBeenCalledTimes(2);
+
+      // Verify datasource context
+      expect(mockCreateContext).toHaveBeenCalledWith(ItemDataType.Datasource, {
+        datasourceName: 'PyroscopeDatasource',
+        datasourceUid: 'test-uid',
+        datasourceType: 'grafana-pyroscope-datasource',
+      });
+
+      // Verify structured context
+      expect(mockCreateContext).toHaveBeenCalledWith(ItemDataType.Structured, {
+        title: 'Analyze Flame Graph',
+        data: {
+          start: request.range.from.valueOf(),
+          end: request.range.to.valueOf(),
+          profile_type_id: 'cpu',
+          label_selector: '{service="test"}',
+          operation: 'execute',
+        },
+      });
+    });
+
+    it('should preserve existing meta.custom properties', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'memory',
+          labelSelector: '{app="backend"}',
+        },
+      ]);
+      const dataFrame = createMockDataFrame('A0', 'flamegraph');
+      dataFrame.meta = {
+        preferredVisualisationType: 'flamegraph',
+        custom: {
+          existingProperty: 'value',
+        },
+      };
+
+      const response = createMockResponse([dataFrame]);
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data[0].meta?.custom?.existingProperty).toBe('value');
+      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined();
+    });
+
+    it('should handle multiple flamegraph data frames', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test1"}',
+        },
+        {
+          refId: 'A1',
+          profileTypeId: 'memory',
+          labelSelector: '{service="test2"}',
+        },
+      ]);
+
+      const dataFrames = [createMockDataFrame('A0', 'flamegraph'), createMockDataFrame('A1', 'flamegraph')];
+      const response = createMockResponse(dataFrames);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data).toHaveLength(2);
+      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined();
+      expect(result.data[1].meta?.custom?.assistantContext).toBeDefined();
+      expect(mockCreateContext).toHaveBeenCalledTimes(4); // 2 contexts per frame
+    });
+  });
+
+  describe('when processing non-flamegraph data frames', () => {
+    it('should not modify data frames without flamegraph visualization type', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+        },
+      ]);
+      const dataFrame = createMockDataFrame('A0', 'table');
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data[0].meta?.custom?.assistantContext).toBeUndefined();
+      expect(mockCreateContext).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('when handling edge cases', () => {
+    it('should not add context data frame without matching query target', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+        },
+      ]);
+      const dataFrame = createMockDataFrame('B0', 'flamegraph'); // Different refId
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data[0]).toBe(dataFrame); // Should remain unchanged
+      expect(mockCreateContext).not.toHaveBeenCalled();
+    });
+
+    it('should not add context when query has no datasource information', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+          datasource: undefined,
+        },
+      ]);
+      const dataFrame = createMockDataFrame('A0', 'flamegraph');
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data[0]).toBe(dataFrame); // Should remain unchanged
+      expect(mockCreateContext).not.toHaveBeenCalled();
+    });
+
+    it('should not add context if query has incomplete datasource information', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+          datasource: {
+            uid: 'test-uid',
+            // Missing type
+          },
+        },
+      ]);
+      const dataFrame = createMockDataFrame('A0', 'flamegraph');
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data[0]).toBe(dataFrame); // Should remain unchanged
+      expect(mockCreateContext).not.toHaveBeenCalled();
+    });
+
+    it('should not add context with empty response data', () => {
+      const request = createMockRequest([]);
+      const response = createMockResponse([]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data).toHaveLength(0);
+      expect(mockCreateContext).not.toHaveBeenCalled();
+    });
+
+    it('should handle mixed data frame types', () => {
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+        },
+        {
+          refId: 'A1',
+          profileTypeId: 'memory',
+          labelSelector: '{service="test2"}',
+        },
+      ]);
+
+      const dataFrames = [
+        createMockDataFrame('A0', 'flamegraph'),
+        createMockDataFrame('A1', 'table'),
+        createMockDataFrame('A0', 'flamegraph'), // Another flamegraph with same refId
+      ];
+      const response = createMockResponse(dataFrames);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      const result = mapper(response);
+
+      expect(result.data).toHaveLength(3);
+      expect(result.data[0].meta?.custom?.assistantContext).toBeDefined(); // Flamegraph
+      expect(result.data[1].meta?.custom?.assistantContext).toBeUndefined(); // Table
+      expect(result.data[2].meta?.custom?.assistantContext).toBeDefined(); // Flamegraph
+      expect(mockCreateContext).toHaveBeenCalledTimes(4); // 2 contexts for each flamegraph
+    });
+  });
+
+  describe('context creation', () => {
+    it('should create context with correct time range values', () => {
+      const fromTime = dateTime('2023-06-15T10:00:00Z');
+      const toTime = dateTime('2023-06-15T11:00:00Z');
+
+      const request = createMockRequest([
+        {
+          refId: 'A0',
+          profileTypeId: 'cpu',
+          labelSelector: '{service="test"}',
+        },
+      ]);
+      request.range.from = fromTime;
+      request.range.to = toTime;
+
+      const dataFrame = createMockDataFrame('A0', 'flamegraph');
+      const response = createMockResponse([dataFrame]);
+
+      const mapper = enrichDataFrameWithAssistantContentMapper(request, 'TestDatasource');
+      mapper(response);
+
+      expect(mockCreateContext).toHaveBeenCalledWith(ItemDataType.Structured, {
+        title: 'Analyze Flame Graph',
+        data: {
+          start: fromTime.valueOf(),
+          end: toTime.valueOf(),
+          profile_type_id: 'cpu',
+          label_selector: '{service="test"}',
+          operation: 'execute',
+        },
+      });
+    });
+  });
+});
--- a/public/app/plugins/datasource/grafana-pyroscope-datasource/utils.ts
+++ b/public/app/plugins/datasource/grafana-pyroscope-datasource/utils.ts
@ -1,7 +1,16 @@
 import { invert } from 'lodash';
 import Prism, { Grammar, Token } from 'prismjs';

-import { AbstractLabelMatcher, AbstractLabelOperator } from '@grafana/data';
+import { createContext, ItemDataType } from '@grafana/assistant';
+import {
+  AbstractLabelMatcher,
+  AbstractLabelOperator,
+  DataFrame,
+  DataQueryResponse,
+  DataQueryRequest,
+} from '@grafana/data';
+
+import { GrafanaPyroscopeDataQuery } from './dataquery.gen';

 export function extractLabelMatchers(tokens: Array<string | Token>): AbstractLabelMatcher[] {
  const labelMatchers: AbstractLabelMatcher[] = [];
@ -131,3 +140,48 @@ export const grammar: Grammar = {
  },
  punctuation: /[{}(),.]/,
 };
+
+export function enrichDataFrameWithAssistantContentMapper(
+  request: DataQueryRequest<GrafanaPyroscopeDataQuery>,
+  datasourceName: string
+) {
+  const validTargets = request.targets;
+  return (response: DataQueryResponse) => {
+    response.data = response.data.map((data: DataFrame) => {
+      if (data.meta?.preferredVisualisationType !== 'flamegraph') {
+        return data;
+      }
+
+      const query = validTargets.find((target) => target.refId === data.refId);
+      if (!query || !query.datasource?.uid || !query.datasource?.type) {
+        return data;
+      }
+
+      const context = [
+        createContext(ItemDataType.Datasource, {
+          datasourceName: datasourceName,
+          datasourceUid: query.datasource.uid,
+          datasourceType: query.datasource.type,
+        }),
+        createContext(ItemDataType.Structured, {
+          title: 'Analyze Flame Graph',
+          data: {
+            start: request.range.from.valueOf(),
+            end: request.range.to.valueOf(),
+            profile_type_id: query.profileTypeId,
+            label_selector: query.labelSelector,
+            operation: 'execute',
+          },
+        }),
+      ];
+
+      data.meta = data.meta || {};
+      data.meta.custom = {
+        ...data.meta.custom,
+        assistantContext: context,
+      };
+      return data;
+    });
+    return response;
+  };
+}
--- a/public/app/plugins/datasource/grafana-pyroscope-datasource/webpack.config.ts
+++ b/public/app/plugins/datasource/grafana-pyroscope-datasource/webpack.config.ts
@ -1,3 +1,5 @@
+import { merge } from 'webpack-merge';
+
 import config, { type Env } from '@grafana/plugin-configs/webpack.config.ts';

 const configWithFallback = async (env: Env) => {
@ -9,7 +11,9 @@ const configWithFallback = async (env: Env) => {
      string_decoder: false,
    };
  }
-  return response;
+  return merge(response, {
+    externals: ['grafana/assistant'],
+  });
 };

 export default configWithFallback;
--- a/yarn.lock
+++ b/yarn.lock
@ -2648,6 +2648,7 @@ __metadata:
    typescript: "npm:5.9.2"
    webpack: "npm:5.101.0"
  peerDependencies:
+    "@grafana/assistant": ^0.0.12
    "@grafana/runtime": "*"
  languageName: unknown
  linkType: soft
@ -3285,6 +3286,7 @@ __metadata:
    tslib: "npm:2.8.1"
    typescript: "npm:5.9.2"
  peerDependencies:
+    "@grafana/assistant": ^0.0.12
    react: ^18.0.0
    react-dom: ^18.0.0
  languageName: unknown