[8.15] [Bedrock chat] Fixes for BedrockChat model + tools (#189227) (#…

…189245) # Backport This will backport the following commits from `main` to `8.15`: - [[Bedrock chat] Fixes for BedrockChat model + tools (#189227)](#189227)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Steph Milovic <[email protected]>
elastic · Jul 26, 2024 · 7c31f6c · 7c31f6c
1 parent a572cea
commit 7c31f6c
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 14 deletions.
diff --git a/x-pack/packages/kbn-langchain/server/language_models/bedrock_chat.ts b/x-pack/packages/kbn-langchain/server/language_models/bedrock_chat.ts
@@ -51,7 +51,13 @@ export class ActionsClientBedrockChatModel extends _BedrockChat {
                 anthropicVersion: inputBody.anthropic_version,
               },
             },
-          })) as { data: Readable };
+          })) as { data: Readable; status: string; message?: string; serviceMessage?: string };
+
+          if (data.status === 'error') {
+            throw new Error(
+              `ActionsClientBedrockChat: action result status is error: ${data?.message} - ${data?.serviceMessage}`
+            );
+          }
 
           return {
             body: Readable.toWeb(data.data),
@@ -72,7 +78,18 @@ export class ActionsClientBedrockChatModel extends _BedrockChat {
               anthropicVersion: inputBody.anthropic_version,
             },
           },
-        })) as { status: string; data: { message: string } };
+        })) as {
+          status: string;
+          data: { message: string };
+          message?: string;
+          serviceMessage?: string;
+        };
+
+        if (data.status === 'error') {
+          throw new Error(
+            `ActionsClientBedrockChat: action result status is error: ${data?.message} - ${data?.serviceMessage}`
+          );
+        }
 
         return {
           ok: data.status === 'ok',

diff --git a/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts b/...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
@@ -14,12 +14,16 @@ import {
   createToolCallingAgent,
 } from 'langchain/agents';
 import { APMTracer } from '@kbn/langchain/server/tracers/apm';
-import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { getLlmClass } from '../../../../routes/utils';
 import { EsAnonymizationFieldsSchema } from '../../../../ai_assistant_data_clients/anonymization_fields/types';
 import { AssistantToolParams } from '../../../../types';
 import { AgentExecutor } from '../../executors/types';
-import { openAIFunctionAgentPrompt, structuredChatAgentPrompt } from './prompts';
+import {
+  bedrockToolCallingAgentPrompt,
+  geminiToolCallingAgentPrompt,
+  openAIFunctionAgentPrompt,
+  structuredChatAgentPrompt,
+} from './prompts';
 import { getDefaultAssistantGraph } from './graph';
 import { invokeGraph, streamGraph } from './helpers';
 import { transformESSearchToAnonymizationFields } from '../../../../ai_assistant_data_clients/anonymization_fields/helpers';
@@ -126,16 +130,8 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     ? createToolCallingAgent({
         llm,
         tools,
-        prompt: ChatPromptTemplate.fromMessages([
-          [
-            'system',
-            'You are a helpful assistant. ALWAYS use the provided tools. Use tools as often as possible, as they have access to the latest data and syntax.\n\n' +
-              `The final response will be the only output the user sees and should be a complete answer to the user's question, as if you were responding to the user's initial question, which is "{input}". The final response should never be empty.`,
-          ],
-          ['placeholder', '{chat_history}'],
-          ['human', '{input}'],
-          ['placeholder', '{agent_scratchpad}'],
-        ]),
+        prompt:
+          llmType === 'bedrock' ? bedrockToolCallingAgentPrompt : geminiToolCallingAgentPrompt,
         streamRunnable: isStream,
       })
     : await createStructuredChatAgent({

diff --git a/.../plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts b/.../plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts
@@ -14,6 +14,27 @@ export const openAIFunctionAgentPrompt = ChatPromptTemplate.fromMessages([
   ['placeholder', '{agent_scratchpad}'],
 ]);
 
+export const bedrockToolCallingAgentPrompt = ChatPromptTemplate.fromMessages([
+  [
+    'system',
+    'You are a helpful assistant. ALWAYS use the provided tools. Use tools as often as possible, as they have access to the latest data and syntax.',
+  ],
+  ['placeholder', '{chat_history}'],
+  ['human', '{input}'],
+  ['placeholder', '{agent_scratchpad}'],
+]);
+
+export const geminiToolCallingAgentPrompt = ChatPromptTemplate.fromMessages([
+  [
+    'system',
+    'You are a helpful assistant. ALWAYS use the provided tools. Use tools as often as possible, as they have access to the latest data and syntax.\n\n' +
+      `The final response will be the only output the user sees and should be a complete answer to the user's question, as if you were responding to the user's initial question, which is "{input}". The final response should never be empty.`,
+  ],
+  ['placeholder', '{chat_history}'],
+  ['human', '{input}'],
+  ['placeholder', '{agent_scratchpad}'],
+]);
+
 export const structuredChatAgentPrompt = ChatPromptTemplate.fromMessages([
   [
     'system',