From 68a5e448643eca1e8b74146bc871c7af41ef56a7 Mon Sep 17 00:00:00 2001 From: Stratoula Kalafateli Date: Tue, 23 Jan 2024 14:47:39 +0200 Subject: [PATCH] Suggest a chart --- .../public/functions/visualize_esql.tsx | 23 +++++++++--- .../server/functions/esql/index.ts | 36 +++++++++++++++++-- 2 files changed, 53 insertions(+), 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 index df7443d9eeb11..32b9775f0f6f9 100644 --- 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 index 40c01ae831704..c4dc3e340daca 100644 --- 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,