[Obs ai assistant][ESQL] Visualize only with ES|QL, no formulas

x-pack/plugins/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx

@@ -54,21 +54,6 @@ export function EsqlCodeBlock({
         </EuiFlexItem>
         <EuiFlexItem>
           <EuiFlexGroup direction="row" gutterSize="s">
-            <EuiFlexItem grow={false}>
-              <EuiButtonEmpty
-                data-test-subj="observabilityAiAssistantEsqlCodeBlockRunThisQueryButton"
-                size="xs"
-                iconType="play"
-                onClick={() =>
-                  onActionClick({ type: ChatActionClickType.executeEsqlQuery, query: value })
-                }
-                disabled={actionsDisabled}
-              >
-                {i18n.translate('xpack.observabilityAiAssistant.runThisQuery', {
-                  defaultMessage: 'Run this query',
-                })}
-              </EuiButtonEmpty>
-            </EuiFlexItem>
             <EuiFlexItem grow={false}>
               <EuiButtonEmpty
                 data-test-subj="observabilityAiAssistantEsqlCodeBlockVisualizeThisQueryButton"

x-pack/plugins/observability_ai_assistant/public/functions/visualize_esql.tsx

@@ -26,7 +26,6 @@ import type {
 import React, { useState, useEffect, useCallback, useMemo } from 'react';
 import ReactDOM from 'react-dom';
 import useAsync from 'react-use/lib/useAsync';
-import type { VisualizeESQLFunctionArguments } from '../../common/functions/visualize_esql';
 import type {
   ObservabilityAIAssistantPluginStartDependencies,
   ObservabilityAIAssistantService,
@@ -43,6 +42,12 @@ interface VisualizeLensResponse {
   content: DatatableColumn[];
 }
 
+interface VisualizeESQLFunctionArguments {
+  [x: string]: unknown;
+  query: string;
+  chartType?: string;
+}
+
 interface VisualizeESQLProps {
   /** Lens start contract, get the ES|QL charts suggestions api */
   lens: LensPublicStart;
@@ -66,6 +71,8 @@ interface VisualizeESQLProps {
    * if this is addressed https://github.com/elastic/eui/issues/7443
    */
   chatFlyoutSecondSlotHandler?: ChatFlyoutSecondSlotHandler;
+  /** User's preferation chart type as it comes from the model */
+  preferredChartType?: string;
 }
 
 function generateId() {
@@ -81,6 +88,7 @@ export function VisualizeESQL({
   onActionClick,
   initialInput,
   chatFlyoutSecondSlotHandler,
+  preferredChartType,
 }: VisualizeESQLProps) {
   // fetch the pattern from the query
   const indexPattern = getIndexPatternFromESQLQuery(query);
@@ -133,14 +141,24 @@ export function VisualizeESQL({
 
       const chartSuggestions = lensHelpersAsync.value.suggestions(context, dataViewAsync.value);
       if (chartSuggestions?.length) {
-        const [firstSuggestion] = chartSuggestions;
+        let [suggestion] = chartSuggestions;
+
+        if (chartSuggestions.length > 1 && preferredChartType) {
+          const suggestionFromModel = chartSuggestions.find(
+            (s) =>
+              s.title.includes(preferredChartType) || s.visualizationId.includes(preferredChartType)
+          );
+          if (suggestionFromModel) {
+            suggestion = suggestionFromModel;
+          }
+        }
 
         const attrs = getLensAttributesFromSuggestion({
           filters: [],
           query: {
             esql: query,
           },
-          suggestion: firstSuggestion,
+          suggestion,
           dataView: dataViewAsync.value,
         }) as TypedLensByValueInput['attributes'];
 
@@ -151,7 +169,7 @@ export function VisualizeESQL({
         setLensInput(lensEmbeddableInput);
       }
     }
-  }, [columns, dataViewAsync.value, lensHelpersAsync.value, lensInput, query]);
+  }, [columns, dataViewAsync.value, lensHelpersAsync.value, lensInput, preferredChartType, query]);
 
   // trigger options to open the inline editing flyout correctly
   const triggerOptions: InlineEditLensEmbeddableContext | undefined = useMemo(() => {
@@ -286,7 +304,7 @@ export function registerVisualizeQueryRenderFunction({
   registerRenderFunction(
     'visualize_query',
     ({
-      arguments: { query, newInput },
+      arguments: { query, newInput, chartType },
       response,
       onActionClick,
       chatFlyoutSecondSlotHandler,
@@ -302,6 +320,7 @@ export function registerVisualizeQueryRenderFunction({
           onActionClick={onActionClick}
           initialInput={newInput}
           chatFlyoutSecondSlotHandler={chatFlyoutSecondSlotHandler}
+          preferredChartType={chartType}
         />
       );
     }

x-pack/plugins/observability_ai_assistant/server/functions/esql/index.ts

@@ -11,6 +11,8 @@ import pLimit from 'p-limit';
 import Path from 'path';
 import { lastValueFrom, Observable } from 'rxjs';
 import { promisify } from 'util';
+import { esFieldTypeToKibanaFieldType } from '@kbn/field-types';
+import type { ESQLSearchReponse } from '@kbn/es-types';
 import type { FunctionRegistrationParameters } from '..';
 import {
   ChatCompletionChunkEvent,
@@ -21,6 +23,17 @@ import { concatenateChatCompletionChunks } from '../../../common/utils/concatena
 import { emitWithConcatenatedMessage } from '../../../common/utils/emit_with_concatenated_message';
 import { correctCommonEsqlMistakes } from './correct_common_esql_mistakes';
 
+enum ChartType {
+  XY = 'XY',
+  Bar = 'Bar',
+  Line = 'Line',
+  Donut = 'Donut',
+  Heatmap = 'Heat map',
+  Treemap = 'Treemap',
+  Tagcloud = 'Tag cloud',
+  Waffle = 'Waffle',
+}
+
 const readFile = promisify(Fs.readFile);
 const readdir = promisify(Fs.readdir);
 
@@ -69,33 +82,46 @@ export function registerEsqlFunction({
 }: FunctionRegistrationParameters) {
   registerFunction(
     {
-      name: 'execute_query',
-      contexts: ['core'],
-      visibility: FunctionVisibility.User,
-      description: 'Execute an ES|QL query.',
+      name: 'visualize_query',
+      description:
+        'Use this function to visualize charts for ES|QL queries. The visualisation is displayed to the user above your reply, DO NOT try to generate or display an image yourself.',
+      descriptionForUser: 'Use this function to visualize charts for ES|QL queries.',
       parameters: {
         type: 'object',
-        additionalProperties: false,
+        additionalProperties: true,
         properties: {
           query: {
             type: 'string',
           },
+          chartType: {
+            type: 'string',
+          },
         },
         required: ['query'],
       } as const,
+      contexts: ['core'],
     },
-    async ({ arguments: { query } }) => {
-      const response = await (
-        await resources.context.core
+    async ({ arguments: { query }, connectorId, messages }, signal) => {
+      // With limit 0 I get only the columns, it is much more performant
+      const performantQuery = `${query} | limit 0`;
+      const coreContext = await resources.context.core;
+
+      const response = (await (
+        await coreContext
       ).elasticsearch.client.asCurrentUser.transport.request({
         method: 'POST',
         path: '_query',
         body: {
-          query,
+          query: performantQuery,
         },
-      });
-
-      return { content: response };
+      })) as ESQLSearchReponse;
+      const columns =
+        response.columns?.map(({ name, type }) => ({
+          id: name,
+          name,
+          meta: { type: esFieldTypeToKibanaFieldType(type) },
+        })) ?? [];
+      return { content: columns };
     }
   );
 
@@ -140,12 +166,15 @@ export function registerEsqlFunction({
             Extract data? Request \`DISSECT\` AND \`GROK\`.
             Convert a column based on a set of conditionals? Request \`EVAL\` and \`CASE\`.
 
-            Examples for determining whether the user wants to execute a query:
+            Examples for determining whether the user wants to visualize a query:
             - "Show me the avg of x"
             - "Give me the results of y"
             - "Display the sum of z"
+            - "I want to visualize ..."
+            - "I want to display the avg of ..."
+            - I want a chart ..."
 
-            Examples for determining whether the user does not want to execute a query:
+            Examples for determining whether the user does not want to visualize a query:
             - "I want a query that ..."
             - "... Just show me the query"
             - "Create a query that ..."`
@@ -156,7 +185,7 @@ export function registerEsqlFunction({
               name: 'classify_esql',
               description: `Use this function to determine:
               - what ES|QL functions and commands are candidates for answering the user's question
-              - whether the user has requested a query, and if so, it they want it to be executed, or just shown.
+              - whether the user has requested a query, and if so, it they want it to be visualized, or just shown.
               `,
               parameters: {
                 type: 'object',
@@ -178,7 +207,20 @@ export function registerEsqlFunction({
                   execute: {
                     type: 'boolean',
                     description:
-                      'Whether the user wants to execute a query (true) or just wants the query to be displayed (false)',
+                      'Whether the user wants to visualize a query (true) or just wants the query to be displayed (false)',
+                  },
+                  chartType: {
+                    type: 'string',
+                    enum: [
+                      ChartType.XY,
+                      ChartType.Bar,
+                      ChartType.Line,
+                      ChartType.Donut,
+                      ChartType.Treemap,
+                      ChartType.Heatmap,
+                      ChartType.Tagcloud,
+                      ChartType.Waffle,
+                    ],
                   },
                 },
                 required: ['commands', 'functions', 'execute'],
@@ -195,6 +237,7 @@ export function registerEsqlFunction({
         commands: string[];
         functions: string[];
         execute: boolean;
+        chartType?: ChartType;
       };
 
       const keywords = args.commands.concat(args.functions).concat('SYNTAX').concat('OVERVIEW');
@@ -291,8 +334,11 @@ export function registerEsqlFunction({
                   id,
                   message: {
                     function_call: {
-                      name: 'execute_query',
-                      arguments: JSON.stringify({ query: esqlQuery }),
+                      name: 'visualize_query',
+                      arguments: JSON.stringify({
+                        query: esqlQuery,
+                        chartType: args.chartType,
+                      }),
                     },
                   },
                   type: StreamingChatResponseEventType.ChatCompletionChunk,

x-pack/plugins/observability_ai_assistant/server/functions/index.ts

@@ -13,9 +13,7 @@ import { registerAlertsFunction } from './alerts';
 import { registerElasticsearchFunction } from './elasticsearch';
 import { registerEsqlFunction } from './esql';
 import { registerGetDatasetInfoFunction } from './get_dataset_info';
-import { registerLensFunction } from './lens';
 import { registerKibanaFunction } from './kibana';
-import { registerVisualizeESQLFunction } from './visualize_esql';
 
 export type FunctionRegistrationParameters = Omit<
   Parameters<ChatRegistrationFunction>[0],
@@ -76,8 +74,6 @@ export const registerFunctions: ChatRegistrationFunction = async ({
 
       registerSummarizationFunction(registrationParameters);
       registerRecallFunction(registrationParameters);
-      registerLensFunction(registrationParameters);
-      registerVisualizeESQLFunction(registrationParameters);
     } else {
       description += `You do not have a working memory. Don't try to recall information via the "recall" function.  If the user expects you to remember the previous conversations, tell them they can set up the knowledge base. A banner is available at the top of the conversation to set this up.`;
     }