[Security solution] Bedrock connector update (elastic#168145)

dej611 · Oct 17, 2023 · 9aef9c4 · 9aef9c4
1 parent 691cc54
commit 9aef9c4
Show file tree

Hide file tree

Showing 8 changed files with 90 additions and 9 deletions.
diff --git a/x-pack/plugins/actions/server/lib/axios_utils.test.ts b/x-pack/plugins/actions/server/lib/axios_utils.test.ts
@@ -260,6 +260,66 @@ describe('request', () => {
       data: { incidentId: '123' },
     });
   });
+
+  test('it uses timeout argument when one is provided', async () => {
+    configurationUtilities.getProxySettings.mockReturnValue({
+      proxySSLSettings: {
+        verificationMode: 'full',
+      },
+      proxyUrl: 'https://elastic.proxy.co',
+      proxyBypassHosts: new Set(['elastic.co']),
+      proxyOnlyHosts: undefined,
+    });
+
+    await request({
+      axios,
+      url: TestUrl,
+      logger,
+      configurationUtilities,
+    });
+
+    await request({
+      axios,
+      url: TestUrl,
+      logger,
+      configurationUtilities,
+      timeout: 55,
+    });
+
+    expect(axiosMock.mock.calls.length).toBe(2);
+    expect(axiosMock.mock.calls[0][1].timeout).toBe(360000);
+    expect(axiosMock.mock.calls[1][1].timeout).toBe(360000);
+  });
+
+  test('it uses timeout argument when one is provided that is greater than settings timeout', async () => {
+    configurationUtilities.getProxySettings.mockReturnValue({
+      proxySSLSettings: {
+        verificationMode: 'full',
+      },
+      proxyUrl: 'https://elastic.proxy.co',
+      proxyBypassHosts: new Set(['elastic.co']),
+      proxyOnlyHosts: undefined,
+    });
+
+    await request({
+      axios,
+      url: TestUrl,
+      logger,
+      configurationUtilities,
+    });
+
+    await request({
+      axios,
+      url: TestUrl,
+      logger,
+      configurationUtilities,
+      timeout: 360001,
+    });
+
+    expect(axiosMock.mock.calls.length).toBe(2);
+    expect(axiosMock.mock.calls[0][1].timeout).toBe(360000);
+    expect(axiosMock.mock.calls[1][1].timeout).toBe(360001);
+  });
 });
 
 describe('patch', () => {

diff --git a/x-pack/plugins/actions/server/lib/axios_utils.ts b/x-pack/plugins/actions/server/lib/axios_utils.ts
@@ -28,6 +28,7 @@ export const request = async <T = unknown>({
   configurationUtilities,
   headers,
   sslOverrides,
+  timeout,
   ...config
 }: {
   axios: AxiosInstance;
@@ -37,6 +38,7 @@ export const request = async <T = unknown>({
   data?: T;
   configurationUtilities: ActionsConfigurationUtilities;
   headers?: Record<string, AxiosHeaderValue>;
+  timeout?: number;
   sslOverrides?: SSLSettings;
 } & AxiosRequestConfig): Promise<AxiosResponse> => {
   const { httpAgent, httpsAgent } = getCustomAgents(
@@ -45,7 +47,8 @@ export const request = async <T = unknown>({
     url,
     sslOverrides
   );
-  const { maxContentLength, timeout } = configurationUtilities.getResponseSettings();
+  const { maxContentLength, timeout: settingsTimeout } =
+    configurationUtilities.getResponseSettings();
 
   return await axios(url, {
     ...config,
@@ -57,7 +60,7 @@ export const request = async <T = unknown>({
     httpsAgent,
     proxy: false,
     maxContentLength,
-    timeout,
+    timeout: Math.max(settingsTimeout, timeout ?? 0),
   });
 };
 

diff --git a/x-pack/plugins/actions/server/sub_action_framework/sub_action_connector.ts b/x-pack/plugins/actions/server/sub_action_framework/sub_action_connector.ts
@@ -118,6 +118,7 @@ export abstract class SubActionConnector<Config, Secrets> {
     method = 'get',
     responseSchema,
     headers,
+    timeout,
     ...config
   }: SubActionRequestParams<R>): Promise<AxiosResponse<R>> {
     try {
@@ -138,6 +139,7 @@ export abstract class SubActionConnector<Config, Secrets> {
         data: this.normalizeData(data),
         configurationUtilities: this.configurationUtilities,
         headers: this.getHeaders(headers as AxiosHeaders),
+        timeout,
       });
 
       this.validateResponse(responseSchema, res.data);

diff --git a/x-pack/plugins/stack_connectors/common/bedrock/constants.ts b/x-pack/plugins/stack_connectors/common/bedrock/constants.ts
@@ -20,6 +20,7 @@ export enum SUB_ACTION {
   TEST = 'test',
 }
 
+export const DEFAULT_TOKEN_LIMIT = 8191;
 export const DEFAULT_BEDROCK_MODEL = 'anthropic.claude-v2';
 
 export const DEFAULT_BEDROCK_URL = `https://bedrock.us-east-1.amazonaws.com` as const;
diff --git a/x-pack/plugins/stack_connectors/public/connector_types/bedrock/constants.tsx b/x-pack/plugins/stack_connectors/public/connector_types/bedrock/constants.tsx
@@ -9,15 +9,19 @@ import React from 'react';
 import { ConfigFieldSchema, SecretsFieldSchema } from '@kbn/triggers-actions-ui-plugin/public';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { EuiLink } from '@elastic/eui';
-import { DEFAULT_BEDROCK_MODEL, DEFAULT_BEDROCK_URL } from '../../../common/bedrock/constants';
+import {
+  DEFAULT_BEDROCK_MODEL,
+  DEFAULT_BEDROCK_URL,
+  DEFAULT_TOKEN_LIMIT,
+} from '../../../common/bedrock/constants';
 import * as i18n from './translations';
 
 const human = '\n\nHuman:';
 const assistant = '\n\nAssistant:';
 
 export const DEFAULT_BODY = JSON.stringify({
   prompt: `${human} Hello world! ${assistant}`,
-  max_tokens_to_sample: 300,
+  max_tokens_to_sample: DEFAULT_TOKEN_LIMIT,
   stop_sequences: [human],
 });
 

diff --git a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.test.ts
@@ -14,6 +14,7 @@ import {
   BEDROCK_CONNECTOR_ID,
   DEFAULT_BEDROCK_MODEL,
   DEFAULT_BEDROCK_URL,
+  DEFAULT_TOKEN_LIMIT,
 } from '../../../common/bedrock/constants';
 import { DEFAULT_BODY } from '../../../public/connector_types/bedrock/constants';
 import { AxiosError } from 'axios';
@@ -65,6 +66,7 @@ describe('BedrockConnector', () => {
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
           signed: true,
+          timeout: 120000,
           url: `${DEFAULT_BEDROCK_URL}/model/${DEFAULT_BEDROCK_MODEL}/invoke`,
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -96,12 +98,14 @@ describe('BedrockConnector', () => {
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
           signed: true,
+          timeout: 120000,
           url: `${DEFAULT_BEDROCK_URL}/model/${DEFAULT_BEDROCK_MODEL}/invoke`,
           method: 'post',
           responseSchema: RunActionResponseSchema,
           data: JSON.stringify({
             prompt: '\n\nHuman:Hello world \n\nAssistant:',
-            max_tokens_to_sample: 300,
+            max_tokens_to_sample: DEFAULT_TOKEN_LIMIT,
+            temperature: 0.5,
             stop_sequences: ['\n\nHuman:'],
           }),
         });
@@ -132,13 +136,15 @@ describe('BedrockConnector', () => {
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
           signed: true,
+          timeout: 120000,
           url: `${DEFAULT_BEDROCK_URL}/model/${DEFAULT_BEDROCK_MODEL}/invoke`,
           method: 'post',
           responseSchema: RunActionResponseSchema,
           data: JSON.stringify({
             prompt:
               '\n\nHuman:Hello world\n\nHuman:Be a good chatbot\n\nAssistant:Hi, I am a good chatbot\n\nHuman:What is 2+2? \n\nAssistant:',
-            max_tokens_to_sample: 300,
+            max_tokens_to_sample: DEFAULT_TOKEN_LIMIT,
+            temperature: 0.5,
             stop_sequences: ['\n\nHuman:'],
           }),
         });

diff --git a/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts b/x-pack/plugins/stack_connectors/server/connector_types/bedrock/bedrock.ts
@@ -21,7 +21,7 @@ import type {
   InvokeAIActionParams,
   InvokeAIActionResponse,
 } from '../../../common/bedrock/types';
-import { SUB_ACTION } from '../../../common/bedrock/constants';
+import { SUB_ACTION, DEFAULT_TOKEN_LIMIT } from '../../../common/bedrock/constants';
 
 interface SignedRequest {
   host: string;
@@ -116,6 +116,8 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
       method: 'post',
       responseSchema: RunActionResponseSchema,
       data: body,
+      // give up to 2 minutes for response
+      timeout: 120000,
     });
     return response.data;
   }
@@ -141,7 +143,8 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
     const req = {
       // end prompt in "Assistant:" to avoid the model starting its message with "Assistant:"
       prompt: `${combinedMessages} \n\nAssistant:`,
-      max_tokens_to_sample: 300,
+      max_tokens_to_sample: DEFAULT_TOKEN_LIMIT,
+      temperature: 0.5,
       // prevent model from talking to itself
       stop_sequences: ['\n\nHuman:'],
     };

diff --git a/...rting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts b/...rting_api_integration/security_and_spaces/group2/tests/actions/connector_types/bedrock.ts
@@ -11,6 +11,7 @@ import {
   BedrockSimulator,
   bedrockSuccessResponse,
 } from '@kbn/actions-simulators-plugin/server/bedrock_simulation';
+import { DEFAULT_TOKEN_LIMIT } from '@kbn/stack-connectors-plugin/common/bedrock/constants';
 import { FtrProviderContext } from '../../../../../common/ftr_provider_context';
 import { getUrlPrefix, ObjectRemover } from '../../../../../common/lib';
 
@@ -396,7 +397,8 @@ export default function bedrockTest({ getService }: FtrProviderContext) {
             expect(simulator.requestData).to.eql({
               prompt:
                 '\n\nHuman:Hello world\n\nHuman:Be a good chatbot\n\nAssistant:Hi, I am a good chatbot\n\nHuman:What is 2+2? \n\nAssistant:',
-              max_tokens_to_sample: 300,
+              max_tokens_to_sample: DEFAULT_TOKEN_LIMIT,
+              temperature: 0.5,
               stop_sequences: ['\n\nHuman:'],
             });
             expect(body).to.eql({