From 2cf5bd0d5d594023ab84c5e9bba13af37e39ce56 Mon Sep 17 00:00:00 2001
From: ivanahuckova <ivana.huckova@gmail.com>
Date: Wed, 14 Jan 2026 17:21:55 +0100
Subject: [PATCH] feat(query): query with Assistant action

---
 .../components/QueryActionAssistantButton.tsx | 133 ++++++++++++++++++
 .../query/components/QueryEditorRow.tsx       |  10 +-
 2 files changed, 142 insertions(+), 1 deletion(-)
 create mode 100644 public/app/features/query/components/QueryActionAssistantButton.tsx

diff --git a/public/app/features/query/components/QueryActionAssistantButton.tsx b/public/app/features/query/components/QueryActionAssistantButton.tsx
new file mode 100644
index 00000000000..0101714173f
--- /dev/null
+++ b/public/app/features/query/components/QueryActionAssistantButton.tsx
@@ -0,0 +1,133 @@
+import { useAssistant, createAssistantContextItem, OpenAssistantButton } from "@grafana/assistant";
+import { CoreApp, DataSourceApi, DataSourceInstanceSettings } from "@grafana/data";
+import { t } from "@grafana/i18n";
+import { DataQuery, DataSourceJsonData } from "@grafana/schema";
+import { queryIsEmpty } from "app/core/utils/query";
+
+
+interface QueryActionAssistantButtonProps<TQuery extends DataQuery = DataQuery> {
+    query: TQuery;
+    queries: TQuery[];
+    dataSourceInstance: DataSourceInstanceSettings;
+    app?: CoreApp;
+    datasource: DataSourceApi<TQuery, DataSourceJsonData, {}> | null;
+  }
+  
+  export function QueryActionAssistantButton<TQuery extends DataQuery = DataQuery>({
+    query,
+    queries,
+    dataSourceInstance,
+    app,
+    datasource,
+  }: QueryActionAssistantButtonProps<TQuery>) {
+    const { isAvailable } = useAssistant();
+    if (!isAvailable) {
+      return null;
+    }
+  
+    // Only show for Explore and Dashboard apps
+    if (app !== CoreApp.Explore && app !== CoreApp.Dashboard) {
+      return null;
+    }
+  
+    // Only show for loki and prometheus datasources
+    const pluginId = dataSourceInstance.type;
+    if (pluginId !== 'loki' && pluginId !== 'prometheus') {
+      return null;
+    }
+    const origin = `grafana/query-editor/${pluginId}/${app ?? CoreApp.Unknown}`;
+    
+    // Check if current query has content
+    const hasCurrentQuery = !queryIsEmpty(query);
+    const otherQueries = queries.filter((q) => q.refId !== query.refId && !queryIsEmpty(q));
+    
+    // Build context items
+    const context = [
+      createAssistantContextItem('datasource', {
+        datasourceUid: dataSourceInstance.uid,
+      }),
+    ];
+  
+    // Add current query if it has content
+    if (hasCurrentQuery) {
+      context.push(
+        createAssistantContextItem('structured', {
+          title: t('query-operation.header.current-query', 'Current query'),
+          data: query,
+        })
+      );
+    }
+  
+    // Add other queries if they exist
+    if (otherQueries.length > 0) {
+      context.push(
+        createAssistantContextItem('structured', {
+          title: t('query-operation.header.other-queries', 'Other queries'),
+          data: {
+            queries: otherQueries,
+          },
+        })
+      );
+    }
+  
+    // Get query display text to determine if we're creating or updating
+    const queryDisplayText = hasCurrentQuery && datasource?.getQueryDisplayText 
+      ? datasource.getQueryDisplayText(query) 
+      : null;
+    
+    // Determine if we're creating or updating based on queryDisplayText
+    const isUpdating = !!queryDisplayText;
+    const actionText = isUpdating 
+      ? t('query-operation.header.assistant-prompt-update', 'Help me update the current query to answer my questions and provide the insights I need.')
+      : t('query-operation.header.assistant-prompt-create', 'Help me create a new query to answer my questions and provide the insights I need.');
+    
+    // Format app name nicely
+    const appName = app === CoreApp.Explore 
+      ? t('query-operation.header.app-explore', 'Explore')
+      : app === CoreApp.Dashboard
+      ? t('query-operation.header.app-dashboard', 'Dashboard')
+      : '';
+    
+    // Build the prompt with proper formatting
+    const codeBlockLines: string[] = [];
+    
+    if (queryDisplayText) {
+      codeBlockLines.push(
+        t('query-operation.header.current-query-label', 'Current query:') + ` ${queryDisplayText}`
+      );
+    }
+    
+    codeBlockLines.push(
+      t('query-operation.header.selected-datasource-label', 'Selected data source:') + ` ${dataSourceInstance.name}`
+    );
+    
+    if (appName) {
+      codeBlockLines.push(
+        t('query-operation.header.app-label', 'App:') + ` ${appName}`
+      );
+    }
+    
+    // Add actionable sentence to motivate users
+    const actionableSentence = isUpdating
+      ? t('query-operation.header.assistant-actionable-update', 'Please describe what you want to change or improve in this query.')
+      : t('query-operation.header.assistant-actionable-create', 'Please describe what you want to query and what insights you\'re looking for.');
+    
+    // Build final prompt with code block
+    const prompt = [
+      actionText,
+      '```',
+      ...codeBlockLines,
+      '```',
+      actionableSentence,
+    ].join('\n');
+  
+    return (
+      <OpenAssistantButton
+        origin={origin}
+        prompt={prompt}
+        context={context}
+        autoSend={false}
+        title={t('query-operation.header.query-with-assistant', 'Query with Assistant')}
+      />
+    );
+  }
diff --git a/public/app/features/query/components/QueryEditorRow.tsx b/public/app/features/query/components/QueryEditorRow.tsx
index 5cc7f3467d3..5e09ee70272 100644
--- a/public/app/features/query/components/QueryEditorRow.tsx
+++ b/public/app/features/query/components/QueryEditorRow.tsx
@@ -36,6 +36,7 @@ import {
 import { useQueryLibraryContext } from '../../explore/QueryLibrary/QueryLibraryContext';
 import { ExpressionDatasourceUID } from '../../expressions/types';
 
+import { QueryActionAssistantButton } from './QueryActionAssistantButton';
 import { QueryActionComponent, RowActionComponents } from './QueryActionComponent';
 import { QueryEditorRowHeader } from './QueryEditorRowHeader';
 import { QueryErrorAlert } from './QueryErrorAlert';
@@ -399,7 +400,7 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
   };
 
   renderActions = (props: QueryOperationRowRenderProps) => {
-    const { query, hideHideQueryButton: hideHideQueryButton = false, queryLibraryRef, app } = this.props;
+    const { query, queries, hideHideQueryButton: hideHideQueryButton = false, queryLibraryRef, app } = this.props;
     const { datasource, showingHelp } = this.state;
     const isHidden = !!query.hide;
 
@@ -410,6 +411,13 @@ export class QueryEditorRow<TQuery extends DataQuery> extends PureComponent<Prop
 
     return (
       <>
+        <QueryActionAssistantButton
+          query={query}
+          queries={queries}
+          dataSourceInstance={this.props.dataSource}
+          datasource={datasource}
+          app={app}
+        />
         {!isEditingQueryLibrary && !isUnifiedAlerting && !isExpressionQuery && (
           <SavedQueryButtons
             query={{