From 65d1a9de887aa483702ff42d41238c56cfa84f91 Mon Sep 17 00:00:00 2001 From: Garrett Spong Date: Wed, 9 Oct 2024 10:17:47 -0600 Subject: [PATCH] [Security Assistant] V2 Knowledge Base Settings feedback and fixes (#194354) ## Summary This PR is a follow up to #192665 and addresses a bunch of feedback and fixes including: - [X] Adds support for updating/editing entries - [X] Fixes initial loading experience of the KB Settings Setup/Table - [X] Fixes two bugs where `semantic_text` and `text` must be declared for `IndexEntries` to work - [X] Add new Settings Context Menu items for KB and Alerts - [X] Add support for `required` entries in initial prompt * See [this trace](https://smith.langchain.com/public/84a17a31-8ce8-4bd9-911e-38a854484dd8/r) for included knowledge. Note that the KnowledgeBaseRetrievalTool was not selected. * Note: All prompts were updated to include the `{knowledge_history}` placeholder, and _not behind the feature flag_, as this will just be the empty case until the feature flag is enabled. TODO (in this or follow-up PR): - [ ] Add suggestions to `index` and `fields` inputs - [ ] Adds URL deeplinking to securityAssistantManagement - [ ] Fix bug where updating entry does not re-create embeddings (see [comment](https://github.com/elastic/kibana/pull/194354#discussion_r1786475496)) - [ ] Fix loading indicators when adding/editing entries - [ ] API integration tests for update API (@e40pud) ### Checklist Delete any items that are not applicable to this PR. - [X] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md) - [ ] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials * Docs being tracked in https://github.com/elastic/security-docs/issues/5337 for when feature flag is enabled - [ ] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com> Co-authored-by: Patryk Kopycinski (cherry picked from commit 7df36721923159f45bc4fdbd26f76b20ad84249a) --- .../entries/common_attributes.gen.ts | 6 +- .../entries/common_attributes.schema.yaml | 6 + .../impl/assistant/assistant_header/index.tsx | 79 +------- .../assistant_header/translations.ts | 28 +++ .../alerts_settings}/alerts_settings.test.tsx | 4 +- .../alerts_settings}/alerts_settings.tsx | 6 +- .../alerts_settings_management.tsx | 15 +- .../assistant_settings_management.test.tsx | 6 + .../assistant_settings_management.tsx | 8 +- .../settings_context_menu.tsx | 186 ++++++++++++++++++ .../impl/knowledge_base/alerts_range.tsx | 2 +- .../knowledge_base_settings.tsx | 2 +- .../document_entry_editor.tsx | 1 - .../index.tsx | 106 ++++++++-- .../index_entry_editor.tsx | 47 ++++- .../translations.ts | 41 +++- .../use_knowledge_base_table.tsx | 8 +- .../kbn-elastic-assistant/tsconfig.json | 1 + .../create_knowledge_base_entry.ts | 108 +++++++++- .../knowledge_base/helpers.ts | 5 +- .../knowledge_base/index.ts | 53 ++++- .../knowledge_base/types.ts | 33 ++++ .../server/ai_assistant_service/index.ts | 8 +- .../graphs/default_assistant_graph/graph.ts | 2 +- .../nodes/run_agent.ts | 14 ++ .../nodes/translations.ts | 6 +- .../graphs/default_assistant_graph/prompts.ts | 2 + .../entries/bulk_actions_route.ts | 23 ++- .../knowledge_base/entries/create_route.ts | 2 +- .../knowledge_base/entries/find_route.ts | 4 +- .../management_settings.test.tsx | 5 + .../stack_management/management_settings.tsx | 12 +- 32 files changed, 686 insertions(+), 143 deletions(-) rename x-pack/packages/kbn-elastic-assistant/impl/{alerts/settings => assistant/settings/alerts_settings}/alerts_settings.test.tsx (89%) rename x-pack/packages/kbn-elastic-assistant/impl/{alerts/settings => assistant/settings/alerts_settings}/alerts_settings.tsx (92%) rename x-pack/packages/kbn-elastic-assistant/impl/{alerts/settings => assistant/settings/alerts_settings}/alerts_settings_management.tsx (68%) create mode 100644 x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts index 1af5c46b1c130..c32517fec0860 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts @@ -106,7 +106,11 @@ export type BaseCreateProps = z.infer; export const BaseCreateProps = BaseRequiredFields.merge(BaseDefaultableFields); export type BaseUpdateProps = z.infer; -export const BaseUpdateProps = BaseCreateProps.partial(); +export const BaseUpdateProps = BaseCreateProps.partial().merge( + z.object({ + id: NonEmptyString, + }) +); export type BaseResponseProps = z.infer; export const BaseResponseProps = BaseRequiredFields.merge(BaseDefaultableFields.required()); diff --git a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml index c1c551059f04b..af7f4dd8e4221 100644 --- a/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml +++ b/x-pack/packages/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.schema.yaml @@ -112,6 +112,12 @@ components: allOf: - $ref: "#/components/schemas/BaseCreateProps" x-modify: partial + - type: object + properties: + id: + $ref: "../../common_attributes.schema.yaml#/components/schemas/NonEmptyString" + required: + - id BaseResponseProps: x-inline: true diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx index d81a56fb97eef..ef37506f2af17 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/index.tsx @@ -5,16 +5,13 @@ * 2.0. */ -import React, { useState, useMemo, useCallback } from 'react'; +import React, { useMemo, useCallback } from 'react'; import { QueryObserverResult, RefetchOptions, RefetchQueryFilters } from '@tanstack/react-query'; import { EuiFlexGroup, EuiFlexItem, - EuiPopover, - EuiContextMenu, EuiButtonIcon, EuiPanel, - EuiConfirmModal, EuiToolTip, EuiSkeletonTitle, } from '@elastic/eui'; @@ -29,6 +26,7 @@ import { FlyoutNavigation } from '../assistant_overlay/flyout_navigation'; import { AssistantSettingsButton } from '../settings/assistant_settings_button'; import * as i18n from './translations'; import { AIConnector } from '../../connectorland/connector_selector'; +import { SettingsContextMenu } from '../settings/settings_context_menu/settings_context_menu'; interface OwnProps { selectedConversation: Conversation | undefined; @@ -94,21 +92,6 @@ export const AssistantHeader: React.FC = ({ [selectedConversation?.apiConfig?.connectorId] ); - const [isPopoverOpen, setPopover] = useState(false); - - const onButtonClick = useCallback(() => { - setPopover(!isPopoverOpen); - }, [isPopoverOpen]); - - const closePopover = useCallback(() => { - setPopover(false); - }, []); - - const [isResetConversationModalVisible, setIsResetConversationModalVisible] = useState(false); - - const closeDestroyModal = useCallback(() => setIsResetConversationModalVisible(false), []); - const showDestroyModal = useCallback(() => setIsResetConversationModalVisible(true), []); - const onConversationChange = useCallback( (updatedConversation: Conversation) => { onConversationSelected({ @@ -119,32 +102,6 @@ export const AssistantHeader: React.FC = ({ [onConversationSelected] ); - const panels = useMemo( - () => [ - { - id: 0, - items: [ - { - name: i18n.RESET_CONVERSATION, - css: css` - color: ${euiThemeVars.euiColorDanger}; - `, - onClick: showDestroyModal, - icon: 'refresh', - 'data-test-subj': 'clear-chat', - }, - ], - }, - ], - [showDestroyModal] - ); - - const handleReset = useCallback(() => { - onChatCleared(); - closeDestroyModal(); - closePopover(); - }, [onChatCleared, closeDestroyModal, closePopover]); - return ( <> = ({ - - } - isOpen={isPopoverOpen} - closePopover={closePopover} - panelPaddingSize="none" - anchorPosition="downLeft" - > - - + - {isResetConversationModalVisible && ( - -

{i18n.CLEAR_CHAT_CONFIRMATION}

- - )} ); }; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts index 68c926d2aa14c..e4f23e0970eb0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/assistant_header/translations.ts @@ -7,6 +7,34 @@ import { i18n } from '@kbn/i18n'; +export const AI_ASSISTANT_SETTINGS = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.aiAssistantSettings', + { + defaultMessage: 'AI Assistant settings', + } +); + +export const ANONYMIZATION = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.anonymization', + { + defaultMessage: 'Anonymization', + } +); + +export const KNOWLEDGE_BASE = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBase', + { + defaultMessage: 'Knowledge Base', + } +); + +export const ALERTS_TO_ANALYZE = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.alertsToAnalyze', + { + defaultMessage: 'Alerts to analyze', + } +); + export const RESET_CONVERSATION = i18n.translate( 'xpack.elasticAssistant.assistant.settings.resetConversation', { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.test.tsx similarity index 89% rename from x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.test.tsx rename to x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.test.tsx index 3e730451ba1d5..2a5cae76d5e77 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.test.tsx @@ -9,8 +9,8 @@ import { render, screen, fireEvent } from '@testing-library/react'; import React from 'react'; import { AlertsSettings } from './alerts_settings'; -import { KnowledgeBaseConfig } from '../../assistant/types'; -import { DEFAULT_LATEST_ALERTS } from '../../assistant_context/constants'; +import { KnowledgeBaseConfig } from '../../types'; +import { DEFAULT_LATEST_ALERTS } from '../../../assistant_context/constants'; describe('AlertsSettings', () => { beforeEach(() => { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.tsx similarity index 92% rename from x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.tsx rename to x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.tsx index e73bfa15e66be..60078178a1771 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings.tsx @@ -9,9 +9,9 @@ import { EuiFlexGroup, EuiFormRow, EuiFlexItem, EuiSpacer, EuiText } from '@elas import { css } from '@emotion/react'; import React from 'react'; -import { KnowledgeBaseConfig } from '../../assistant/types'; -import { AlertsRange } from '../../knowledge_base/alerts_range'; -import * as i18n from '../../knowledge_base/translations'; +import { KnowledgeBaseConfig } from '../../types'; +import { AlertsRange } from '../../../knowledge_base/alerts_range'; +import * as i18n from '../../../knowledge_base/translations'; export const MIN_LATEST_ALERTS = 10; export const MAX_LATEST_ALERTS = 100; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings_management.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings_management.tsx similarity index 68% rename from x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings_management.tsx rename to x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings_management.tsx index d103c1a8c03c2..1a6f826bd415f 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/alerts/settings/alerts_settings_management.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/alerts_settings/alerts_settings_management.tsx @@ -7,19 +7,24 @@ import { EuiPanel, EuiSpacer, EuiText, EuiTitle } from '@elastic/eui'; import React from 'react'; -import { KnowledgeBaseConfig } from '../../assistant/types'; -import { AlertsRange } from '../../knowledge_base/alerts_range'; -import * as i18n from '../../knowledge_base/translations'; +import { KnowledgeBaseConfig } from '../../types'; +import { AlertsRange } from '../../../knowledge_base/alerts_range'; +import * as i18n from '../../../knowledge_base/translations'; interface Props { knowledgeBase: KnowledgeBaseConfig; setUpdatedKnowledgeBaseSettings: React.Dispatch>; + hasBorder?: boolean; } +/** + * Replaces the AlertsSettings component used in the existing settings modal. Once the modal is + * fully removed we can delete that component in favor of this one. + */ export const AlertsSettingsManagement: React.FC = React.memo( - ({ knowledgeBase, setUpdatedKnowledgeBaseSettings }) => { + ({ knowledgeBase, setUpdatedKnowledgeBaseSettings, hasBorder = true }) => { return ( - +

{i18n.ALERTS_LABEL}

 diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.test.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.test.tsx index d8e207cbb23cd..dd472b3ee87ab 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.test.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.test.tsx @@ -25,6 +25,7 @@ import { SYSTEM_PROMPTS_TAB, } from './const'; import { mockSystemPrompts } from '../../mock/system_prompt'; +import { DataViewsContract } from '@kbn/data-views-plugin/public'; const mockConversations = { [alertConvo.title]: alertConvo, @@ -53,8 +54,13 @@ const mockContext = { }, }; +const mockDataViews = { + getIndices: jest.fn(), +} as unknown as DataViewsContract; + const testProps = { selectedConversation: welcomeConvo, + dataViews: mockDataViews, }; jest.mock('../../assistant_context'); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx index 89c00fbf88773..4c50d14a5662e 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/assistant_settings_management.tsx @@ -9,6 +9,7 @@ import React, { useEffect, useMemo } from 'react'; import { EuiAvatar, EuiPageTemplate, EuiTitle, useEuiShadow, useEuiTheme } from '@elastic/eui'; import { css } from '@emotion/react'; +import { DataViewsContract } from '@kbn/data-views-plugin/public'; import { Conversation } from '../../..'; import * as i18n from './translations'; import { useAssistantContext } from '../../assistant_context'; @@ -33,6 +34,7 @@ import { KnowledgeBaseSettingsManagement } from '../../knowledge_base/knowledge_ import { EvaluationSettings } from '.'; interface Props { + dataViews: DataViewsContract; selectedConversation: Conversation; } @@ -41,7 +43,7 @@ interface Props { * anonymization, knowledge base, and evaluation via the `isModelEvaluationEnabled` feature flag. */ export const AssistantSettingsManagement: React.FC = React.memo( - ({ selectedConversation: defaultSelectedConversation }) => { + ({ dataViews, selectedConversation: defaultSelectedConversation }) => { const { assistantFeatures: { assistantModelEvaluation: modelEvaluatorEnabled }, http, @@ -158,7 +160,9 @@ export const AssistantSettingsManagement: React.FC = React.memo( )} {selectedSettingsTab === QUICK_PROMPTS_TAB && } {selectedSettingsTab === ANONYMIZATION_TAB && } - {selectedSettingsTab === KNOWLEDGE_BASE_TAB && } + {selectedSettingsTab === KNOWLEDGE_BASE_TAB && ( + + )} {selectedSettingsTab === EVALUATION_TAB && } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx new file mode 100644 index 0000000000000..b7f33b9a6af5a --- /dev/null +++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/settings/settings_context_menu/settings_context_menu.tsx @@ -0,0 +1,186 @@ +/* + * 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 React, { ReactElement, useCallback, useMemo, useState } from 'react'; +import { + EuiFlexGroup, + EuiFlexItem, + EuiContextMenuPanel, + EuiContextMenuItem, + EuiConfirmModal, + EuiNotificationBadge, + EuiPopover, + EuiButtonIcon, +} from '@elastic/eui'; +import { css } from '@emotion/react'; +import { euiThemeVars } from '@kbn/ui-theme'; +import { useAssistantContext } from '../../../..'; +import * as i18n from '../../assistant_header/translations'; + +interface Params { + isDisabled?: boolean; + onChatCleared?: () => void; +} + +export const SettingsContextMenu: React.FC = React.memo( + ({ isDisabled = false, onChatCleared }: Params) => { + const { + navigateToApp, + knowledgeBase, + assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, + } = useAssistantContext(); + + const [isPopoverOpen, setPopover] = useState(false); + + const [isResetConversationModalVisible, setIsResetConversationModalVisible] = useState(false); + const closeDestroyModal = useCallback(() => setIsResetConversationModalVisible(false), []); + + const onButtonClick = useCallback(() => { + setPopover(!isPopoverOpen); + }, [isPopoverOpen]); + + const closePopover = useCallback(() => { + setPopover(false); + }, []); + + const showDestroyModal = useCallback(() => { + closePopover?.(); + setIsResetConversationModalVisible(true); + }, [closePopover]); + + const handleNavigateToSettings = useCallback( + () => + navigateToApp('management', { + path: 'kibana/securityAiAssistantManagement', + }), + [navigateToApp] + ); + + const handleNavigateToKnowledgeBase = useCallback( + () => + navigateToApp('management', { + path: 'kibana/securityAiAssistantManagement', + }), + [navigateToApp] + ); + + // We are migrating away from the settings modal in favor of the new Stack Management UI + // Currently behind `assistantKnowledgeBaseByDefault` FF + const newItems: ReactElement[] = useMemo( + () => [ + + {i18n.AI_ASSISTANT_SETTINGS} + , + + {i18n.ANONYMIZATION} + , + + {i18n.KNOWLEDGE_BASE} + , + + + {i18n.ALERTS_TO_ANALYZE} + + + {knowledgeBase.latestAlerts} + + + + , + ], + [handleNavigateToKnowledgeBase, handleNavigateToSettings, knowledgeBase] + ); + + const items = useMemo( + () => [ + ...(enableKnowledgeBaseByDefault ? newItems : []), + + {i18n.RESET_CONVERSATION} + , + ], + + [enableKnowledgeBaseByDefault, newItems, showDestroyModal] + ); + + const handleReset = useCallback(() => { + onChatCleared?.(); + closeDestroyModal(); + closePopover?.(); + }, [onChatCleared, closeDestroyModal, closePopover]); + + return ( + <> + + } + isOpen={isPopoverOpen} + closePopover={closePopover} + panelPaddingSize="none" + anchorPosition="leftUp" + > + + + {isResetConversationModalVisible && ( + +

{i18n.CLEAR_CHAT_CONFIRMATION}

+ + )} + + ); + } +); + +SettingsContextMenu.displayName = 'SettingsContextMenu'; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.tsx index 152f0a91a7d04..63bd86121dcc1 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/alerts_range.tsx @@ -12,7 +12,7 @@ import { MAX_LATEST_ALERTS, MIN_LATEST_ALERTS, TICK_INTERVAL, -} from '../alerts/settings/alerts_settings'; +} from '../assistant/settings/alerts_settings/alerts_settings'; import { KnowledgeBaseConfig } from '../assistant/types'; import { ALERTS_RANGE } from './translations'; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx index df254805d9cee..28ce86a3595a0 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings.tsx @@ -23,7 +23,7 @@ import { import { FormattedMessage } from '@kbn/i18n-react'; import { css } from '@emotion/react'; -import { AlertsSettings } from '../alerts/settings/alerts_settings'; +import { AlertsSettings } from '../assistant/settings/alerts_settings/alerts_settings'; import { useAssistantContext } from '../assistant_context'; import type { KnowledgeBaseConfig } from '../assistant/types'; import * as i18n from './translations'; diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx index 016da27d2c051..b33f221bfde3b 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/document_entry_editor.tsx @@ -127,7 +127,6 @@ export const DocumentEntryEditor: React.FC = React.memo(({ entry, setEntr id="requiredKnowledge" onChange={onRequiredKnowledgeChanged} checked={entry?.required ?? false} - disabled={true} /> diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx index a2097177a2ca4..34e8601e37ce7 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index.tsx @@ -6,8 +6,12 @@ */ import { + EuiButton, + EuiFlexGroup, + EuiFlexItem, EuiInMemoryTable, EuiLink, + EuiLoadingSpinner, EuiPanel, EuiSearchBarProps, EuiSpacer, @@ -23,7 +27,9 @@ import { KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, } from '@kbn/elastic-assistant-common'; -import { AlertsSettingsManagement } from '../../alerts/settings/alerts_settings_management'; +import { css } from '@emotion/react'; +import { DataViewsContract } from '@kbn/data-views-plugin/public'; +import { AlertsSettingsManagement } from '../../assistant/settings/alerts_settings/alerts_settings_management'; import { useKnowledgeBaseEntries } from '../../assistant/api/knowledge_base/entries/use_knowledge_base_entries'; import { useAssistantContext } from '../../assistant_context'; import { useKnowledgeBaseTable } from './use_knowledge_base_table'; @@ -40,7 +46,7 @@ import { useFlyoutModalVisibility } from '../../assistant/common/components/assi import { IndexEntryEditor } from './index_entry_editor'; import { DocumentEntryEditor } from './document_entry_editor'; import { KnowledgeBaseSettings } from '../knowledge_base_settings'; -import { SetupKnowledgeBaseButton } from '../setup_knowledge_base_button'; +import { ESQL_RESOURCE, SetupKnowledgeBaseButton } from '../setup_knowledge_base_button'; import { useDeleteKnowledgeBaseEntries } from '../../assistant/api/knowledge_base/entries/use_delete_knowledge_base_entries'; import { isSystemEntry, @@ -51,14 +57,24 @@ import { useCreateKnowledgeBaseEntry } from '../../assistant/api/knowledge_base/ import { useUpdateKnowledgeBaseEntries } from '../../assistant/api/knowledge_base/entries/use_update_knowledge_base_entries'; import { SETTINGS_UPDATED_TOAST_TITLE } from '../../assistant/settings/translations'; import { KnowledgeBaseConfig } from '../../assistant/types'; +import { + isKnowledgeBaseSetup, + useKnowledgeBaseStatus, +} from '../../assistant/api/knowledge_base/use_knowledge_base_status'; + +interface Params { + dataViews: DataViewsContract; +} -export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { +export const KnowledgeBaseSettingsManagement: React.FC = React.memo(({ dataViews }) => { const { assistantFeatures: { assistantKnowledgeBaseByDefault: enableKnowledgeBaseByDefault }, http, toasts, } = useAssistantContext(); const [hasPendingChanges, setHasPendingChanges] = useState(false); + const { data: kbStatus, isFetched } = useKnowledgeBaseStatus({ http, resource: ESQL_RESOURCE }); + const isKbSetup = isKnowledgeBaseSetup(kbStatus); // Only needed for legacy settings management const { knowledgeBase, setUpdatedKnowledgeBaseSettings, resetSettings, saveSettings } = @@ -123,12 +139,12 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { // Flyout Save/Cancel Actions const onSaveConfirmed = useCallback(() => { - if (isKnowledgeBaseEntryCreateProps(selectedEntry)) { - createEntry(selectedEntry); - closeFlyout(); - } else if (isKnowledgeBaseEntryResponse(selectedEntry)) { + if (isKnowledgeBaseEntryResponse(selectedEntry)) { updateEntries([selectedEntry]); closeFlyout(); + } else if (isKnowledgeBaseEntryCreateProps(selectedEntry)) { + createEntry(selectedEntry); + closeFlyout(); } }, [closeFlyout, selectedEntry, createEntry, updateEntries]); @@ -137,7 +153,11 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { closeFlyout(); }, [closeFlyout]); - const { data: entries } = useKnowledgeBaseEntries({ + const { + data: entries, + isFetching: isFetchingEntries, + refetch: refetchEntries, + } = useKnowledgeBaseEntries({ http, toasts, enabled: enableKnowledgeBaseByDefault, @@ -169,6 +189,9 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { [deleteEntry, entries.data, getColumns, openFlyout] ); + // Refresh button + const handleRefreshTable = useCallback(() => refetchEntries(), [refetchEntries]); + const onDocumentClicked = useCallback(() => { setSelectedEntry({ type: DocumentEntryType.value, kbResource: 'user', source: 'user' }); openFlyout(); @@ -182,7 +205,30 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { const search: EuiSearchBarProps = useMemo( () => ({ toolsRight: ( - + + + + + + + + + + ), box: { incremental: true, @@ -190,7 +236,7 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { }, filters: [], }), - [onDocumentClicked, onIndexClicked] + [isFetchingEntries, handleRefreshTable, onDocumentClicked, onIndexClicked] ); const flyoutTitle = useMemo(() => { @@ -247,15 +293,40 @@ export const KnowledgeBaseSettingsManagement: React.FC = React.memo(() => { ), }} /> - - + + + {!isFetched ? ( + + ) : isKbSetup ? ( + + ) : ( + <> + + + + + + + + + + + + + + )} + + { ) : ( >> } diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx index 19f8cfbbc52ba..f5dd2df3bcaac 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/index_entry_editor.tsx @@ -17,14 +17,16 @@ import { } from '@elastic/eui'; import React, { useCallback } from 'react'; import { IndexEntry } from '@kbn/elastic-assistant-common'; +import { DataViewsContract } from '@kbn/data-views-plugin/public'; import * as i18n from './translations'; interface Props { + dataViews: DataViewsContract; entry?: IndexEntry; setEntry: React.Dispatch>>; } -export const IndexEntryEditor: React.FC = React.memo(({ entry, setEntry }) => { +export const IndexEntryEditor: React.FC = React.memo(({ dataViews, entry, setEntry }) => { // Name const setName = useCallback( (e: React.ChangeEvent) => @@ -74,9 +76,17 @@ export const IndexEntryEditor: React.FC = React.memo(({ entry, setEntry } entry?.users?.length === 0 ? sharingOptions[1].value : sharingOptions[0].value; // Index + // TODO: For index field autocomplete + // const indexOptions = useMemo(() => { + // const indices = await dataViews.getIndices({ + // pattern: e[0]?.value ?? '', + // isRollupIndex: () => false, + // }); + // }, [dataViews]); const setIndex = useCallback( - (e: Array>) => - setEntry((prevEntry) => ({ ...prevEntry, index: e[0].value })), + async (e: Array>) => { + setEntry((prevEntry) => ({ ...prevEntry, index: e[0]?.value })); + }, [setEntry] ); @@ -162,30 +172,51 @@ export const IndexEntryEditor: React.FC = React.memo(({ entry, setEntry } - + - + + + + ); }); diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts index ed4a3676975b8..0cc16089fdaae 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/translations.ts @@ -251,14 +251,44 @@ export const ENTRY_FIELD_INPUT_LABEL = i18n.translate( export const ENTRY_DESCRIPTION_INPUT_LABEL = i18n.translate( 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryDescriptionInputLabel', { - defaultMessage: 'Description', + defaultMessage: 'Data Description', + } +); + +export const ENTRY_DESCRIPTION_HELP_LABEL = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryDescriptionHelpLabel', + { + defaultMessage: + 'A description of the type of data in this index and/or when the assistant should look for data here.', } ); export const ENTRY_QUERY_DESCRIPTION_INPUT_LABEL = i18n.translate( 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryQueryDescriptionInputLabel', { - defaultMessage: 'Query Description', + defaultMessage: 'Query Instruction', + } +); + +export const ENTRY_QUERY_DESCRIPTION_HELP_LABEL = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryQueryDescriptionHelpLabel', + { + defaultMessage: 'Any instructions for extracting the search query from the user request.', + } +); + +export const ENTRY_OUTPUT_FIELDS_INPUT_LABEL = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryOutputFieldsInputLabel', + { + defaultMessage: 'Output Fields', + } +); + +export const ENTRY_OUTPUT_FIELDS_HELP_LABEL = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryOutputFieldsHelpLabel', + { + defaultMessage: + 'What fields should be sent to the LLM. Leave empty to send the entire document.', } ); @@ -269,6 +299,13 @@ export const ENTRY_INPUT_PLACEHOLDER = i18n.translate( } ); +export const ENTRY_FIELD_PLACEHOLDER = i18n.translate( + 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.entryFieldPlaceholder', + { + defaultMessage: 'semantic_text', + } +); + export const KNOWLEDGE_BASE_DOCUMENTATION = i18n.translate( 'xpack.elasticAssistant.assistant.settings.knowledgeBaseSettingsManagement.knowledgeBaseDocumentation', { diff --git a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx index 5af360a598205..d0038169cd597 100644 --- a/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx +++ b/x-pack/packages/kbn-elastic-assistant/impl/knowledge_base/knowledge_base_settings_management/use_knowledge_base_table.tsx @@ -5,7 +5,7 @@ * 2.0. */ -import { EuiAvatar, EuiBadge, EuiBasicTableColumn, EuiIcon, EuiLink, EuiText } from '@elastic/eui'; +import { EuiAvatar, EuiBadge, EuiBasicTableColumn, EuiIcon, EuiText } from '@elastic/eui'; import { css } from '@emotion/react'; import React, { useCallback } from 'react'; import { FormattedDate } from '@kbn/i18n-react'; @@ -32,7 +32,7 @@ export const useKnowledgeBaseTable = () => { if (['esql', 'security_labs'].includes(entry.kbResource)) { return 'logoElastic'; } - return 'visText'; + return 'document'; } else if (entry.type === IndexEntryType.value) { return 'index'; } @@ -61,9 +61,7 @@ export const useKnowledgeBaseTable = () => { }, { name: i18n.COLUMN_NAME, - render: (entry: KnowledgeBaseEntryResponse) => ( - onEntryNameClicked(entry)}>{entry.name} - ), + render: ({ name }: KnowledgeBaseEntryResponse) => name, sortable: ({ name }: KnowledgeBaseEntryResponse) => name, width: '30%', }, diff --git a/x-pack/packages/kbn-elastic-assistant/tsconfig.json b/x-pack/packages/kbn-elastic-assistant/tsconfig.json index ed2631b597bd6..8d19fa86f4d11 100644 --- a/x-pack/packages/kbn-elastic-assistant/tsconfig.json +++ b/x-pack/packages/kbn-elastic-assistant/tsconfig.json @@ -30,5 +30,6 @@ "@kbn/core-doc-links-browser", "@kbn/core", "@kbn/zod", + "@kbn/data-views-plugin", ] } diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts index 7dac58ddecc9b..aef66d406bf74 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry.ts @@ -12,10 +12,11 @@ import { DocumentEntryCreateFields, KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, + KnowledgeBaseEntryUpdateProps, Metadata, } from '@kbn/elastic-assistant-common'; import { getKnowledgeBaseEntry } from './get_knowledge_base_entry'; -import { CreateKnowledgeBaseEntrySchema } from './types'; +import { CreateKnowledgeBaseEntrySchema, UpdateKnowledgeBaseEntrySchema } from './types'; export interface CreateKnowledgeBaseEntryParams { esClient: ElasticsearchClient; @@ -77,6 +78,111 @@ export const createKnowledgeBaseEntry = async ({ } }; +interface TransformToUpdateSchemaProps { + user: AuthenticatedUser; + updatedAt: string; + entry: KnowledgeBaseEntryUpdateProps; + global?: boolean; +} + +export const transformToUpdateSchema = ({ + user, + updatedAt, + entry, + global = false, +}: TransformToUpdateSchemaProps): UpdateKnowledgeBaseEntrySchema => { + const base = { + id: entry.id, + updated_at: updatedAt, + updated_by: user.profile_uid ?? 'unknown', + name: entry.name, + type: entry.type, + users: global + ? [] + : [ + { + id: user.profile_uid, + name: user.username, + }, + ], + }; + + if (entry.type === 'index') { + const { inputSchema, outputFields, queryDescription, ...restEntry } = entry; + return { + ...base, + ...restEntry, + query_description: queryDescription, + input_schema: + entry.inputSchema?.map((schema) => ({ + field_name: schema.fieldName, + field_type: schema.fieldType, + description: schema.description, + })) ?? undefined, + output_fields: outputFields ?? undefined, + }; + } + return { + ...base, + kb_resource: entry.kbResource, + required: entry.required ?? false, + source: entry.source, + text: entry.text, + vector: undefined, + }; +}; + +export const getUpdateScript = ({ + entry, + isPatch, +}: { + entry: UpdateKnowledgeBaseEntrySchema; + isPatch?: boolean; +}) => { + return { + source: ` + if (params.assignEmpty == true || params.containsKey('name')) { + ctx._source.name = params.name; + } + if (params.assignEmpty == true || params.containsKey('type')) { + ctx._source.type = params.type; + } + if (params.assignEmpty == true || params.containsKey('users')) { + ctx._source.users = params.users; + } + if (params.assignEmpty == true || params.containsKey('query_description')) { + ctx._source.query_description = params.query_description; + } + if (params.assignEmpty == true || params.containsKey('input_schema')) { + ctx._source.input_schema = params.input_schema; + } + if (params.assignEmpty == true || params.containsKey('output_fields')) { + ctx._source.output_fields = params.output_fields; + } + if (params.assignEmpty == true || params.containsKey('kb_resource')) { + ctx._source.kb_resource = params.kb_resource; + } + if (params.assignEmpty == true || params.containsKey('required')) { + ctx._source.required = params.required; + } + if (params.assignEmpty == true || params.containsKey('source')) { + ctx._source.source = params.source; + } + if (params.assignEmpty == true || params.containsKey('text')) { + ctx._source.text = params.text; + } + ctx._source.updated_at = params.updated_at; + ctx._source.updated_by = params.updated_by; + `, + lang: 'painless', + params: { + ...entry, // when assigning undefined in painless, it will remove property and wil set it to null + // for patch we don't want to remove unspecified value in payload + assignEmpty: !(isPatch ?? true), + }, + }; +}; + interface TransformToCreateSchemaProps { createdAt: string; spaceId: string; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index 8ff8de6cfb408..de76a38135f0b 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -6,6 +6,7 @@ */ import { z } from '@kbn/zod'; +import { get } from 'lodash'; import { DynamicStructuredTool } from '@langchain/core/tools'; import { errors } from '@elastic/elasticsearch'; import { QueryDslQueryContainer, SearchRequest } from '@elastic/elasticsearch/lib/api/types'; @@ -189,7 +190,7 @@ export const getStructuredToolForIndexEntry = ({ standard: { query: { nested: { - path: 'semantic_text.inference.chunks', + path: `${indexEntry.field}.inference.chunks`, query: { sparse_vector: { inference_id: elserId, @@ -220,7 +221,7 @@ export const getStructuredToolForIndexEntry = ({ }, {}); } return { - text: (hit._source as { text: string }).text, + text: get(hit._source, `${indexEntry.field}.inference.chunks[0].text`), }; }); diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 7f665fa7f9a16..8d8bf069d779e 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -15,6 +15,7 @@ import { Document } from 'langchain/document'; import type { SavedObjectsClientContract } from '@kbn/core-saved-objects-api-server'; import { DocumentEntryType, + DocumentEntry, IndexEntry, KnowledgeBaseEntryCreateProps, KnowledgeBaseEntryResponse, @@ -444,7 +445,9 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { ); this.options.logger.debug( () => - `getKnowledgeBaseDocuments() - Similarity Search Results:\n ${JSON.stringify(results)}` + `getKnowledgeBaseDocuments() - Similarity Search returned [${JSON.stringify( + results.length + )}] results` ); return results; @@ -454,6 +457,47 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { } }; + /** + * Returns all global and current user's private `required` document entries. + */ + public getRequiredKnowledgeBaseDocumentEntries = async (): Promise => { + const user = this.options.currentUser; + if (user == null) { + throw new Error( + 'Authenticated user not found! Ensure kbDataClient was initialized from a request.' + ); + } + + try { + const userFilter = getKBUserFilter(user); + const results = await this.findDocuments({ + // Note: This is a magic number to set some upward bound as to not blow the context with too + // many historical KB entries. Ideally we'd query for all and token trim. + perPage: 100, + page: 1, + sortField: 'created_at', + sortOrder: 'asc', + filter: `${userFilter} AND type:document AND kb_resource:user AND required:true`, + }); + this.options.logger.debug( + `kbDataClient.getRequiredKnowledgeBaseDocumentEntries() - results:\n${JSON.stringify( + results + )}` + ); + + if (results) { + return transformESSearchToKnowledgeBaseEntry(results.data) as DocumentEntry[]; + } + } catch (e) { + this.options.logger.error( + `kbDataClient.getRequiredKnowledgeBaseDocumentEntries() - Failed to fetch DocumentEntries` + ); + return []; + } + + return []; + }; + /** * Creates a new Knowledge Base Entry. * @@ -492,7 +536,10 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { }; /** - * Returns AssistantTools for any 'relevant' KB IndexEntries that exist in the knowledge base + * Returns AssistantTools for any 'relevant' KB IndexEntries that exist in the knowledge base. + * + * Note: Accepts esClient so retrieval can be scoped to the current user as esClient on kbDataClient + * is scoped to system user. */ public getAssistantTools = async ({ assistantToolParams, @@ -520,7 +567,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { page: 1, sortField: 'created_at', sortOrder: 'asc', - filter: `${userFilter}${` AND type:index`}`, // TODO: Support global tools (no user filter), and filter by space as well + filter: `${userFilter} AND type:index`, }); this.options.logger.debug( `kbDataClient.getAssistantTools() - results:\n${JSON.stringify(results, null, 2)}` diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts index ecf9260e999d2..3de1a15d79b2a 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/types.ts @@ -82,6 +82,39 @@ export interface LegacyEsKnowledgeBaseEntrySchema { model_id: string; }; } +export interface UpdateKnowledgeBaseEntrySchema { + id: string; + created_at?: string; + created_by?: string; + updated_at?: string; + updated_by?: string; + users?: Array<{ + id?: string; + name?: string; + }>; + name?: string; + type?: string; + // Document Entry Fields + kb_resource?: string; + required?: boolean; + source?: string; + text?: string; + vector?: { + tokens: Record; + model_id: string; + }; + // Index Entry Fields + index?: string; + field?: string; + description?: string; + query_description?: string; + input_schema?: Array<{ + field_name: string; + field_type: string; + description: string; + }>; + output_fields?: string[]; +} export interface CreateKnowledgeBaseEntrySchema { '@timestamp'?: string; diff --git a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 942f94c203873..08912f41a8bbc 100644 --- a/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -84,6 +84,7 @@ export class AIAssistantService { private isKBSetupInProgress: boolean = false; // Temporary 'feature flag' to determine if we should initialize the new kb mappings, toggled when accessing kbDataClient private v2KnowledgeBaseEnabled: boolean = false; + private hasInitializedV2KnowledgeBase: boolean = false; constructor(private readonly options: AIAssistantServiceOpts) { this.initialized = false; @@ -363,8 +364,13 @@ export class AIAssistantService { // If either v2 KB or a modelIdOverride is provided, we need to reinitialize all persistence resources to make sure // they're using the correct model/mappings. Technically all existing KB data is stale since it was created // with a different model/mappings, but modelIdOverride is only intended for testing purposes at this time - if (opts.v2KnowledgeBaseEnabled || opts.modelIdOverride != null) { + // Added hasInitializedV2KnowledgeBase to prevent the console noise from re-init on each KB request + if ( + !this.hasInitializedV2KnowledgeBase && + (opts.v2KnowledgeBaseEnabled || opts.modelIdOverride != null) + ) { await this.initializeResources(); + this.hasInitializedV2KnowledgeBase = true; } const res = await this.checkResourcesInstallation(opts); diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts index 8f2f713c170ed..f8e21da57e5b9 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/graph.ts @@ -141,7 +141,7 @@ export const getDefaultAssistantGraph = ({ }) ) .addNode(NodeType.AGENT, (state: AgentState) => - runAgent({ ...nodeParams, state, agentRunnable }) + runAgent({ ...nodeParams, state, agentRunnable, kbDataClient: dataClients?.kbDataClient }) ) .addNode(NodeType.TOOLS, (state: AgentState) => executeTools({ ...nodeParams, state, tools })) .addNode(NodeType.RESPOND, (state: AgentState) => diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts index 2d076f6bd1472..053254a1d99b3 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/run_agent.ts @@ -10,15 +10,20 @@ import { AgentRunnableSequence } from 'langchain/dist/agents/agent'; import { formatLatestUserMessage } from '../prompts'; import { AgentState, NodeParamsBase } from '../types'; import { NodeType } from '../constants'; +import { AIAssistantKnowledgeBaseDataClient } from '../../../../../ai_assistant_data_clients/knowledge_base'; export interface RunAgentParams extends NodeParamsBase { state: AgentState; config?: RunnableConfig; agentRunnable: AgentRunnableSequence; + kbDataClient?: AIAssistantKnowledgeBaseDataClient; } export const AGENT_NODE_TAG = 'agent_run'; +const KNOWLEDGE_HISTORY_PREFIX = 'Knowledge History:'; +const NO_KNOWLEDGE_HISTORY = '[No existing knowledge history]'; + /** * Node to run the agent * @@ -26,18 +31,27 @@ export const AGENT_NODE_TAG = 'agent_run'; * @param state - The current state of the graph * @param config - Any configuration that may've been supplied * @param agentRunnable - The agent to run + * @param kbDataClient - Data client for accessing the Knowledge Base on behalf of the current user */ export async function runAgent({ logger, state, agentRunnable, config, + kbDataClient, }: RunAgentParams): Promise> { logger.debug(() => `${NodeType.AGENT}: Node state:\n${JSON.stringify(state, null, 2)}`); + const knowledgeHistory = await kbDataClient?.getRequiredKnowledgeBaseDocumentEntries(); + const agentOutcome = await agentRunnable.withConfig({ tags: [AGENT_NODE_TAG] }).invoke( { ...state, + knowledge_history: `${KNOWLEDGE_HISTORY_PREFIX}\n${ + knowledgeHistory?.length + ? JSON.stringify(knowledgeHistory.map((e) => e.text)) + : NO_KNOWLEDGE_HISTORY + }`, // prepend any user prompt (gemini) input: formatLatestUserMessage(state.input, state.llmType), chat_history: state.messages, // TODO: Message de-dupe with ...state spread diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts index e55e1081e6474..e5a1c14846e23 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/translations.ts @@ -8,8 +8,10 @@ const YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT = 'You are a security analyst and expert in resolving security incidents. Your role is to assist by answering questions about Elastic Security.'; const IF_YOU_DONT_KNOW_THE_ANSWER = 'Do not answer questions unrelated to Elastic Security.'; +export const KNOWLEDGE_HISTORY = + 'If available, use the Knowledge History provided to try and answer the question. If not provided, you can try and query for additional knowledge via the KnowledgeBaseRetrievalTool.'; -export const DEFAULT_SYSTEM_PROMPT = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}`; +export const DEFAULT_SYSTEM_PROMPT = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER} ${KNOWLEDGE_HISTORY}`; // system prompt from @afirstenberg const BASE_GEMINI_PROMPT = 'You are an assistant that is an expert at using tools and Elastic Security, doing your best to use these tools to answer questions or follow instructions. It is very important to use tools to answer the question or follow the instructions rather than coming up with your own answer. Tool calls are good. Sometimes you may need to make several tool calls to accomplish the task or get an answer to the question that was asked. Use as many tool calls as necessary.'; @@ -19,7 +21,7 @@ export const GEMINI_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} ${KB_CATCH}`; export const BEDROCK_SYSTEM_PROMPT = `Use tools as often as possible, as they have access to the latest data and syntax. Always return value from ESQLKnowledgeBaseTool as is. Never return tags in the response, but make sure to include tags content in the response. Do not reflect on the quality of the returned search results in your response.`; export const GEMINI_USER_PROMPT = `Now, always using the tools at your disposal, step by step, come up with a response to this request:\n\n`; -export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. You have access to the following tools: +export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. ${KNOWLEDGE_HISTORY} You have access to the following tools: {tools} diff --git a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts index 883047ed7b9df..05cc8b50852f5 100644 --- a/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts +++ b/x-pack/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/prompts.ts @@ -17,6 +17,7 @@ import { export const formatPrompt = (prompt: string, additionalPrompt?: string) => ChatPromptTemplate.fromMessages([ ['system', additionalPrompt ? `${prompt}\n\n${additionalPrompt}` : prompt], + ['placeholder', '{knowledge_history}'], ['placeholder', '{chat_history}'], ['human', '{input}'], ['placeholder', '{agent_scratchpad}'], @@ -39,6 +40,7 @@ export const geminiToolCallingAgentPrompt = formatPrompt(systemPrompts.gemini); export const formatPromptStructured = (prompt: string, additionalPrompt?: string) => ChatPromptTemplate.fromMessages([ ['system', additionalPrompt ? `${prompt}\n\n${additionalPrompt}` : prompt], + ['placeholder', '{knowledge_history}'], ['placeholder', '{chat_history}'], [ 'human', diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts index 96045b17e6171..ce3f0c8c92693 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts @@ -22,11 +22,18 @@ import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/ import { performChecks } from '../../helpers'; import { KNOWLEDGE_BASE_ENTRIES_TABLE_MAX_PAGE_SIZE } from '../../../../common/constants'; -import { EsKnowledgeBaseEntrySchema } from '../../../ai_assistant_data_clients/knowledge_base/types'; +import { + EsKnowledgeBaseEntrySchema, + UpdateKnowledgeBaseEntrySchema, +} from '../../../ai_assistant_data_clients/knowledge_base/types'; import { ElasticAssistantPluginRouter } from '../../../types'; import { buildResponse } from '../../utils'; import { transformESSearchToKnowledgeBaseEntry } from '../../../ai_assistant_data_clients/knowledge_base/transforms'; -import { transformToCreateSchema } from '../../../ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry'; +import { + getUpdateScript, + transformToCreateSchema, + transformToUpdateSchema, +} from '../../../ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry'; export interface BulkOperationError { message: string; @@ -210,7 +217,17 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug }) ), documentsToDelete: body.delete?.ids, - documentsToUpdate: [], // TODO: Support bulk update + documentsToUpdate: body.update?.map((entry) => + // TODO: KB-RBAC check, required when users != null as entry will either be created globally if empty + transformToUpdateSchema({ + user: authenticatedUser, + updatedAt: changedAt, + entry, + global: entry.users != null && entry.users.length === 0, + }) + ), + getUpdateScript: (entry: UpdateKnowledgeBaseEntrySchema) => + getUpdateScript({ entry, isPatch: true }), authenticatedUser, }); const created = diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts index 3dbb5a9cf930e..51e3d48505ec2 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/create_route.ts @@ -66,7 +66,7 @@ export const createKnowledgeBaseEntryRoute = (router: ElasticAssistantPluginRout logger.debug(() => `Creating KB Entry:\n${JSON.stringify(request.body)}`); const createResponse = await kbDataClient?.createKnowledgeBaseEntry({ knowledgeBaseEntry: request.body, - // TODO: KB-RBAC check, required when users != null as entry will either be created globally if empty, or for specific users (only admin API feature) + // TODO: KB-RBAC check, required when users != null as entry will either be created globally if empty global: request.body.users != null && request.body.users.length === 0, }); diff --git a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts index 78b7ac720dd1a..7bafff17352dc 100644 --- a/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts +++ b/x-pack/plugins/elastic_assistant/server/routes/knowledge_base/entries/find_route.ts @@ -74,7 +74,7 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout }); const currentUser = ctx.elasticAssistant.getCurrentUser(); const userFilter = getKBUserFilter(currentUser); - const systemFilter = ` AND kb_resource:"user"`; + const systemFilter = ` AND (kb_resource:"user" OR type:"index")`; const additionalFilter = query.filter ? ` AND ${query.filter}` : ''; const result = await kbDataClient?.findDocuments({ @@ -166,7 +166,7 @@ export const findKnowledgeBaseEntriesRoute = (router: ElasticAssistantPluginRout body: { perPage: result.perPage, page: result.page, - total: result.total, + total: result.total + systemEntries.length, data: [...transformESSearchToKnowledgeBaseEntry(result.data), ...systemEntries], }, }); diff --git a/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.test.tsx b/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.test.tsx index 1c988d14e845f..65a0ab84d3412 100644 --- a/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.test.tsx +++ b/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.test.tsx @@ -77,6 +77,11 @@ describe('ManagementSettings', () => { securitySolutionAssistant: { 'ai-assistant': false }, }, }, + data: { + dataViews: { + getIndices: jest.fn(), + }, + }, security: { userProfiles: { getCurrent: jest.fn().mockResolvedValue({ data: { color: 'blue', initials: 'P' } }), diff --git a/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.tsx b/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.tsx index 90e39398474ec..48d89e02dfc71 100644 --- a/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.tsx +++ b/x-pack/plugins/security_solution/public/assistant/stack_management/management_settings.tsx @@ -37,6 +37,7 @@ export const ManagementSettings = React.memo(() => { securitySolutionAssistant: { 'ai-assistant': securityAIAssistantEnabled }, }, }, + data: { dataViews }, security, } = useKibana().services; @@ -46,8 +47,8 @@ export const ManagementSettings = React.memo(() => { security?.userProfiles.getCurrent<{ avatar: UserAvatar }>({ dataPath: 'avatar', }), - select: (data) => { - return data.data.avatar; + select: (d) => { + return d.data.avatar; }, keepPreviousData: true, refetchOnWindowFocus: false, @@ -79,7 +80,12 @@ export const ManagementSettings = React.memo(() => { } if (conversations) { - return ; + return ( + + ); } return <>;