[Security solution] AI Assistant, replace LLM with SimpleChatModel + Bedrock streaming

.eslintrc.js

@@ -1058,6 +1058,7 @@ module.exports = {
         'x-pack/plugins/elastic_assistant/**/*.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant/**/*.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant-common/**/*.{ts,tsx}',
+        'x-pack/packages/kbn-langchain/**/*.{ts,tsx}',
         'x-pack/packages/security-solution/**/*.{ts,tsx}',
         'x-pack/plugins/security_solution/**/*.{ts,tsx}',
         'x-pack/plugins/security_solution_ess/**/*.{ts,tsx}',
@@ -1071,6 +1072,7 @@ module.exports = {
         'x-pack/plugins/elastic_assistant/**/*.{test,mock,test_helper}.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant/**/*.{test,mock,test_helper}.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant-common/**/*.{test,mock,test_helper}.{ts,tsx}',
+        'x-pack/packages/kbn-langchain/**/*.{test,mock,test_helper}.{ts,tsx}',
         'x-pack/packages/security-solution/**/*.{test,mock,test_helper}.{ts,tsx}',
         'x-pack/plugins/security_solution/**/*.{test,mock,test_helper}.{ts,tsx}',
         'x-pack/plugins/security_solution_ess/**/*.{test,mock,test_helper}.{ts,tsx}',
@@ -1090,6 +1092,7 @@ module.exports = {
         'x-pack/plugins/elastic_assistant/**/*.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant/**/*.{ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant-common/**/*.{ts,tsx}',
+        'x-pack/packages/kbn-langchain/**/*.{ts,tsx}',
         'x-pack/packages/security-solution/**/*.{ts,tsx}',
         'x-pack/plugins/security_solution/**/*.{ts,tsx}',
         'x-pack/plugins/security_solution_ess/**/*.{ts,tsx}',
@@ -1128,6 +1131,7 @@ module.exports = {
         'x-pack/plugins/elastic_assistant/**/*.{js,mjs,ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant/**/*.{js,mjs,ts,tsx}',
         'x-pack/packages/kbn-elastic-assistant-common/**/*.{js,mjs,ts,tsx}',
+        'x-pack/packages/kbn-langchain/**/*.{js,mjs,ts,tsx}',
         'x-pack/packages/security-solution/**/*.{js,mjs,ts,tsx}',
         'x-pack/plugins/security_solution/**/*.{js,mjs,ts,tsx}',
         'x-pack/plugins/security_solution_ess/**/*.{js,mjs,ts,tsx}',

.github/CODEOWNERS

@@ -517,6 +517,7 @@ src/plugins/kibana_react @elastic/appex-sharedux
 src/plugins/kibana_usage_collection @elastic/kibana-core
 src/plugins/kibana_utils @elastic/appex-sharedux
 x-pack/plugins/kubernetes_security @elastic/kibana-cloud-security-posture
+x-pack/packages/kbn-langchain @elastic/security-generative-ai
 packages/kbn-language-documentation-popover @elastic/kibana-esql
 x-pack/examples/lens_config_builder_example @elastic/kibana-visualizations
 packages/kbn-lens-embeddable-utils @elastic/obs-ux-infra_services-team @elastic/kibana-visualizations

package.json

@@ -544,6 +544,7 @@
     "@kbn/kibana-usage-collection-plugin": "link:src/plugins/kibana_usage_collection",
     "@kbn/kibana-utils-plugin": "link:src/plugins/kibana_utils",
     "@kbn/kubernetes-security-plugin": "link:x-pack/plugins/kubernetes_security",
+    "@kbn/langchain": "link:x-pack/packages/kbn-langchain",
     "@kbn/language-documentation-popover": "link:packages/kbn-language-documentation-popover",
     "@kbn/lens-config-builder-example-plugin": "link:x-pack/examples/lens_config_builder_example",
     "@kbn/lens-embeddable-utils": "link:packages/kbn-lens-embeddable-utils",

tsconfig.base.json

@@ -1028,6 +1028,8 @@
       "@kbn/kibana-utils-plugin/*": ["src/plugins/kibana_utils/*"],
       "@kbn/kubernetes-security-plugin": ["x-pack/plugins/kubernetes_security"],
       "@kbn/kubernetes-security-plugin/*": ["x-pack/plugins/kubernetes_security/*"],
+      "@kbn/langchain": ["x-pack/packages/kbn-langchain"],
+      "@kbn/langchain/*": ["x-pack/packages/kbn-langchain/*"],
       "@kbn/language-documentation-popover": ["packages/kbn-language-documentation-popover"],
       "@kbn/language-documentation-popover/*": ["packages/kbn-language-documentation-popover/*"],
       "@kbn/lens-config-builder-example-plugin": ["x-pack/examples/lens_config_builder_example"],

x-pack/packages/kbn-elastic-assistant-common/impl/utils/bedrock.ts

@@ -5,10 +5,30 @@
  * 2.0.
  */
 
-import { Logger } from '@kbn/core/server';
+import { Logger } from '@kbn/logging';
 import { EventStreamCodec } from '@smithy/eventstream-codec';
 import { fromUtf8, toUtf8 } from '@smithy/util-utf8';
 
+/**
+ * Parses a Bedrock buffer from an array of chunks.
+ *
+ * @param {Uint8Array[]} chunks - Array of Uint8Array chunks to be parsed.
+ * @returns {string} - Parsed string from the Bedrock buffer.
+ */
+export const parseBedrockBuffer = (chunks: Uint8Array[], logger: Logger): string => {
+  // Initialize an empty Uint8Array to store the concatenated buffer.
+  let bedrockBuffer: Uint8Array = new Uint8Array(0);
+
+  // Map through each chunk to process the Bedrock buffer.
+  return chunks
+    .map((chunk) => {
+      const processedChunk = handleBedrockChunk({ chunk, bedrockBuffer, logger });
+      bedrockBuffer = processedChunk.bedrockBuffer;
+      return processedChunk.decodedChunk;
+    })
+    .join('');
+};
+
 /**
  * Handle a chunk of data from the bedrock API.
  * @param chunk - The chunk of data to process.
@@ -55,7 +75,9 @@ export const handleBedrockChunk = ({
         Buffer.from(JSON.parse(new TextDecoder().decode(event.body)).bytes, 'base64').toString()
       );
       const decodedContent = prepareBedrockOutput(body, logger);
-      if (chunkHandler) chunkHandler(decodedContent);
+      if (chunkHandler) {
+        chunkHandler(decodedContent);
+      }
       return decodedContent;
     })
     .join('');

x-pack/packages/kbn-elastic-assistant-common/index.ts

@@ -23,6 +23,5 @@ export {
 } from './impl/data_anonymization/helpers';
 
 export { transformRawData } from './impl/data_anonymization/transform_raw_data';
-export { handleBedrockChunk } from './impl/utils/bedrock';
-
+export { parseBedrockBuffer, handleBedrockChunk } from './impl/utils/bedrock';
 export * from './constants';

x-pack/packages/kbn-elastic-assistant-common/tsconfig.json

@@ -19,7 +19,6 @@
     "@kbn/zod-helpers",
     "@kbn/securitysolution-io-ts-utils",
     "@kbn/core",
-    "@kbn/actions-plugin",
-    "@kbn/logging-mocks",
+    "@kbn/logging",
   ]
 }

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.test.tsx

@@ -24,7 +24,7 @@ const mockHttp = {
   fetch: jest.fn(),
 } as unknown as HttpSetup;
 
-const apiConfig: Record<'openai' | 'bedrock', ApiConfig> = {
+const apiConfig: Record<'openai' | 'bedrock' | 'gemini', ApiConfig> = {
   openai: {
     connectorId: 'foo',
     actionTypeId: '.gen-ai',
@@ -35,6 +35,10 @@ const apiConfig: Record<'openai' | 'bedrock', ApiConfig> = {
     connectorId: 'foo',
     actionTypeId: '.bedrock',
   },
+  gemini: {
+    connectorId: 'foo',
+    actionTypeId: '.gemini',
+  },
 };
 
 const fetchConnectorArgs: FetchConnectorExecuteAction = {
@@ -94,7 +98,7 @@ describe('API tests', () => {
       );
     });
 
-    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is true', async () => {
+    it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is true', async () => {
       const testProps: FetchConnectorExecuteAction = {
         ...fetchConnectorArgs,
         apiConfig: apiConfig.bedrock,
@@ -105,13 +109,13 @@ describe('API tests', () => {
       expect(mockHttp.fetch).toHaveBeenCalledWith(
         '/internal/elastic_assistant/actions/connector/foo/_execute',
         {
-          ...staticDefaults,
-          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
+          ...streamingDefaults,
+          body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
         }
       );
     });
 
-    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
+    it('calls the stream API when assistantStreamingEnabled is true and actionTypeId is bedrock and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
       const testProps: FetchConnectorExecuteAction = {
         ...fetchConnectorArgs,
         apiConfig: apiConfig.bedrock,
@@ -121,11 +125,47 @@ describe('API tests', () => {
 
       await fetchConnectorExecuteAction(testProps);
 
+      expect(mockHttp.fetch).toHaveBeenCalledWith(
+        '/internal/elastic_assistant/actions/connector/foo/_execute',
+        {
+          ...streamingDefaults,
+          body: '{"message":"This is a test","subAction":"invokeStream","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
+        }
+      );
+    });
+
+    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is true', async () => {
+      const testProps: FetchConnectorExecuteAction = {
+        ...fetchConnectorArgs,
+        apiConfig: apiConfig.gemini,
+      };
+
+      await fetchConnectorExecuteAction(testProps);
+
+      expect(mockHttp.fetch).toHaveBeenCalledWith(
+        '/internal/elastic_assistant/actions/connector/foo/_execute',
+        {
+          ...staticDefaults,
+          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":true,"isEnabledRAGAlerts":false}',
+        }
+      );
+    });
+
+    it('calls the non-stream API when assistantStreamingEnabled is true and actionTypeId is gemini and isEnabledKnowledgeBase is false and isEnabledRAGAlerts is true', async () => {
+      const testProps: FetchConnectorExecuteAction = {
+        ...fetchConnectorArgs,
+        apiConfig: apiConfig.gemini,
+        isEnabledKnowledgeBase: false,
+        isEnabledRAGAlerts: true,
+      };
+
+      await fetchConnectorExecuteAction(testProps);
+
       expect(mockHttp.fetch).toHaveBeenCalledWith(
         '/internal/elastic_assistant/actions/connector/foo/_execute',
         {
           ...staticDefaults,
-          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".bedrock","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
+          body: '{"message":"This is a test","subAction":"invokeAI","conversationId":"test","actionTypeId":".gemini","replacements":{},"isEnabledKnowledgeBase":false,"isEnabledRAGAlerts":true}',
         }
       );
     });

x-pack/packages/kbn-elastic-assistant/impl/assistant/api/index.tsx

@@ -63,12 +63,14 @@ export const fetchConnectorExecuteAction = async ({
   size,
   traceOptions,
 }: FetchConnectorExecuteAction): Promise<FetchConnectorExecuteResponse> => {
+  // TODO add streaming support for gemini with langchain on
   const isStream =
     assistantStreamingEnabled &&
     (apiConfig.actionTypeId === '.gen-ai' ||
-      // TODO add streaming support for bedrock with langchain on
+      apiConfig.actionTypeId === '.bedrock' ||
+      // TODO add streaming support for gemini with langchain on
       // tracked here: https://github.com/elastic/security-team/issues/7363
-      (apiConfig.actionTypeId === '.bedrock' && !isEnabledRAGAlerts && !isEnabledKnowledgeBase));
+      (apiConfig.actionTypeId === '.gemini' && !isEnabledRAGAlerts && !isEnabledKnowledgeBase));
 
   const optionalRequestParams = getOptionalRequestParams({
     isEnabledRAGAlerts,

x-pack/packages/kbn-elastic-assistant/impl/content/prompts/system/translations.ts

@@ -15,13 +15,6 @@ export const YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT = i18n.translate(
   }
 );
 
-export const USE_THE_FOLLOWING_CONTEXT_TO_ANSWER = i18n.translate(
-  'xpack.elasticAssistant.assistant.content.prompts.system.useTheFollowingContextToAnswer',
-  {
-    defaultMessage: 'Use the following context to answer questions:',
-  }
-);
-
 export const IF_YOU_DONT_KNOW_THE_ANSWER = i18n.translate(
   'xpack.elasticAssistant.assistant.content.prompts.system.ifYouDontKnowTheAnswer',
   {
@@ -37,8 +30,7 @@ export const SUPERHERO_PERSONALITY = i18n.translate(
   }
 );
 
-export const DEFAULT_SYSTEM_PROMPT_NON_I18N = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}
-${USE_THE_FOLLOWING_CONTEXT_TO_ANSWER}`;
+export const DEFAULT_SYSTEM_PROMPT_NON_I18N = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}`;
 
 export const DEFAULT_SYSTEM_PROMPT_NAME = i18n.translate(
   'xpack.elasticAssistant.assistant.content.prompts.system.defaultSystemPromptName',
@@ -48,8 +40,7 @@ export const DEFAULT_SYSTEM_PROMPT_NAME = i18n.translate(
 );
 
 export const SUPERHERO_SYSTEM_PROMPT_NON_I18N = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}
-${SUPERHERO_PERSONALITY}
-${USE_THE_FOLLOWING_CONTEXT_TO_ANSWER}`;
+${SUPERHERO_PERSONALITY}`;
 
 export const SUPERHERO_SYSTEM_PROMPT_NAME = i18n.translate(
   'xpack.elasticAssistant.assistant.content.prompts.system.superheroSystemPromptName',

x-pack/packages/kbn-langchain/README.md

@@ -0,0 +1,5 @@
+# @kbn/langchain
+
+Contains LangChain language models to be used with Kibana connectors
+
+The package does not expose `index.ts` at its root, instead there's a `server` directory you should deep-import from.

x-pack/packages/kbn-langchain/jest.config.js

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+module.exports = {
+  coverageDirectory: '<rootDir>/target/kibana-coverage/jest/x-pack/packages/kbn_langchain',
+  coverageReporters: ['text', 'html'],
+  collectCoverageFrom: [
+    '<rootDir>/x-pack/packages/kbn-langchain/server/**/*.{ts}',
+    '!<rootDir>/x-pack/packages/kbn-langchain/server/{__test__,__snapshots__,__examples__,*mock*,tests,test_helpers,integration_tests,types}/**/*',
+    '!<rootDir>/x-pack/packages/kbn-langchain/server/*mock*.{ts}',
+    '!<rootDir>/x-pack/packages/kbn-langchain/server/*.test.{ts}',
+    '!<rootDir>/x-pack/packages/kbn-langchain/server/*.d.ts',
+    '!<rootDir>/x-pack/packages/kbn-langchain/server/*.config.ts',
+  ],
+  preset: '@kbn/test',
+  rootDir: '../../..',
+  roots: ['<rootDir>/x-pack/packages/kbn-langchain'],
+};

x-pack/packages/kbn-langchain/kibana.jsonc

@@ -0,0 +1,5 @@
+{
+  "type": "shared-server",
+  "id": "@kbn/langchain",
+  "owner": "@elastic/security-generative-ai"
+}

x-pack/packages/kbn-langchain/package.json

@@ -0,0 +1,7 @@
+{
+  "name": "@kbn/langchain",
+  "private": true,
+  "version": "1.0.0",
+  "license": "Elastic License 2.0",
+  "sideEffects": false
+}

x-pack/packages/kbn-langchain/server/index.ts

@@ -0,0 +1,20 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { ActionsClientChatOpenAI } from './language_models/chat_openai';
+import { ActionsClientLlm } from './language_models/llm';
+import { ActionsClientSimpleChatModel } from './language_models/simple_chat_model';
+import { parseBedrockStream } from './utils/bedrock';
+import { getDefaultArguments } from './language_models/constants';
+
+export {
+  parseBedrockStream,
+  getDefaultArguments,
+  ActionsClientChatOpenAI,
+  ActionsClientLlm,
+  ActionsClientSimpleChatModel,
+};