diff --git a/x-pack/plugins/observability_ai_assistant/common/functions/visualize_esql.tsx b/x-pack/plugins/observability_ai_assistant/common/functions/visualize_esql.tsx deleted file mode 100644 index 48f9b2a4cf374..0000000000000 --- a/x-pack/plugins/observability_ai_assistant/common/functions/visualize_esql.tsx +++ /dev/null @@ -1,27 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ -import { FromSchema } from 'json-schema-to-ts'; - -export const visualizeESQLFunction = { - name: 'visualize_query', - description: - 'Use this function to visualize charts for ES|QL queries. DO NOT run the lens function after.', - descriptionForUser: 'Use this function to visualize charts for ES|QL queries.', - parameters: { - type: 'object', - additionalProperties: true, - properties: { - query: { - type: 'string', - }, - }, - required: ['query'], - } as const, - contexts: ['core'], -}; - -export type VisualizeESQLFunctionArguments = FromSchema; diff --git a/x-pack/plugins/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx b/x-pack/plugins/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx index 458674b8f5f79..214dfe5e1bb78 100644 --- a/x-pack/plugins/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx +++ b/x-pack/plugins/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx @@ -54,21 +54,6 @@ export function EsqlCodeBlock({ - - - onActionClick({ type: ChatActionClickType.executeEsqlQuery, query: value }) - } - disabled={actionsDisabled} - > - {i18n.translate('xpack.observabilityAiAssistant.runThisQuery', { - defaultMessage: 'Run this query', - })} - - 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} /> ); } 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 fcab5fc074285..a1537ad60488b 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 @@ -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, diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/index.ts b/x-pack/plugins/observability_ai_assistant/server/functions/index.ts index 759aa8bb8bea6..6183d249efd19 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/index.ts +++ b/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[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.`; } diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/visualize_esql.ts b/x-pack/plugins/observability_ai_assistant/server/functions/visualize_esql.ts deleted file mode 100644 index d4a1f867e3868..0000000000000 --- a/x-pack/plugins/observability_ai_assistant/server/functions/visualize_esql.ts +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ -import { esFieldTypeToKibanaFieldType } from '@kbn/field-types'; -import type { ESQLSearchReponse } from '@kbn/es-types'; -import { visualizeESQLFunction } from '../../common/functions/visualize_esql'; -import type { FunctionRegistrationParameters } from '.'; - -export function registerVisualizeESQLFunction({ - client, - registerFunction, - resources, -}: FunctionRegistrationParameters) { - registerFunction( - visualizeESQLFunction, - 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: performantQuery, - }, - })) as ESQLSearchReponse; - const columns = - response.columns?.map(({ name, type }) => ({ - id: name, - name, - meta: { type: esFieldTypeToKibanaFieldType(type) }, - })) ?? []; - return { content: columns }; - } - ); -}