[Security solution] AI connector changes (elastic#170150)

joeypoon · Oct 30, 2023 · 75586c8 · 75586c8
1 parent f5f5338
commit 75586c8
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 0 deletions.
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
@@ -93,6 +93,8 @@ export class BedrockConnector extends SubActionConnector<Config, Secrets> {
         },
         body,
         path,
+        // Despite AWS docs, this value does not always get inferred. We need to always send it
+        service: 'bedrock',
       },
       {
         secretAccessKey: this.secrets.secret,

diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.test.ts
@@ -85,6 +85,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleOpenAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -102,6 +103,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(requestBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -118,6 +120,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleOpenAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -148,6 +151,7 @@ describe('OpenAIConnector', () => {
         });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -269,6 +273,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.invokeAI(sampleOpenAiBody);
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://api.openai.com/v1/chat/completions',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -379,6 +384,7 @@ describe('OpenAIConnector', () => {
         const response = await connector.runApi({ body: JSON.stringify(sampleAzureAiBody) });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://My-test-resource-123.openai.azure.com/openai/deployments/NEW-DEPLOYMENT-321/chat/completions?api-version=2023-05-15',
           method: 'post',
           responseSchema: RunActionResponseSchema,
@@ -405,6 +411,7 @@ describe('OpenAIConnector', () => {
         });
         expect(mockRequest).toBeCalledTimes(1);
         expect(mockRequest).toHaveBeenCalledWith({
+          timeout: 120000,
           url: 'https://My-test-resource-123.openai.azure.com/openai/deployments/NEW-DEPLOYMENT-321/chat/completions?api-version=2023-05-15',
           method: 'post',
           responseSchema: RunActionResponseSchema,

diff --git a/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts b/x-pack/plugins/stack_connectors/server/connector_types/openai/openai.ts
@@ -114,6 +114,8 @@ export class OpenAIConnector extends SubActionConnector<Config, Secrets> {
       method: 'post',
       responseSchema: RunActionResponseSchema,
       data: sanitizedBody,
+      // give up to 2 minutes for response
+      timeout: 120000,
       ...axiosOptions,
     });
     return response.data;
-Original file line number
+Diff line change
@@ Expand Up @@
             },
             body,
             path,
+            // Despite AWS docs, this value does not always get inferred. We need to always send it
+            service: 'bedrock',
           },
           {
             secretAccessKey: this.secrets.secret,
@@ Expand Down @@