[Security solution] AWS Bedrock connector

package.json

@@ -838,6 +838,7 @@
     "antlr4ts": "^0.5.0-alpha.3",
     "archiver": "^5.3.1",
     "async": "^3.2.3",
+    "aws4": "^1.12.0",
     "axios": "^1.4.0",
     "base64-js": "^1.3.1",
     "bitmap-sdf": "^1.0.3",

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

@@ -43,15 +43,14 @@ export const fetchConnectorExecuteAction = async ({
           temperature: 0.2,
         }
       : {
+          // Azure OpenAI and Bedrock invokeAI both expect this body format
           messages: outboundMessages,
         };
 
   const requestBody = {
     params: {
-      subActionParams: {
-        body: JSON.stringify(body),
-      },
-      subAction: 'test',
+      subActionParams: body,
+      subAction: 'invokeAI',
     },
   };
 
@@ -60,28 +59,23 @@ export const fetchConnectorExecuteAction = async ({
       ? `/internal/elastic_assistant/actions/connector/${apiConfig?.connectorId}/_execute`
       : `/api/actions/connector/${apiConfig?.connectorId}/_execute`;
 
-    // TODO: Find return type for this API
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    const response = await http.fetch<any>(path, {
-      method: 'POST',
-      headers: {
-        'Content-Type': 'application/json',
-      },
-      body: JSON.stringify(requestBody),
-      signal,
-    });
+    const response = await http.fetch<{ connector_id: string; status: string; data: string }>(
+      path,
+      {
+        method: 'POST',
+        headers: {
+          'Content-Type': 'application/json',
+        },
+        body: JSON.stringify(requestBody),
+        signal,
+      }
+    );
 
-    const data = response.data;
     if (response.status !== 'ok') {
       return API_ERROR;
     }
 
-    if (data.choices && data.choices.length > 0 && data.choices[0].message.content) {
-      const result = data.choices[0].message.content.trim();
-      return result;
-    } else {
-      return API_ERROR;
-    }
+    return response.data;
   } catch (error) {
     return API_ERROR;
   }

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

@@ -110,7 +110,6 @@ export const AssistantTitle: React.FC<{
           <EuiFlexItem grow={false}>
             <ConnectorSelectorInline
               isDisabled={isDisabled || selectedConversation === undefined}
-              onConnectorModalVisibilityChange={() => {}}
               selectedConnectorId={selectedConnectorId}
               selectedConversation={selectedConversation}
             />

x-pack/packages/kbn-elastic-assistant/impl/assistant/conversations/conversation_settings/conversation_settings.tsx

@@ -8,7 +8,6 @@
 import { EuiFormRow, EuiLink, EuiTitle, EuiText, EuiHorizontalRule, EuiSpacer } from '@elastic/eui';
 import React, { useCallback, useMemo } from 'react';
 
-import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
 import { HttpSetup } from '@kbn/core-http-browser';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/public/common';
@@ -27,7 +26,6 @@ import { useLoadConnectors } from '../../../connectorland/use_load_connectors';
 import { getGenAiConfig } from '../../../connectorland/helpers';
 
 export interface ConversationSettingsProps {
-  actionTypeRegistry: ActionTypeRegistryContract;
   allSystemPrompts: Prompt[];
   conversationSettings: UseAssistantContext['conversations'];
   defaultConnectorId?: string;
@@ -46,7 +44,6 @@ export interface ConversationSettingsProps {
  */
 export const ConversationSettings: React.FC<ConversationSettingsProps> = React.memo(
   ({
-    actionTypeRegistry,
     allSystemPrompts,
     defaultConnectorId,
     defaultProvider,
@@ -250,10 +247,7 @@ export const ConversationSettings: React.FC<ConversationSettingsProps> = React.m
           }
         >
           <ConnectorSelector
-            actionTypeRegistry={actionTypeRegistry}
-            http={http}
             isDisabled={selectedConversation == null}
-            onConnectorModalVisibilityChange={() => {}}
             onConnectorSelectionChange={handleOnConnectorSelectionChange}
             selectedConnectorId={selectedConnector?.id}
           />

x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings.tsx

@@ -76,13 +76,8 @@ export const AssistantSettings: React.FC<Props> = React.memo(
     selectedConversation: defaultSelectedConversation,
     setSelectedConversationId,
   }) => {
-    const {
-      actionTypeRegistry,
-      assistantLangChain,
-      http,
-      selectedSettingsTab,
-      setSelectedSettingsTab,
-    } = useAssistantContext();
+    const { assistantLangChain, http, selectedSettingsTab, setSelectedSettingsTab } =
+      useAssistantContext();
     const {
       conversationSettings,
       defaultAllow,
@@ -267,7 +262,6 @@ export const AssistantSettings: React.FC<Props> = React.memo(
                     conversationSettings={conversationSettings}
                     setUpdatedConversationSettings={setUpdatedConversationSettings}
                     allSystemPrompts={systemPromptSettings}
-                    actionTypeRegistry={actionTypeRegistry}
                     selectedConversation={selectedConversation}
                     onSelectedConversationChange={onHandleSelectedConversationChange}
                     http={http}

x-pack/packages/kbn-elastic-assistant/impl/assistant_context/types.tsx

@@ -45,6 +45,7 @@ export interface ConversationTheme {
 export interface Conversation {
   apiConfig: {
     connectorId?: string;
+    connectorTypeTitle?: string;
     defaultSystemPromptId?: string;
     provider?: OpenAiProviderType;
     model?: string;

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

@@ -8,61 +8,66 @@
 import { EuiButtonEmpty, EuiFlexGroup, EuiFlexItem, EuiSuperSelect, EuiText } from '@elastic/eui';
 import React, { useCallback, useMemo, useState } from 'react';
 
-import {
-  ActionConnector,
-  ActionTypeRegistryContract,
-} from '@kbn/triggers-actions-ui-plugin/public';
+import { ActionConnector, ActionType } from '@kbn/triggers-actions-ui-plugin/public';
 
-import { HttpSetup } from '@kbn/core-http-browser';
 import { ConnectorAddModal } from '@kbn/triggers-actions-ui-plugin/public/common/constants';
-import { GEN_AI_CONNECTOR_ID } from '@kbn/stack-connectors-plugin/public/common';
+import { ActionTypeSelectorModal } from '../connector_selector_inline/action_type_selector_modal';
 import { useLoadConnectors } from '../use_load_connectors';
 import * as i18n from '../translations';
 import { useLoadActionTypes } from '../use_load_action_types';
 import { useAssistantContext } from '../../assistant_context';
-import { getGenAiConfig } from '../helpers';
+import { getActionTypeTitle, getGenAiConfig } from '../helpers';
 
 export const ADD_NEW_CONNECTOR = 'ADD_NEW_CONNECTOR';
+
 interface Props {
-  actionTypeRegistry: ActionTypeRegistryContract;
-  http: HttpSetup;
   isDisabled?: boolean;
-  onConnectorSelectionChange: (connector: ActionConnector | undefined) => void;
+  isOpen?: boolean;
+  onConnectorSelectionChange: (connector: AIConnector) => void;
   selectedConnectorId?: string;
-  onConnectorModalVisibilityChange?: (isVisible: boolean) => void;
+  displayFancy?: (displayText: string) => React.ReactNode;
+  setIsOpen?: (isOpen: boolean) => void;
 }
 
+export type AIConnector = ActionConnector & {
+  connectorTypeTitle: string;
+};
+
 export const ConnectorSelector: React.FC<Props> = React.memo(
   ({
-    actionTypeRegistry,
-    http,
     isDisabled = false,
-    onConnectorModalVisibilityChange,
+    isOpen = false,
+    displayFancy,
     selectedConnectorId,
     onConnectorSelectionChange,
+    setIsOpen,
   }) => {
-    const { assistantAvailability } = useAssistantContext();
+    const { actionTypeRegistry, http, assistantAvailability } = useAssistantContext();
     // Connector Modal State
     const [isConnectorModalVisible, setIsConnectorModalVisible] = useState<boolean>(false);
     const { data: actionTypes } = useLoadActionTypes({ http });
-    const actionType = actionTypes?.find((at) => at.id === GEN_AI_CONNECTOR_ID) ?? {
-      enabledInConfig: true,
-      enabledInLicense: true,
-      minimumLicenseRequired: 'platinum',
-      supportedFeatureIds: ['general'],
-      isSystemActionType: false,
-      id: '.gen-ai',
-      name: 'Generative AI',
-      enabled: true,
-    };
+
+    const [selectedActionType, setSelectedActionType] = useState<ActionType | null>(null);
 
     const {
-      data: connectors,
-      isLoading: isLoadingActionTypes,
-      isFetching: isFetchingActionTypes,
+      data: connectorsWithoutActionContext,
+      isLoading: isLoadingConnectors,
+      isFetching: isFetchingConnectors,
       refetch: refetchConnectors,
     } = useLoadConnectors({ http });
-    const isLoading = isLoadingActionTypes || isFetchingActionTypes;
+
+    const aiConnectors: AIConnector[] = useMemo(
+      () =>
+        connectorsWithoutActionContext
+          ? connectorsWithoutActionContext.map((c) => ({
+              ...c,
+              connectorTypeTitle: getActionTypeTitle(actionTypeRegistry.get(c.actionTypeId)),
+            }))
+          : [],
+      [actionTypeRegistry, connectorsWithoutActionContext]
+    );
+
+    const isLoading = isLoadingConnectors || isFetchingConnectors;
     const localIsDisabled = isDisabled || !assistantAvailability.hasConnectorsReadPrivilege;
 
     const addNewConnectorOption = useMemo(() => {
@@ -72,7 +77,12 @@ export const ConnectorSelector: React.FC<Props> = React.memo(
         dropdownDisplay: (
           <EuiFlexGroup gutterSize="none" key={ADD_NEW_CONNECTOR}>
             <EuiFlexItem grow={true}>
-              <EuiButtonEmpty href="#" iconType="plus" size="xs">
+              <EuiButtonEmpty
+                data-test-subj="addNewConnectorButton"
+                href="#"
+                iconType="plus"
+                size="xs"
+              >
                 {i18n.ADD_NEW_CONNECTOR}
               </EuiButtonEmpty>
             </EuiFlexItem>
@@ -85,16 +95,17 @@ export const ConnectorSelector: React.FC<Props> = React.memo(
       };
     }, []);
 
-    const connectorOptions = useMemo(() => {
-      return (
-        connectors?.map((connector) => {
-          const apiProvider = getGenAiConfig(connector)?.apiProvider;
+    const connectorOptions = useMemo(
+      () =>
+        aiConnectors.map((connector) => {
+          const connectorTypeTitle =
+            getGenAiConfig(connector)?.apiProvider ?? connector.connectorTypeTitle;
           const connectorDetails = connector.isPreconfigured
             ? i18n.PRECONFIGURED_CONNECTOR
-            : apiProvider;
+            : connectorTypeTitle;
           return {
             value: connector.id,
-            inputDisplay: connector.name,
+            inputDisplay: displayFancy ? displayFancy(connector.name) : connector.name,
             dropdownDisplay: (
               <React.Fragment key={connector.id}>
                 <strong>{connector.name}</strong>
@@ -106,9 +117,9 @@ export const ConnectorSelector: React.FC<Props> = React.memo(
               </React.Fragment>
             ),
           };
-        }) ?? []
-      );
-    }, [connectors]);
+        }),
+      [aiConnectors, displayFancy]
+    );
 
     // Only include add new connector option if user has privilege
     const allConnectorOptions = useMemo(
@@ -120,22 +131,39 @@ export const ConnectorSelector: React.FC<Props> = React.memo(
     );
 
     const cleanupAndCloseModal = useCallback(() => {
-      onConnectorModalVisibilityChange?.(false);
+      setIsOpen?.(false);
       setIsConnectorModalVisible(false);
-    }, [onConnectorModalVisibilityChange]);
+      setSelectedActionType(null);
+    }, []);
+
+    const [modalForceOpen, setModalForceOpen] = useState(isOpen);
 
     const onChange = useCallback(
       (connectorId: string) => {
         if (connectorId === ADD_NEW_CONNECTOR) {
-          onConnectorModalVisibilityChange?.(true);
+          setModalForceOpen(false);
           setIsConnectorModalVisible(true);
           return;
         }
 
-        const connector = connectors?.find((c) => c.id === connectorId);
-        onConnectorSelectionChange(connector);
+        const connector = aiConnectors.find((c) => c.id === connectorId);
+        if (connector) {
+          onConnectorSelectionChange(connector);
+        }
       },
-      [connectors, onConnectorSelectionChange, onConnectorModalVisibilityChange]
+      [aiConnectors, onConnectorSelectionChange]
+    );
+
+    const onSaveConnector = useCallback(
+      (connector: ActionConnector) => {
+        onConnectorSelectionChange({
+          ...connector,
+          connectorTypeTitle: getActionTypeTitle(actionTypeRegistry.get(connector.actionTypeId)),
+        });
+        refetchConnectors?.();
+        cleanupAndCloseModal();
+      },
+      [actionTypeRegistry, cleanupAndCloseModal, onConnectorSelectionChange, refetchConnectors]
     );
 
     return (
@@ -146,19 +174,24 @@ export const ConnectorSelector: React.FC<Props> = React.memo(
           disabled={localIsDisabled}
           hasDividers={true}
           isLoading={isLoading}
+          isOpen={modalForceOpen}
           onChange={onChange}
           options={allConnectorOptions}
           valueOfSelected={selectedConnectorId ?? ''}
         />
-        {isConnectorModalVisible && (
+        {isConnectorModalVisible && !selectedActionType && (
+          <ActionTypeSelectorModal
+            actionTypes={actionTypes}
+            actionTypeRegistry={actionTypeRegistry}
+            onClose={() => setIsConnectorModalVisible(false)}
+            onSelect={(actionType: ActionType) => setSelectedActionType(actionType)}
+          />
+        )}
+        {isConnectorModalVisible && selectedActionType && (
           <ConnectorAddModal
-            actionType={actionType}
+            actionType={selectedActionType}
             onClose={cleanupAndCloseModal}
-            postSaveEventHandler={(connector: ActionConnector) => {
-              onConnectorSelectionChange(connector);
-              refetchConnectors?.();
-              cleanupAndCloseModal();
-            }}
+            postSaveEventHandler={onSaveConnector}
             actionTypeRegistry={actionTypeRegistry}
           />
         )}