From 75d35d214329b2b15d6535e6b0e6a0065ae9f8f2 Mon Sep 17 00:00:00 2001 From: Viduni Wickramarachchi Date: Fri, 11 Jul 2025 08:25:04 -0400 Subject: [PATCH] [Obs AI Assistant] Inject user prompt before tool call when query actions are clicked (#227462) Closes https://github.com/elastic/obs-ai-assistant-team/issues/309 ## Summary This PR improves the `Display Results` and `visualize this query` actions by injecting a user prompt first and then the tool call to maintain correct message ordering. This fixes the empty message that was observed when visualize this query was clicked. However, in any case, if an empty message is received this should be handled separately to avoid breaking the conversation flow. ## Testing instructions: 1. Generate some logs 2. Pick the LLM (Needs to be tested with Elastic Managed LLM, Claude 3.5, Claude 3.7, GPT-4, etc) 3. Ask the assistant to generate a query 4. Click on `Display results` -- the results should be shown and the user should be able to continue the conversation 5. Click on `Visualize this query` -- the visualization should be shown and the user should be able to continue the conversation ## Traces with each model: - EIS (Visualize this query) - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/b1dcb77fb39ad7efbf82d9c4a1feb0e0?selectedSpanNodeId=U3BhbjoyOTk3ODQ%3D image - EIS (Display results) - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/e7d8c42774491e6e2bcf328c34203c14?selectedSpanNodeId=U3BhbjozMDAxOTg%3D https://github.com/user-attachments/assets/6aa520db-2b40-4317-ace0-7be206ebb669 - Bedrock+Claude 3.7 - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/f1bc4f490fcc990db9dbddf0782cff7e?selectedSpanNodeId=U3BhbjoyOTk4MDA%3D image - Bedrock+Claude 3.5 - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/457a7cc939ce3b18ae3a9e2b066c3a7f?selectedSpanNodeId=U3BhbjoyOTk4MTY%3D image - AzureOpenAI+GPT 4.1 - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/de9fc09a6f0162b7b2d6f52fbb188469?selectedSpanNodeId=U3BhbjoyOTk4MzI%3D image - Gemini 2.0 Flash - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/982da6b0ff4bae04f45af0271aa4cd84?selectedSpanNodeId=U3BhbjoyOTk4NTI%3D image - Gemini 2.5 Flash - https://35-187-109-62.sslip.io/projects/UHJvamVjdDo5/traces/9f7115cc79bb2e929dbeb93670555e81?selectedSpanNodeId=U3BhbjoyOTk4NjU%3D image ### Checklist - [x] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) - [x] Review the [backport guidelines](https://docs.google.com/document/d/1VyN5k91e5OVumlc0Gb9RPa3h1ewuPE705nRtioPiTvY/edit?usp=sharing) and apply applicable `backport:*` labels. (cherry picked from commit 021b30816e86ddb2e1fa2a8c911ff85ab07e1d43) --- .../kbn-ai-assistant/src/chat/chat_body.tsx | 76 ++++++++++++------- .../server/functions/visualize_esql.ts | 16 +++- 2 files changed, 62 insertions(+), 30 deletions(-) diff --git a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx index 8dedc7d77c600..eb4adcb13bde9 100644 --- a/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx +++ b/x-pack/platform/packages/shared/kbn-ai-assistant/src/chat/chat_body.tsx @@ -290,24 +290,35 @@ export function ChatBody({ ({ message, payload }: { message: Message; payload: ChatActionClickPayload }) => { setStickToBottom(true); switch (payload.type) { - case ChatActionClickType.executeEsqlQuery: + case ChatActionClickType.executeEsqlQuery: { + const now = new Date().toISOString(); next( - messages.concat({ - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.Assistant, - content: '', - function_call: { - name: 'execute_query', - arguments: JSON.stringify({ - query: payload.query, - }), - trigger: MessageRole.User, + messages.concat([ + { + '@timestamp': now, + message: { + role: MessageRole.User, + content: `Display results for the following ES|QL query:\n\n\`\`\`esql\n${payload.query}\n\`\`\``, }, }, - }) + { + '@timestamp': now, + message: { + role: MessageRole.Assistant, + content: '', + function_call: { + name: 'execute_query', + arguments: JSON.stringify({ + query: payload.query, + }), + trigger: MessageRole.User, + }, + }, + }, + ]) ); break; + } case ChatActionClickType.updateVisualization: const visualizeQueryResponse = message; @@ -331,25 +342,36 @@ export function ChatBody({ }) ); break; - case ChatActionClickType.visualizeEsqlQuery: + case ChatActionClickType.visualizeEsqlQuery: { + const now = new Date().toISOString(); next( - messages.concat({ - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.Assistant, - content: '', - function_call: { - name: 'visualize_query', - arguments: JSON.stringify({ - query: payload.query, - intention: VisualizeESQLUserIntention.visualizeAuto, - }), - trigger: MessageRole.User, + messages.concat([ + { + '@timestamp': now, + message: { + role: MessageRole.User, + content: `Visualize the following ES|QL query:\n\n\`\`\`esql\n${payload.query}\n\`\`\``, }, }, - }) + { + '@timestamp': now, + message: { + role: MessageRole.Assistant, + content: '', + function_call: { + name: 'visualize_query', + arguments: JSON.stringify({ + query: payload.query, + intention: VisualizeESQLUserIntention.visualizeAuto, + }), + trigger: MessageRole.User, + }, + }, + }, + ]) ); break; + } } }, [messages, next] diff --git a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts index e7a4e8c0dc029..c96db9a713288 100644 --- a/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts +++ b/x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/visualize_esql.ts @@ -21,10 +21,20 @@ const getMessageForLLM = ( if (hasErrors) { return 'The query has syntax errors'; } - return intention === VisualizeESQLUserIntention.executeAndReturnResults || + + if ( + intention === VisualizeESQLUserIntention.executeAndReturnResults || intention === VisualizeESQLUserIntention.generateQueryOnly - ? 'These results are not visualized' - : 'Only following query is visualized: ```esql\n' + query + '\n```'; + ) { + return 'These results are not visualized.'; + } + + // This message is added to avoid the model echoing the full ES|QL query back to the user. + // The UI already shows the chart. + return `Only the following query is visualized: \`\`\`esql\n' + ${query} + '\n\`\`\`\n + If the query is visualized once, don't attempt to visualize the same query again immediately. + After calling visualize_query you are done - **do NOT repeat the ES|QL query or add any further + explanation unless the user explicitly asks for it again.** Mention that the query is visualized.`; }; export function registerVisualizeESQLFunction({