Suggest a chart

stratoula · Jan 23, 2024 · 68a5e44 · 68a5e44
1 parent 2006470
commit 68a5e44
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 6 deletions.
diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/visualize_esql.tsx b/x-pack/plugins/observability_ai_assistant/public/functions/visualize_esql.tsx
@@ -45,6 +45,7 @@ interface VisualizeLensResponse {
 interface VisualizeESQLFunctionArguments {
   [x: string]: unknown;
   query: string;
+  chartType?: string;
 }
 
 interface VisualizeESQLProps {
@@ -70,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() {
@@ -85,6 +88,7 @@ export function VisualizeESQL({
   onActionClick,
   initialInput,
   chatFlyoutSecondSlotHandler,
+  preferredChartType,
 }: VisualizeESQLProps) {
   // fetch the pattern from the query
   const indexPattern = getIndexPatternFromESQLQuery(query);
@@ -137,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'];
 
@@ -155,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(() => {
@@ -290,7 +304,7 @@ export function registerVisualizeQueryRenderFunction({
   registerRenderFunction(
     'visualize_query',
     ({
-      arguments: { query, newInput },
+      arguments: { query, newInput, chartType },
       response,
       onActionClick,
       chatFlyoutSecondSlotHandler,
@@ -306,6 +320,7 @@ export function registerVisualizeQueryRenderFunction({
           onActionClick={onActionClick}
           initialInput={newInput}
           chatFlyoutSecondSlotHandler={chatFlyoutSecondSlotHandler}
+          preferredChartType={chartType}
         />
       );
     }

diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/esql/index.ts b/x-pack/plugins/observability_ai_assistant/server/functions/esql/index.ts
@@ -23,6 +23,17 @@ import { concatenateOpenAiChunks } from '../../../common/utils/concatenate_opena
 import { processOpenAiStream } from '../../../common/utils/process_openai_stream';
 import { streamIntoObservable } from '../../service/util/stream_into_observable';
 
+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);
 
@@ -73,7 +84,7 @@ export function registerEsqlFunction({
     {
       name: 'visualize_query',
       description:
-        'Use this function to visualize charts for ES|QL queries. Do NOT run the lens function and do not try to visualize it on your own',
+        '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',
@@ -82,6 +93,9 @@ export function registerEsqlFunction({
           query: {
             type: 'string',
           },
+          chartType: {
+            type: 'string',
+          },
         },
         required: ['query'],
       } as const,
@@ -189,6 +203,7 @@ export function registerEsqlFunction({
             - "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 visualize a query:
@@ -227,6 +242,19 @@ export function registerEsqlFunction({
                     description:
                       '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'],
               },
@@ -242,6 +270,7 @@ export function registerEsqlFunction({
         commands: string[];
         functions: string[];
         execute: boolean;
+        chartType?: ChartType;
       };
 
       const keywords = args.commands.concat(args.functions).concat('SYNTAX').concat('OVERVIEW');
@@ -381,7 +410,10 @@ export function registerEsqlFunction({
                       delta: {
                         function_call: {
                           name: 'visualize_query',
-                          arguments: JSON.stringify({ query: esqlQuery }),
+                          arguments: JSON.stringify({
+                            query: esqlQuery,
+                            chartType: args.chartType,
+                          }),
                         },
                       },
                       index: 0,