From 1a563002b70bb28f5f02d6ab28fe6a01ca12bbbd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Arturo=20Lidue=C3=B1a?= Date: Mon, 21 Apr 2025 10:51:18 +0200 Subject: [PATCH] [Obs AI Assistant]fixing error - Display results and Visualize query Bedrock Error (#218213) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Fix: Bedrock Streaming Error on ES|QL Actions #### Summary When an ES|QL is generated, we present two action buttons: - Visualize Query - Display Results These actions were not working as expected when using Bedrock as the model provider. #### Error Details ```txt Encountered error in Bedrock stream of type validationException messages.8: Did not find 1 `tool_result` block(s) at the beginning of this message. Messages following `tool_use` blocks must begin with a matching number of `tool_result` blocks. ``` #### Root Cause We were sending a tool_use block in the assistant message without immediately following it with the corresponding tool_result block. This violates Bedrock’s message protocol. (cherry picked from commit 33993b7123bc0d6c85d9c42b15610cc0d5092281) --- .../kbn-ai-assistant/src/chat/chat_body.tsx | 129 +++++++++--------- .../components/message_panel/message_text.tsx | 10 +- 2 files changed, 67 insertions(+), 72 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 03ae82dbd400a..73f460c33b620 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 @@ -285,75 +285,74 @@ export function ChatBody({ } }); - const handleActionClick = ({ - message, - payload, - }: { - message: Message; - payload: ChatActionClickPayload; - }) => { - setStickToBottom(true); - switch (payload.type) { - case ChatActionClickType.executeEsqlQuery: - 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, + const handleActionClick = useCallback( + ({ message, payload }: { message: Message; payload: ChatActionClickPayload }) => { + setStickToBottom(true); + switch (payload.type) { + case ChatActionClickType.executeEsqlQuery: + 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, + }, }, - }, - }) - ); - break; - - case ChatActionClickType.updateVisualization: - const visualizeQueryResponse = message; - - const visualizeQueryResponseData = JSON.parse(visualizeQueryResponse.message.data ?? '{}'); - - next( - messages.slice(0, messages.indexOf(visualizeQueryResponse)).concat({ - '@timestamp': new Date().toISOString(), - message: { - name: 'visualize_query', - content: visualizeQueryResponse.message.content, - data: JSON.stringify({ - ...visualizeQueryResponseData, - userOverrides: payload.userOverrides, - }), - role: MessageRole.User, - }, - }) - ); - break; - case ChatActionClickType.visualizeEsqlQuery: - next( - messages.concat({ - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.Assistant, - content: '', - function_call: { + }) + ); + break; + + case ChatActionClickType.updateVisualization: + const visualizeQueryResponse = message; + + const visualizeQueryResponseData = JSON.parse( + visualizeQueryResponse.message.data ?? '{}' + ); + + next( + messages.slice(0, messages.indexOf(visualizeQueryResponse)).concat({ + '@timestamp': new Date().toISOString(), + message: { name: 'visualize_query', - arguments: JSON.stringify({ - query: payload.query, - intention: VisualizeESQLUserIntention.visualizeAuto, + content: visualizeQueryResponse.message.content, + data: JSON.stringify({ + ...visualizeQueryResponseData, + userOverrides: payload.userOverrides, }), - trigger: MessageRole.User, + role: MessageRole.User, }, - }, - }) - ); - break; - } - }; + }) + ); + break; + case ChatActionClickType.visualizeEsqlQuery: + 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, + }, + }, + }) + ); + break; + } + }, + [messages, next] + ); const handleConversationAccessUpdate = async (access: ConversationAccess) => { await updateConversationAccess(access); diff --git a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx index 7565cfee2e1a8..a4e652eda80cd 100644 --- a/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx +++ b/x-pack/platform/plugins/shared/observability_ai_assistant/public/components/message_panel/message_text.tsx @@ -18,7 +18,7 @@ import { import { css } from '@emotion/css'; import classNames from 'classnames'; import type { Code, InlineCode, Parent, Text } from 'mdast'; -import React, { useMemo, useRef } from 'react'; +import React, { useMemo } from 'react'; import type { Node } from 'unist'; import { ChatActionClickHandler } from '../chat/types'; import { CodeBlock, EsqlCodeBlock } from './esql_code_block'; @@ -120,10 +120,6 @@ export function MessageText({ loading, content, onActionClick }: Props) { overflow-wrap: anywhere; `; - const onActionClickRef = useRef(onActionClick); - - onActionClickRef.current = onActionClick; - const { parsingPluginList, processingPluginList } = useMemo(() => { const parsingPlugins = getDefaultEuiMarkdownParsingPlugins(); @@ -149,7 +145,7 @@ export function MessageText({ loading, content, onActionClick }: Props) { value={props.value} lang={props.lang} actionsDisabled={loading} - onActionClick={onActionClickRef.current} + onActionClick={onActionClick} /> @@ -187,7 +183,7 @@ export function MessageText({ loading, content, onActionClick }: Props) { parsingPluginList: [loadingCursorPlugin, esqlLanguagePlugin, ...parsingPlugins], processingPluginList: processingPlugins, }; - }, [loading]); + }, [loading, onActionClick]); return (