[8.17] [Security solution] Use `BedrockRuntimeClient` to in…

…teract with converse APIs (elastic#201046) (elastic#201288) # Backport This will backport the following commits from `main` to `8.17`: - [[Security solution] Use `BedrockRuntimeClient` to interact with converse APIs (elastic#201046)](elastic#201046)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Steph Milovic <[email protected]>
crespocarlos · Nov 21, 2024 · 1dfae0f · 1dfae0f
1 parent d292cec
commit 1dfae0f
Show file tree

Hide file tree

Showing 13 changed files with 454 additions and 1,359 deletions.
diff --git a/...ages/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts b/...ages/kbn-langchain/server/language_models/chat_bedrock_converse/bedrock_runtime_client.ts
@@ -8,12 +8,17 @@
 import {
   BedrockRuntimeClient as _BedrockRuntimeClient,
   BedrockRuntimeClientConfig,
+  ConverseCommand,
+  ConverseResponse,
+  ConverseStreamCommand,
+  ConverseStreamResponse,
 } from '@aws-sdk/client-bedrock-runtime';
 import { constructStack } from '@smithy/middleware-stack';
+import { HttpHandlerOptions } from '@smithy/types';
 import { PublicMethodsOf } from '@kbn/utility-types';
 import type { ActionsClient } from '@kbn/actions-plugin/server';
 
-import { NodeHttpHandler } from './node_http_handler';
+import { prepareMessages } from '../../utils/bedrock';
 
 export interface CustomChatModelInput extends BedrockRuntimeClientConfig {
   actionsClient: PublicMethodsOf<ActionsClient>;
@@ -23,15 +28,51 @@ export interface CustomChatModelInput extends BedrockRuntimeClientConfig {
 
 export class BedrockRuntimeClient extends _BedrockRuntimeClient {
   middlewareStack: _BedrockRuntimeClient['middlewareStack'];
+  streaming: boolean;
+  actionsClient: PublicMethodsOf<ActionsClient>;
+  connectorId: string;
 
   constructor({ actionsClient, connectorId, ...fields }: CustomChatModelInput) {
     super(fields ?? {});
-    this.config.requestHandler = new NodeHttpHandler({
-      streaming: fields.streaming ?? true,
-      actionsClient,
-      connectorId,
-    });
+    this.streaming = fields.streaming ?? true;
+    this.actionsClient = actionsClient;
+    this.connectorId = connectorId;
     // eliminate middleware steps that handle auth as Kibana connector handles auth
     this.middlewareStack = constructStack() as _BedrockRuntimeClient['middlewareStack'];
   }
+
+  public async send(
+    command: ConverseCommand | ConverseStreamCommand,
+    optionsOrCb?: HttpHandlerOptions | ((err: unknown, data: unknown) => void)
+  ) {
+    const options = typeof optionsOrCb !== 'function' ? optionsOrCb : {};
+    if (command.input.messages) {
+      // without this, our human + human messages do not work and result in error:
+      // A conversation must alternate between user and assistant roles.
+      command.input.messages = prepareMessages(command.input.messages);
+    }
+    const data = (await this.actionsClient.execute({
+      actionId: this.connectorId,
+      params: {
+        subAction: 'bedrockClientSend',
+        subActionParams: {
+          command,
+          signal: options?.abortSignal,
+        },
+      },
+    })) as {
+      data: ConverseResponse | ConverseStreamResponse;
+      status: string;
+      message?: string;
+      serviceMessage?: string;
+    };
+
+    if (data.status === 'error') {
+      throw new Error(
+        `ActionsClient BedrockRuntimeClient: action result status is error: ${data?.message} - ${data?.serviceMessage}`
+      );
+    }
+
+    return data.data;
+  }
 }
diff --git a/...ages/kbn-langchain/server/language_models/chat_bedrock_converse/node_http_handler.test.ts b/...ages/kbn-langchain/server/language_models/chat_bedrock_converse/node_http_handler.test.ts
diff --git a/.../packages/kbn-langchain/server/language_models/chat_bedrock_converse/node_http_handler.ts b/.../packages/kbn-langchain/server/language_models/chat_bedrock_converse/node_http_handler.ts
diff --git a/x-pack/packages/kbn-langchain/server/utils/bedrock.ts b/x-pack/packages/kbn-langchain/server/utils/bedrock.ts
@@ -10,6 +10,7 @@ import { finished } from 'stream/promises';
 import { Logger } from '@kbn/core/server';
 import { EventStreamCodec } from '@smithy/eventstream-codec';
 import { fromUtf8, toUtf8 } from '@smithy/util-utf8';
+import { Message } from '@aws-sdk/client-bedrock-runtime';
 import { StreamParser } from './types';
 
 export const parseBedrockStreamAsAsyncIterator = async function* (
@@ -227,7 +228,7 @@ function parseContent(content: Array<{ text?: string; type: string }>): string {
  * Prepare messages for the bedrock API by combining messages from the same role
  * @param messages
  */
-export const prepareMessages = (messages: Array<{ role: string; content: string[] }>) =>
+export const prepareMessages = (messages: Message[]) =>
   messages.reduce((acc, { role, content }) => {
     const lastMessage = acc[acc.length - 1];
 
@@ -236,13 +237,13 @@ export const prepareMessages = (messages: Array<{ role: string; content: string[
       return acc;
     }
 
-    if (lastMessage.role === role) {
-      acc[acc.length - 1].content = lastMessage.content.concat(content);
+    if (lastMessage.role === role && lastMessage.content) {
+      acc[acc.length - 1].content = lastMessage.content.concat(content || []);
       return acc;
     }
 
     return acc;
-  }, [] as Array<{ role: string; content: string[] }>);
+  }, [] as Message[]);
 
 export const DEFAULT_BEDROCK_MODEL = 'anthropic.claude-3-5-sonnet-20240620-v1:0';
 export const DEFAULT_BEDROCK_REGION = 'us-east-1';