From 02b9f8f24997ee2be82ee26eb4278cfd594c7a6f Mon Sep 17 00:00:00 2001 From: Joe Reuter Date: Tue, 4 Mar 2025 17:46:14 +0100 Subject: [PATCH 1/9] =?UTF-8?q?=F0=9F=8C=8A=20Streams:=20Disable=20AI=20su?= =?UTF-8?q?ggestions=20button=20if=20there=20is=20no=20sample=20data=20(#2?= =?UTF-8?q?13113)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This PR makes the AI suggestions button more stable in case of misconfigured fields: * Only make the button clickable if there are sample values * Filter out sample documents that don't have the required field on the server (would have broken the request before) Screenshot 2025-03-04 at 15 43 23 --- .../streams/processing/suggestions_handler.ts | 22 ++++--- .../processors/grok/grok_ai_suggestions.tsx | 59 +++++++++++++------ 2 files changed, 56 insertions(+), 25 deletions(-) diff --git a/x-pack/platform/plugins/shared/streams/server/routes/streams/processing/suggestions_handler.ts b/x-pack/platform/plugins/shared/streams/server/routes/streams/processing/suggestions_handler.ts index a953a86ec9b35..121f401cc934d 100644 --- a/x-pack/platform/plugins/shared/streams/server/routes/streams/processing/suggestions_handler.ts +++ b/x-pack/platform/plugins/shared/streams/server/routes/streams/processing/suggestions_handler.ts @@ -66,14 +66,20 @@ export function extractAndGroupPatterns(samples: FlattenRecord[], field: string) const NUMBER_PATTERN_CATEGORIES = 5; const NUMBER_SAMPLES_PER_PATTERN = 8; - const samplesWithPatterns = samples.map((sample) => { - const pattern = evalPattern(get(sample, field) as string); - return { - document: sample, - fullPattern: pattern, - truncatedPattern: pattern.slice(0, 10), - fieldValue: get(sample, field) as string, - }; + const samplesWithPatterns = samples.flatMap((sample) => { + const value = get(sample, field); + if (typeof value !== 'string') { + return []; + } + const pattern = evalPattern(value); + return [ + { + document: sample, + fullPattern: pattern, + truncatedPattern: pattern.slice(0, 10), + fieldValue: get(sample, field) as string, + }, + ]; }); // Group samples by their truncated patterns diff --git a/x-pack/platform/plugins/shared/streams_app/public/components/data_management/stream_detail_enrichment/processors/grok/grok_ai_suggestions.tsx b/x-pack/platform/plugins/shared/streams_app/public/components/data_management/stream_detail_enrichment/processors/grok/grok_ai_suggestions.tsx index 4db576375530e..6a61317c01c6d 100644 --- a/x-pack/platform/plugins/shared/streams_app/public/components/data_management/stream_detail_enrichment/processors/grok/grok_ai_suggestions.tsx +++ b/x-pack/platform/plugins/shared/streams_app/public/components/data_management/stream_detail_enrichment/processors/grok/grok_ai_suggestions.tsx @@ -5,7 +5,7 @@ * 2.0. */ -import React, { useCallback, useState } from 'react'; +import React, { useCallback, useMemo, useState } from 'react'; import { EuiBadge, EuiButton, @@ -40,12 +40,14 @@ const RefreshButton = ({ selectConnector, currentConnector, isLoading, + hasValidField, }: { generatePatterns: () => void; selectConnector?: UseGenAIConnectorsResult['selectConnector']; connectors?: FindActionResult[]; currentConnector?: string; isLoading: boolean; + hasValidField: boolean; }) => { const [isPopoverOpen, { off: closePopover, toggle: togglePopover }] = useBoolean(false); const splitButtonPopoverId = useGeneratedHtmlId({ @@ -55,21 +57,34 @@ const RefreshButton = ({ return ( - - {i18n.translate( - 'xpack.streams.streamDetailView.managementTab.enrichment.processorFlyout.refreshSuggestions', - { - defaultMessage: 'Generate patterns', - } - )} - + + {i18n.translate( + 'xpack.streams.streamDetailView.managementTab.enrichment.processorFlyout.refreshSuggestions', + { + defaultMessage: 'Generate patterns', + } + )} + + {connectors && connectors.length > 1 && ( @@ -205,7 +220,16 @@ function InnerGrokAiSuggestions({ content = {suggestionsError.message}; } - const currentPatterns = form.getValues().patterns; + const { field: currentFieldName, patterns: currentPatterns } = form.getValues(); + + const hasValidField = useMemo(() => { + return Boolean( + currentFieldName && + filteredSamples.some( + (sample) => sample[currentFieldName] && typeof sample[currentFieldName] === 'string' + ) + ); + }, [filteredSamples, currentFieldName]); const filteredSuggestions = suggestions?.patterns .map((pattern, i) => ({ @@ -323,6 +347,7 @@ function InnerGrokAiSuggestions({ connectors={genAiConnectors?.connectors} selectConnector={genAiConnectors?.selectConnector} currentConnector={currentConnector} + hasValidField={hasValidField} /> @@ -365,7 +390,7 @@ export function GrokAiSuggestions() { ); } - if (!isAiEnabled || !props.filteredSamples.length) { + if (!isAiEnabled) { return null; } return ; From 5903c7a5520a32c482ebf846d5458ba49769a622 Mon Sep 17 00:00:00 2001 From: Mason Herron <46727170+Supplementing@users.noreply.github.com> Date: Tue, 4 Mar 2025 10:18:35 -0700 Subject: [PATCH 2/9] [Fleet] Improve validation for dynamic Kafka topics (#212422) Closes #206194 ## Summary - Removed hardcoded wrapping of user-entered topics with `%{[]}` to fix issues arising from the user pre-wrapping, and also allow greater flexibility in naming - Added validation rules to check for unclosed brackets & brackets with missing `%` preceding - Added the auto-wrapping to the `value` field of items chosen from the dropdown to ensure they were always wrapped as intended ### Checklist Check the PR satisfies following conditions. Reviewers should verify this PR satisfies this list as well. - [ ] 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/src/platform/packages/shared/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 - [ ] [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 - [ ] If a plugin configuration key changed, check if it needs to be allowlisted in the cloud and added to the [docker list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker) - [ ] This was checked for breaking HTTP API changes, and any breaking changes have been approved by the breaking-change committee. The `release_note:breaking` label should be applied in these situations. - [ ] [Flaky Test Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was used on any tests changed - [ ] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) ### Identify risks n/a --------- Co-authored-by: Elastic Machine --- .../shared/kbn-doc-links/src/get_doc_links.ts | 1 + .../shared/kbn-doc-links/src/types.ts | 1 + .../translations/translations/fr-FR.json | 1 - .../translations/translations/ja-JP.json | 1 - .../translations/translations/zh-CN.json | 1 - .../output_form_kafka_topics.tsx | 21 +++++++++++-- .../output_form_validators.test.tsx | 25 +++++++++++++++ .../output_form_validators.tsx | 31 ++++++++++++++++++- .../edit_output_flyout/use_output_form.tsx | 2 +- 9 files changed, 76 insertions(+), 8 deletions(-) diff --git a/src/platform/packages/shared/kbn-doc-links/src/get_doc_links.ts b/src/platform/packages/shared/kbn-doc-links/src/get_doc_links.ts index 7c06c923af754..8ce744a20857f 100644 --- a/src/platform/packages/shared/kbn-doc-links/src/get_doc_links.ts +++ b/src/platform/packages/shared/kbn-doc-links/src/get_doc_links.ts @@ -108,6 +108,7 @@ export const getDocLinks = ({ kibanaBranch, buildFlavor }: GetDocLinkOptions): D configuration: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/configuring-howto-filebeat.html`, elasticsearchModule: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/filebeat-module-elasticsearch.html`, elasticsearchOutput: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/elasticsearch-output.html`, + kafkaOutput: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/kafka-output.html`, startup: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/filebeat-starting.html`, exportedFields: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/exported-fields.html`, suricataModule: `${ELASTIC_WEBSITE_URL}guide/en/beats/filebeat/${DOC_LINK_VERSION}/filebeat-module-suricata.html`, diff --git a/src/platform/packages/shared/kbn-doc-links/src/types.ts b/src/platform/packages/shared/kbn-doc-links/src/types.ts index f70a951a3a3a1..261f7403bb40d 100644 --- a/src/platform/packages/shared/kbn-doc-links/src/types.ts +++ b/src/platform/packages/shared/kbn-doc-links/src/types.ts @@ -70,6 +70,7 @@ export interface DocLinks { readonly configuration: string; readonly elasticsearchOutput: string; readonly elasticsearchModule: string; + readonly kafkaOutput: string; readonly startup: string; readonly exportedFields: string; readonly suricataModule: string; diff --git a/x-pack/platform/plugins/private/translations/translations/fr-FR.json b/x-pack/platform/plugins/private/translations/translations/fr-FR.json index 09d029aadf792..eb7be40181a10 100644 --- a/x-pack/platform/plugins/private/translations/translations/fr-FR.json +++ b/x-pack/platform/plugins/private/translations/translations/fr-FR.json @@ -19015,7 +19015,6 @@ "xpack.fleet.settings.editOutputFlyout.kafkaCompressionTitle": "Compression", "xpack.fleet.settings.editOutputFlyout.kafkaConnectionTypeLabel": "Connexion", "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicHelptext": "Sélectionnez un sujet dans la liste. Si un sujet n'est pas disponible, créez un sujet personnalisé.", - "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicLabel": "Sujet du champ", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderKeyInputLabel": "Clé", "xpack.fleet.settings.editOutputFlyout.kafkaHeadersTitle": "En-têtes", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderValueInputLabel": "Valeur", diff --git a/x-pack/platform/plugins/private/translations/translations/ja-JP.json b/x-pack/platform/plugins/private/translations/translations/ja-JP.json index 6651f416d039d..4b9f7c3ce3cbf 100644 --- a/x-pack/platform/plugins/private/translations/translations/ja-JP.json +++ b/x-pack/platform/plugins/private/translations/translations/ja-JP.json @@ -18876,7 +18876,6 @@ "xpack.fleet.settings.editOutputFlyout.kafkaCompressionTitle": "圧縮", "xpack.fleet.settings.editOutputFlyout.kafkaConnectionTypeLabel": "接続", "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicHelptext": "リストからトピックを選択してください。トピックがない場合は、カスタムトピックを作成してください。", - "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicLabel": "フィールドからのトピック", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderKeyInputLabel": "キー", "xpack.fleet.settings.editOutputFlyout.kafkaHeadersTitle": "ヘッダー", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderValueInputLabel": "値", diff --git a/x-pack/platform/plugins/private/translations/translations/zh-CN.json b/x-pack/platform/plugins/private/translations/translations/zh-CN.json index 4c944837ad693..5dd2806726719 100644 --- a/x-pack/platform/plugins/private/translations/translations/zh-CN.json +++ b/x-pack/platform/plugins/private/translations/translations/zh-CN.json @@ -18575,7 +18575,6 @@ "xpack.fleet.settings.editOutputFlyout.kafkaCompressionTitle": "压缩", "xpack.fleet.settings.editOutputFlyout.kafkaConnectionTypeLabel": "连接", "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicHelptext": "从列表中选择主题。如果主题不可用，请创建定制主题。", - "xpack.fleet.settings.editOutputFlyout.kafkaDynamicTopicLabel": "来自字段的主题", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderKeyInputLabel": "钥匙", "xpack.fleet.settings.editOutputFlyout.kafkaHeadersTitle": "标题", "xpack.fleet.settings.editOutputFlyout.kafkaHeaderValueInputLabel": "值", diff --git a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_kafka_topics.tsx b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_kafka_topics.tsx index df6befc104f7b..8b4f9b3265f77 100644 --- a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_kafka_topics.tsx +++ b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_kafka_topics.tsx @@ -15,10 +15,13 @@ import { EuiTitle, EuiRadioGroup, EuiComboBox, + EuiLink, } from '@elastic/eui'; import { FormattedMessage } from '@kbn/i18n-react'; import { i18n } from '@kbn/i18n'; +import { useStartServices } from '../../../../hooks'; + import { kafkaTopicsType, KAFKA_DYNAMIC_FIELDS, @@ -30,10 +33,12 @@ import type { OutputFormInputsType } from './use_output_form'; export const OutputFormKafkaTopics: React.FunctionComponent<{ inputs: OutputFormInputsType }> = ({ inputs, }) => { + const { docLinks } = useStartServices(); + const dynamicOptions: Array> = useMemo(() => { const options = KAFKA_DYNAMIC_FIELDS.map((option) => ({ label: option, - value: option, + value: `%{[${option}]}`, })); return options; }, []); @@ -73,7 +78,17 @@ export const OutputFormKafkaTopics: React.FunctionComponent<{ inputs: OutputForm label={ + + + ), + }} /> } {...inputs.kafkaDynamicTopicInput.formRowProps} @@ -83,7 +98,7 @@ export const OutputFormKafkaTopics: React.FunctionComponent<{ inputs: OutputForm fullWidth isClearable={true} options={dynamicOptions} - customOptionText="Use custom field (not recommended)" + customOptionText="Use custom field" singleSelection={{ asPlainText: true }} {...inputs.kafkaDynamicTopicInput.props} /> diff --git a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.test.tsx b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.test.tsx index 888b815172bc2..b0fce70d831ab 100644 --- a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.test.tsx +++ b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.test.tsx @@ -14,6 +14,7 @@ import { validateKafkaHosts, validateKibanaURL, validateKibanaAPIKey, + validateDynamicKafkaTopics, } from './output_form_validators'; describe('Output form validation', () => { @@ -336,4 +337,28 @@ describe('Output form validation', () => { ]); }); }); + + describe('validateDynamicKafkaTopics', () => { + const validTopics = [ + { label: 'field1', value: '%{[field]}' }, + { label: 'field2', value: 'field2' }, + { label: 'field3', value: '%{[field2]}-%{[field3]}' }, + ]; + const invalidBracketTopic = [{ label: '%{[field}', value: '%{[field}' }]; + const invalidPercentTopic = [{ label: '{[field]}', value: '{[field]}' }]; + it('should work with valid topics', () => { + const res = validateDynamicKafkaTopics(validTopics); + expect(res).toBeUndefined(); + }); + it("should return error with missing brackets in topic's name", () => { + const res = validateDynamicKafkaTopics(invalidBracketTopic); + expect(res).toEqual([ + 'The topic should have a matching number of opening and closing brackets', + ]); + }); + it("should return error with missing percent sign before opening brackets in topic's name", () => { + const res = validateDynamicKafkaTopics(invalidPercentTopic); + expect(res).toEqual(['Opening brackets should be preceded by a percent sign']); + }); + }); }); diff --git a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.tsx b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.tsx index 14d2430f5121a..2410ddfe2e256 100644 --- a/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.tsx +++ b/x-pack/platform/plugins/shared/fleet/public/applications/fleet/sections/settings/components/edit_output_flyout/output_form_validators.tsx @@ -19,6 +19,15 @@ const toSecretValidator = return validator(value ?? ''); }; +const getAllIndices = (str: string, substring: string): number[] => { + const indices = []; + let index = str.indexOf(substring); + while (index !== -1) { + indices.push(index); + index = str.indexOf(substring, index + 1); + } + return indices; +}; export function validateKafkaHosts(value: string[]) { const res: Array<{ message: string; index?: number }> = []; const urlIndexes: { [key: string]: number[] } = {}; @@ -362,12 +371,31 @@ export function validateKafkaStaticTopic(value: string) { export function validateDynamicKafkaTopics(value: Array>) { const res = []; value.forEach((val, idx) => { - if (!val) { + if (!val || !val.value) { res.push( i18n.translate('xpack.fleet.settings.outputForm.kafkaTopicFieldRequiredMessage', { defaultMessage: 'Topic is required', }) ); + } else { + const openingBrackets = getAllIndices(val.value, '{['); + const closingBrackets = getAllIndices(val.value, ']}'); + if (openingBrackets.length !== closingBrackets.length) { + res.push( + i18n.translate('xpack.fleet.settings.outputForm.kafkaTopicBracketsError', { + defaultMessage: + 'The topic should have a matching number of opening and closing brackets', + }) + ); + } + // check for preceding percent sign + if (!openingBrackets.every((item) => val?.value![item - 1] === '%')) { + res.push( + i18n.translate('xpack.fleet.settings.outputForm.kafkaTopicPercentError', { + defaultMessage: 'Opening brackets should be preceded by a percent sign', + }) + ); + } } }); @@ -378,6 +406,7 @@ export function validateDynamicKafkaTopics(value: Array void, output?: Output, defaultOupu } : kafkaTopicsInput.value === kafkaTopicsType.Dynamic && kafkaDynamicTopicInput.value ? { - topic: `%{[${kafkaDynamicTopicInput.value}]}`, + topic: kafkaDynamicTopicInput.value, } : {}), headers: kafkaHeadersInput.value, From 28dc0f6ffcf82279dc3fcaee477b4727b8bbdd0a Mon Sep 17 00:00:00 2001 From: Krzysztof Kowalczyk Date: Tue, 4 Mar 2025 18:20:43 +0100 Subject: [PATCH 3/9] [Lens] Fix partition chart color assignments (#207178) Fixes the color assignment for partition charts consistent with the legend ordering. Aligns legacy and color mapping color logic. --- .../public/utils/layers/get_color.test.ts | 661 ++++++++---------- .../public/utils/layers/get_color.ts | 79 +-- .../public/utils/layers/get_layers.ts | 14 +- 3 files changed, 341 insertions(+), 413 deletions(-) diff --git a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.test.ts b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.test.ts index dee61c2fcdc0d..c717b000bdce8 100644 --- a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.test.ts +++ b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.test.ts @@ -12,401 +12,352 @@ import { chartPluginMock } from '@kbn/charts-plugin/public/mocks'; import { byDataColorPaletteMap, SimplifiedArrayNode } from './get_color'; import type { SeriesLayer } from '@kbn/coloring'; import { dataPluginMock } from '@kbn/data-plugin/public/mocks'; -import { fieldFormatsMock } from '@kbn/field-formats-plugin/common/mocks'; import type { DataPublicPluginStart } from '@kbn/data-plugin/public'; import { getColor } from './get_color'; import { createMockVisData, createMockBucketColumns, createMockPieParams } from '../../mocks'; -import { generateFormatters } from '../formatters'; import { ChartTypes } from '../../../common/types'; import { getDistinctSeries } from '..'; +import { getColorCategories } from '@kbn/chart-expressions-common'; -describe('#byDataColorPaletteMap', () => { - let paletteDefinition: PaletteDefinition; - let palette: PaletteOutput; - const visData = createMockVisData(); - const defaultFormatter = jest.fn((...args) => fieldFormatsMock.deserialize(...args)); - const formatters = generateFormatters(visData, defaultFormatter); +describe('get color', () => { + describe('#byDataColorPaletteMap', () => { + let paletteDefinition: PaletteDefinition; + let palette: PaletteOutput; + let colorIndexMap: Map; + const visData = createMockVisData(); + const categories = getColorCategories(visData.rows, visData.columns[1]?.id); - beforeEach(() => { - paletteDefinition = chartPluginMock.createPaletteRegistry().get('default'); - palette = { type: 'palette' } as PaletteOutput; - }); + beforeEach(() => { + paletteDefinition = chartPluginMock.createPaletteRegistry().get('default'); + palette = { type: 'palette' } as PaletteOutput; + colorIndexMap = new Map(categories.map((c, i) => [String(c), i])); + }); - it('should create byDataColorPaletteMap', () => { - expect( - byDataColorPaletteMap( - visData.rows, - visData.columns[0], - paletteDefinition, - palette, - formatters, - fieldFormatsMock - ) - ).toMatchInlineSnapshot(` - Object { - "getColor": [Function], - } - `); - }); + it('should create byDataColorPaletteMap', () => { + const colorPaletteMap = byDataColorPaletteMap(paletteDefinition, palette, colorIndexMap); + expect(colorPaletteMap.getColor).toEqual(expect.any(Function)); + }); - it('should get color', () => { - const colorPaletteMap = byDataColorPaletteMap( - visData.rows, - visData.columns[0], - paletteDefinition, - palette, - formatters, - fieldFormatsMock - ); + it('should get color', () => { + const colorPaletteMap = byDataColorPaletteMap(paletteDefinition, palette, colorIndexMap); - expect(colorPaletteMap.getColor('Logstash Airways')).toBe('black'); - }); + expect(colorPaletteMap.getColor('Logstash Airways')).toBe('black'); + }); - it('should return undefined in case if values not in datatable', () => { - const colorPaletteMap = byDataColorPaletteMap( - visData.rows, - visData.columns[0], - paletteDefinition, - palette, - formatters, - fieldFormatsMock - ); + it('should cache duplicate values', () => { + const colorPaletteMap = byDataColorPaletteMap(paletteDefinition, palette, colorIndexMap); + colorPaletteMap.getColor('JetBeats'); + colorPaletteMap.getColor('JetBeats'); - expect(colorPaletteMap.getColor('wrong')).toBeUndefined(); - }); + expect(paletteDefinition.getCategoricalColor).toBeCalledTimes(1); + }); + + it('should order rankAtDepth based on colorIndexMap for each value', () => { + const customColorIndexMap = new Map(); + customColorIndexMap.set('Logstash Airways', 0); + customColorIndexMap.set('empty - 1', 1); + customColorIndexMap.set('JetBeats', 2); + customColorIndexMap.set('empty - 2', 3); + customColorIndexMap.set('__other__', 5); - it('should increase rankAtDepth for each new value', () => { - const colorPaletteMap = byDataColorPaletteMap( - visData.rows, - visData.columns[0], - paletteDefinition, - palette, - formatters, - fieldFormatsMock - ); - colorPaletteMap.getColor('Logstash Airways'); - colorPaletteMap.getColor('JetBeats'); + const colorPaletteMap = byDataColorPaletteMap( + paletteDefinition, + palette, + customColorIndexMap + ); + + // should produce same color regardless of call order + ['__other__', 'JetBeats', 'Logstash Airways'].forEach((key) => { + colorPaletteMap.getColor(key); + }); + + expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith( + 1, + [{ name: '__other__', rankAtDepth: 5, totalSeriesAtDepth: 5 }], + { behindText: false }, + undefined + ); - expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith( - 1, - [{ name: 'Logstash Airways', rankAtDepth: 0, totalSeriesAtDepth: 4 }], - { behindText: false }, - undefined - ); + expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith( + 2, + [{ name: 'JetBeats', rankAtDepth: 2, totalSeriesAtDepth: 5 }], + { behindText: false }, + undefined + ); - expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith( - 2, - [{ name: 'JetBeats', rankAtDepth: 1, totalSeriesAtDepth: 4 }], - { behindText: false }, - undefined - ); + expect(paletteDefinition.getCategoricalColor).toHaveBeenNthCalledWith( + 3, + [{ name: 'Logstash Airways', rankAtDepth: 0, totalSeriesAtDepth: 5 }], + { behindText: false }, + undefined + ); + }); }); -}); -describe('getColor', () => { - const visData = createMockVisData(); - const buckets = createMockBucketColumns(); - const visParams = createMockPieParams(); - const colors = ['color1', 'color2', 'color3', 'color4']; - const dataMock = dataPluginMock.createStartContract(); - interface RangeProps { - gte: number; - lt: number; - } - const defaultFormatter = jest.fn((...args) => fieldFormatsMock.deserialize(...args)); - const formatters = generateFormatters(visData, defaultFormatter); - const distinctSeries = getDistinctSeries(visData.rows, buckets); - const dataLength = { columnsLength: buckets.length, rowsLength: visData.rows.length }; + describe('#getColor', () => { + const visData = createMockVisData(); + const buckets = createMockBucketColumns(); + const visParams = createMockPieParams(); + const colors = ['color1', 'color2', 'color3', 'color4']; + const categories = (chartType?: ChartTypes) => + chartType === ChartTypes.MOSAIC && visData.columns.length === 2 + ? getColorCategories(visData.rows, visData.columns[1]?.id) + : getColorCategories(visData.rows, visData.columns[0]?.id); + const colorIndexMap = (chartType?: ChartTypes) => + new Map(categories(chartType).map((d, i) => [d[0], i])); + const dataMock = dataPluginMock.createStartContract(); + interface RangeProps { + gte: number; + lt: number; + } + const distinctSeries = getDistinctSeries(visData.rows, buckets); + const dataLength = { columnsLength: buckets.length, rowsLength: visData.rows.length }; - dataMock.fieldFormats = { - deserialize: jest.fn(() => ({ - convert: jest.fn((s: RangeProps) => { - return `≥ ${s.gte} and < ${s.lt}`; - }), - })), - } as unknown as DataPublicPluginStart['fieldFormats']; + dataMock.fieldFormats = { + deserialize: jest.fn(() => ({ + convert: jest.fn((s: RangeProps) => { + return `≥ ${s.gte} and < ${s.lt}`; + }), + })), + } as unknown as DataPublicPluginStart['fieldFormats']; - const getPaletteRegistry = () => { - const mockPalette1: jest.Mocked = { - id: 'default', - title: 'My Palette', - getCategoricalColor: jest.fn((layer: SeriesLayer[]) => colors[layer[0].rankAtDepth]), - getCategoricalColors: jest.fn((num: number) => colors), - toExpression: jest.fn(() => ({ - type: 'expression', - chain: [ - { - type: 'function', - function: 'system_palette', - arguments: { - name: ['default'], + const getPaletteRegistry = () => { + const mockPalette1: jest.Mocked = { + id: 'default', + title: 'My Palette', + getCategoricalColor: jest.fn((layer: SeriesLayer[]) => colors[layer[0].rankAtDepth]), + getCategoricalColors: jest.fn((num: number) => colors), + toExpression: jest.fn(() => ({ + type: 'expression', + chain: [ + { + type: 'function', + function: 'system_palette', + arguments: { + name: ['default'], + }, }, - }, - ], - })), - }; + ], + })), + }; - return { - get: () => mockPalette1, - getAll: () => [mockPalette1], + return { + get: () => mockPalette1, + getAll: () => [mockPalette1], + }; }; - }; - it('should return the correct color based on the parent sortIndex', () => { - const d: SimplifiedArrayNode = { - depth: 1, - sortIndex: 0, - parent: { - children: [ - ['ES-Air', undefined], - ['Kibana Airlines', undefined], - ], - depth: 0, + it('should return the correct color based on the parent sortIndex', () => { + const d: SimplifiedArrayNode = { + depth: 1, sortIndex: 0, - }, - children: [], - }; + parent: { + children: [ + ['ES-Air', undefined], + ['Kibana Airlines', undefined], + ], + depth: 0, + sortIndex: 0, + }, + children: [], + }; - const color = getColor( - ChartTypes.PIE, - 'ES-Air', - d, - 0, - false, - {}, - distinctSeries, - dataLength, - visParams, - getPaletteRegistry(), - { getColor: () => undefined }, - false, - false, - dataMock.fieldFormats, - visData.columns[0], - formatters - ); - expect(color).toEqual(colors[0]); - }); + const color = getColor( + ChartTypes.PIE, + 'ES-Air', + d, + 0, + false, + {}, + distinctSeries, + dataLength, + visParams, + getPaletteRegistry(), + { getColor: () => undefined }, + false, + false, + dataMock.fieldFormats, + visData.columns[0], + colorIndexMap(ChartTypes.PIE) + ); + expect(color).toEqual(colors[0]); + }); - it('slices with the same label should have the same color for small multiples', () => { - const d: SimplifiedArrayNode = { - depth: 1, - sortIndex: 0, - parent: { - children: [ - ['ES-Air', undefined], - ['Kibana Airlines', undefined], - ], - depth: 0, - sortIndex: 0, - }, - children: [], - }; - const color = getColor( - ChartTypes.PIE, - 'ES-Air', - d, - 0, - true, - {}, - distinctSeries, - dataLength, - visParams, - getPaletteRegistry(), - { getColor: () => undefined }, - false, - false, - dataMock.fieldFormats, - visData.columns[0], - formatters - ); - expect(color).toEqual('color3'); - }); - it('returns the overwriteColor if exists', () => { - const d: SimplifiedArrayNode = { - depth: 1, - sortIndex: 0, - parent: { - children: [ - ['ES-Air', undefined], - ['Kibana Airlines', undefined], - ], - depth: 0, + it('slices with the same label should have the same color for small multiples', () => { + const d: SimplifiedArrayNode = { + depth: 1, sortIndex: 0, - }, - children: [], - }; - const color = getColor( - ChartTypes.PIE, - 'ES-Air', - d, - 0, - true, - { 'ES-Air': '#000028' }, - distinctSeries, - dataLength, - visParams, - getPaletteRegistry(), - { getColor: () => undefined }, - false, - false, - dataMock.fieldFormats, - visData.columns[0], - formatters - ); - expect(color).toEqual('#000028'); - }); - - it('returns the overwriteColor for older visualizations with formatted values', () => { - const d: SimplifiedArrayNode = { - depth: 1, - sortIndex: 0, - parent: { - children: [ - [ - { - gte: 1000, - lt: 2000, - }.toString(), - undefined, - ], - [ - { - gte: 2000, - lt: 3000, - }.toString(), - undefined, + parent: { + children: [ + ['ES-Air', undefined], + ['Kibana Airlines', undefined], ], - ], - depth: 0, - sortIndex: 0, - }, - children: [], - }; - const visParamsNew = { - ...visParams, - distinctColors: true, - }; - const column = { - ...visData.columns[0], - format: { - id: 'range', - params: { - id: 'number', + depth: 0, + sortIndex: 0, }, - }, - }; - const color = getColor( - ChartTypes.PIE, - // There is the unhandled situation that the categoricalName passed is not a plain string but a RangeKey - // In this case, the internal code, thankfully, requires the stringified version of it and/or the formatted one - // handling also this badly configured type - // FIXME when getColor could handle both strings and RangeKey - { gte: 1000, lt: 2000 } as unknown as string, - d, - 0, - true, - { '≥ 1000 and < 2000': '#3F6833' }, - distinctSeries, - dataLength, - visParamsNew, - getPaletteRegistry(), - { getColor: () => undefined }, - false, - false, - dataMock.fieldFormats, - column, - formatters - ); - expect(color).toEqual('#3F6833'); - }); - - it('should only pass the second layer for mosaic', () => { - const d: SimplifiedArrayNode = { - depth: 2, - sortIndex: 0, - parent: { - children: [ - ['Second level 1', undefined], - ['Second level 2', undefined], - ], + children: [], + }; + const color = getColor( + ChartTypes.PIE, + 'ES-Air', + d, + 0, + true, + {}, + distinctSeries, + dataLength, + visParams, + getPaletteRegistry(), + { getColor: () => undefined }, + false, + false, + dataMock.fieldFormats, + visData.columns[0], + colorIndexMap(ChartTypes.PIE) + ); + expect(color).toEqual('color3'); + }); + it('returns the overwriteColor if exists', () => { + const d: SimplifiedArrayNode = { depth: 1, sortIndex: 0, parent: { - children: [['First level', undefined]], + children: [ + ['ES-Air', undefined], + ['Kibana Airlines', undefined], + ], depth: 0, sortIndex: 0, }, - }, - children: [], - }; - const registry = getPaletteRegistry(); - getColor( - ChartTypes.MOSAIC, - 'Second level 1', - d, - 1, - true, - {}, - distinctSeries, - dataLength, - visParams, - registry, - undefined, - true, - false, - dataMock.fieldFormats, - visData.columns[0], - formatters - ); - expect(registry.get().getCategoricalColor).toHaveBeenCalledWith( - [expect.objectContaining({ name: 'Second level 1' })], - expect.anything(), - expect.anything() - ); - }); + children: [], + }; + const color = getColor( + ChartTypes.PIE, + 'ES-Air', + d, + 0, + true, + { 'ES-Air': '#000028' }, + distinctSeries, + dataLength, + visParams, + getPaletteRegistry(), + { getColor: () => undefined }, + false, + false, + dataMock.fieldFormats, + visData.columns[0], + colorIndexMap(ChartTypes.PIE) + ); + expect(color).toEqual('#000028'); + }); - it('should only pass the first layer for treemap', () => { - const d: SimplifiedArrayNode = { - depth: 2, - sortIndex: 0, - parent: { - children: [ - ['Second level 1', undefined], - ['Second level 2', undefined], - ], + it('returns the overwriteColor for older visualizations with formatted values', () => { + const d: SimplifiedArrayNode = { depth: 1, sortIndex: 0, parent: { - children: [['First level', undefined]], + children: [ + [ + { + gte: 1000, + lt: 2000, + }.toString(), + undefined, + ], + [ + { + gte: 2000, + lt: 3000, + }.toString(), + undefined, + ], + ], depth: 0, sortIndex: 0, }, - }, - children: [], - }; - const registry = getPaletteRegistry(); - getColor( - ChartTypes.TREEMAP, - 'Second level 1', - d, - 1, - true, - {}, - distinctSeries, - dataLength, - visParams, - registry, - undefined, - true, - false, - dataMock.fieldFormats, - visData.columns[0], - formatters - ); - expect(registry.get().getCategoricalColor).toHaveBeenCalledWith( - [expect.objectContaining({ name: 'First level' })], - expect.anything(), - expect.anything() - ); + children: [], + }; + const visParamsNew = { + ...visParams, + distinctColors: true, + }; + const column = { + ...visData.columns[0], + format: { + id: 'range', + params: { + id: 'number', + }, + }, + }; + const color = getColor( + ChartTypes.PIE, + // There is the unhandled situation that the categoricalName passed is not a plain string but a RangeKey + // In this case, the internal code, thankfully, requires the stringified version of it and/or the formatted one + // handling also this badly configured type + // FIXME when getColor could handle both strings and RangeKey + { gte: 1000, lt: 2000 } as unknown as string, + d, + 0, + true, + { '≥ 1000 and < 2000': '#3F6833' }, + distinctSeries, + dataLength, + visParamsNew, + getPaletteRegistry(), + { getColor: () => undefined }, + false, + false, + dataMock.fieldFormats, + column, + colorIndexMap(ChartTypes.PIE) + ); + expect(color).toEqual('#3F6833'); + }); + + it('should only pass the second layer for mosaic', () => { + const d: SimplifiedArrayNode = { + depth: 2, + sortIndex: 0, + parent: { + children: [ + ['Second level 1', undefined], + ['Second level 2', undefined], + ], + depth: 1, + sortIndex: 0, + parent: { + children: [['First level', undefined]], + depth: 0, + sortIndex: 0, + }, + }, + children: [], + }; + const registry = getPaletteRegistry(); + getColor( + ChartTypes.MOSAIC, + 'Second level 1', + d, + 1, + true, + {}, + distinctSeries, + dataLength, + visParams, + registry, + undefined, + true, + false, + dataMock.fieldFormats, + visData.columns[0], + colorIndexMap(ChartTypes.MOSAIC) + ); + expect(registry.get().getCategoricalColor).toHaveBeenCalledWith( + [expect.objectContaining({ name: 'Second level 1' })], + expect.anything(), + expect.anything() + ); + }); }); }); diff --git a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.ts b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.ts index 6bd8d188f78c5..6b1953857a3e2 100644 --- a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.ts +++ b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_color.ts @@ -11,60 +11,42 @@ import { ArrayNode } from '@elastic/charts'; import { isEqual } from 'lodash'; import type { PaletteRegistry, SeriesLayer, PaletteOutput, PaletteDefinition } from '@kbn/coloring'; import type { FieldFormatsStart } from '@kbn/field-formats-plugin/public'; -import type { FieldFormat } from '@kbn/field-formats-plugin/common'; import { lightenColor } from '@kbn/charts-plugin/public'; -import type { Datatable } from '@kbn/expressions-plugin/public'; import { BucketColumns, ChartTypes, PartitionVisParams } from '../../../common/types'; import { DistinctSeries } from '../get_distinct_series'; -import { getNodeLabel } from './get_node_labels'; const isTreemapOrMosaicChart = (shape: ChartTypes) => [ChartTypes.MOSAIC, ChartTypes.TREEMAP].includes(shape); export const byDataColorPaletteMap = ( - rows: Datatable['rows'], - column: Partial, paletteDefinition: PaletteDefinition, { params }: PaletteOutput, - formatters: Record, - formatter: FieldFormatsStart + colorIndexMap: Map ) => { - const colorMap = new Map( - rows.map((item) => { - const formattedName = getNodeLabel( - item[column.id ?? ''], - column, - formatters, - formatter.deserialize - ); - return [formattedName, undefined]; - }) - ); - let rankAtDepth = 0; + const colorCache = new Map(); return { - getColor: (item: unknown) => { + getColor: (item: string) => { const key = String(item); - if (!colorMap.has(key)) return; + let color = colorCache.get(key); + + if (color) return color; - let color = colorMap.get(key); - if (color) { - return color; - } + const colorIndex = colorIndexMap.get(key) ?? -1; color = paletteDefinition.getCategoricalColor( [ { name: key, - totalSeriesAtDepth: colorMap.size, - rankAtDepth: rankAtDepth++, + totalSeriesAtDepth: colorIndexMap.size, + rankAtDepth: colorIndex, }, ], { behindText: false }, params ) || undefined; - colorMap.set(key, color); + colorCache.set(key, color); return color; }, }; @@ -136,21 +118,20 @@ export interface SimplifiedArrayNode { * (a node of a hierarchical tree, currently a partition tree) up to the root of the hierarchy tree. * The resulting array only shows, for each parent, the name of the node, its child index within the parent branch * (called rankInDepth) and the total number of children of the parent. - * */ const createSeriesLayers = ( arrayNode: SimplifiedArrayNode, parentSeries: DistinctSeries['parentSeries'], isSplitChart: boolean, - formatters: Record, - formatter: FieldFormatsStart, - column: Partial + colorIndexMap: Map ): SeriesLayer[] => { const seriesLayers: SeriesLayer[] = []; let tempParent: typeof arrayNode | (typeof arrayNode)['parent'] = arrayNode; + while (tempParent.parent && tempParent.depth > 0) { const nodeKey = tempParent.parent.children[tempParent.sortIndex][0]; const seriesName = String(nodeKey); + /** * FIXME this is a bad implementation: The `parentSeries` is an array of both `string` and `RangeKey` even if its type * is marked as `string[]` in `DistinctSeries`. Here instead we are checking if a stringified `RangeKey` is included into this array that @@ -158,15 +139,14 @@ const createSeriesLayers = ( * see https://github.com/elastic/kibana/issues/153437 */ const isSplitParentLayer = isSplitChart && parentSeries.includes(seriesName); - const formattedName = getNodeLabel(nodeKey, column, formatters, formatter.deserialize); + const colorIndex = colorIndexMap.get(seriesName) ?? tempParent.sortIndex; + seriesLayers.unshift({ - // by construction and types `formattedName` should be always be a string, but I leave this Nullish Coalescing - // because I don't trust much our formatting functions - name: formattedName ?? seriesName, + name: seriesName, rankAtDepth: isSplitParentLayer ? // FIXME as described above this will not work correctly if the `nodeKey` is a `RangeKey` parentSeries.findIndex((name) => name === seriesName) - : tempParent.sortIndex, + : colorIndex, totalSeriesAtDepth: isSplitParentLayer ? parentSeries.length : tempParent.parent.children.length, @@ -213,7 +193,7 @@ export const getColor = ( isDarkMode: boolean, formatter: FieldFormatsStart, column: Partial, - formatters: Record + colorIndexMap: Map ) => { // Mind the difference here: the contrast computation for the text ignores the alpha/opacity // therefore change it for dark mode @@ -242,9 +222,7 @@ export const getColor = ( arrayNode, distinctSeries.parentSeries, isSplitChart, - formatters, - formatter, - column + colorIndexMap ); const overriddenColor = overrideColors(seriesLayers, overwriteColors, name); @@ -257,18 +235,13 @@ export const getColor = ( return byDataPalette.getColor(seriesLayers[1].name) || defaultColor; } - if (isTreemapOrMosaicChart(chartType)) { - if (layerIndex < columnsLength - 1) { - return defaultColor; - } - // for treemap use the top layer for coloring, for mosaic use the second layer - if (seriesLayers.length > 1) { - if (chartType === ChartTypes.MOSAIC) { - seriesLayers.shift(); - } else { - seriesLayers.pop(); - } - } + if (chartType === ChartTypes.MOSAIC && layerIndex < columnsLength - 1) { + return defaultColor; + } + + // Mosaic - use the second layer for color + if (chartType === ChartTypes.MOSAIC && seriesLayers.length > 1) { + seriesLayers.shift(); } const outputColor = paletteService?.get(visParams.palette.name).getCategoricalColor( diff --git a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_layers.ts b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_layers.ts index 4346dd6ae4928..756fa182fb231 100644 --- a/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_layers.ts +++ b/src/platform/plugins/shared/chart_expressions/expression_partition_vis/public/utils/layers/get_layers.ts @@ -53,16 +53,20 @@ export const getLayers = ( fillLabel.valueFormatter = () => ''; } + const categories = + chartType === ChartTypes.MOSAIC && columns.length === 2 + ? getColorCategories(rows, columns[1]?.id) + : getColorCategories(rows, columns[0]?.id); + + const colorIndexMap = new Map(categories.map((c, i) => [String(c), i])); + const isSplitChart = Boolean(visParams.dimensions.splitColumn || visParams.dimensions.splitRow); let byDataPalette: ReturnType; if (!syncColors && columns[1]?.id && paletteService && visParams.palette) { byDataPalette = byDataColorPaletteMap( - rows, - columns[1], paletteService?.get(visParams.palette.name), visParams.palette, - formatters, - formatter + colorIndexMap ); } @@ -113,7 +117,7 @@ export const getLayers = ( isDarkMode, formatter, col, - formatters + colorIndexMap ), }, }; From 4447a7050a67daf801d799190b1533aa60e4c791 Mon Sep 17 00:00:00 2001 From: Krzysztof Kowalczyk Date: Tue, 4 Mar 2025 18:31:52 +0100 Subject: [PATCH 4/9] [Global Search] Add ILM keyword (#213122) ## Summary This PR adds an `ILM` keyword to the ILM application, so searching for the keyword in Global Search would match it. ![Screenshot 2025-03-04 at 16 23 52](https://github.com/user-attachments/assets/36b12d59-62df-41c9-be48-5e2615aad299) Closes: #150424 --- .../plugins/private/index_lifecycle_management/public/plugin.tsx | 1 + 1 file changed, 1 insertion(+) diff --git a/x-pack/platform/plugins/private/index_lifecycle_management/public/plugin.tsx b/x-pack/platform/plugins/private/index_lifecycle_management/public/plugin.tsx index 2fe173e917624..38f7345b7f648 100644 --- a/x-pack/platform/plugins/private/index_lifecycle_management/public/plugin.tsx +++ b/x-pack/platform/plugins/private/index_lifecycle_management/public/plugin.tsx @@ -48,6 +48,7 @@ export class IndexLifecycleManagementPlugin id: PLUGIN.ID, title: PLUGIN.TITLE, order: 2, + keywords: ['ilm'], mount: async ({ element, history, setBreadcrumbs }) => { const [coreStart, { licensing }] = await getStartServices(); const { From d3d44defa4e729e4149d58a08f9ddcdbaf603def Mon Sep 17 00:00:00 2001 From: Nicolas Chaulet Date: Tue, 4 Mar 2025 13:15:34 -0500 Subject: [PATCH 5/9] [Fleet] Use streaming for package install instead of an assetsMap with everything loaded in memory (#211961) --- .../shared/fleet/common/types/models/epm.ts | 10 +- .../routes/epm/kibana_assets_handler.ts | 1 - .../services/epm/archive/archive_iterator.ts | 9 +- .../elasticsearch/datastream_ilm/install.ts | 23 +- .../services/epm/elasticsearch/ilm/install.ts | 18 +- .../ingest_pipeline/install.test.ts | 30 ++- .../elasticsearch/ingest_pipeline/install.ts | 17 +- .../epm/elasticsearch/ml_model/install.ts | 16 +- .../elasticsearch/template/install.test.ts | 36 ++- .../epm/elasticsearch/template/install.ts | 62 ++++- .../epm/elasticsearch/transform/install.ts | 47 +++- .../transform/legacy_transforms.test.ts | 33 +++ .../elasticsearch/transform/mappings.test.ts | 41 ++-- .../epm/elasticsearch/transform/mappings.ts | 9 +- .../transform/transforms.test.ts | 173 +++++++------- .../fleet/server/services/epm/fields/field.ts | 8 +- .../services/epm/kibana/assets/install.ts | 157 ++++++------ .../epm/kibana/assets/tag_assets.test.ts | 226 ++++++++---------- .../services/epm/kibana/assets/tag_assets.ts | 41 ++-- .../server/services/epm/package_service.ts | 1 - .../server/services/epm/packages/install.ts | 11 +- .../install_index_template_pipeline.ts | 2 +- .../_state_machine_package_install.test.ts | 5 - .../step_create_restart_installation.test.ts | 4 - .../steps/step_install_ilm_policies.test.ts | 1 - .../steps/step_install_kibana_assets.test.ts | 4 - .../steps/step_save_archive_entries.ts | 69 +++--- .../steps/step_save_system_object.test.ts | 2 - .../update_latest_executed_state.test.ts | 4 - .../experimental_datastream_features.ts | 2 +- .../epm/__snapshots__/bulk_get_assets.snap | 80 +++---- .../apis/epm/update_assets.ts | 93 +++---- 32 files changed, 694 insertions(+), 541 deletions(-) diff --git a/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts b/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts index 512a37a327618..73791e6c74b10 100644 --- a/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts +++ b/x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts @@ -132,17 +132,15 @@ export interface ArchiveEntry { } export interface ArchiveIterator { - traverseEntries: (onEntry: (entry: ArchiveEntry) => Promise) => Promise; + traverseEntries: ( + onEntry: (entry: ArchiveEntry) => Promise, + readBuffer?: (path: string) => boolean + ) => Promise; getPaths: () => Promise; } export interface PackageInstallContext { packageInfo: InstallablePackage; - /** - * @deprecated Use `archiveIterator` to access the package archive entries - * without loading them all into memory at once. - */ - assetsMap: AssetsMap; paths: string[]; archiveIterator: ArchiveIterator; } diff --git a/x-pack/platform/plugins/shared/fleet/server/routes/epm/kibana_assets_handler.ts b/x-pack/platform/plugins/shared/fleet/server/routes/epm/kibana_assets_handler.ts index 57880d5f08397..550377a058371 100644 --- a/x-pack/platform/plugins/shared/fleet/server/routes/epm/kibana_assets_handler.ts +++ b/x-pack/platform/plugins/shared/fleet/server/routes/epm/kibana_assets_handler.ts @@ -68,7 +68,6 @@ export const installPackageKibanaAssetsHandler: FleetRequestHandler< packageInstallContext: { packageInfo, paths: installedPkgWithAssets.paths, - assetsMap: installedPkgWithAssets.assetsMap, archiveIterator: createArchiveIteratorFromMap(installedPkgWithAssets.assetsMap), }, }); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/archive_iterator.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/archive_iterator.ts index bf7fda919e9df..4775add398392 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/archive_iterator.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/archive/archive_iterator.ts @@ -74,10 +74,15 @@ export const createArchiveIterator = ( */ export const createArchiveIteratorFromMap = (assetsMap: AssetsMap): ArchiveIterator => { const traverseEntries = async ( - onEntry: (entry: ArchiveEntry) => Promise + onEntry: (entry: ArchiveEntry) => Promise, + readBuffer?: (path: string) => boolean ): Promise => { for (const [path, buffer] of assetsMap) { - await onEntry({ path, buffer }); + if (readBuffer && !readBuffer(path)) { + await onEntry({ path }); + } else { + await onEntry({ path, buffer }); + } } }; diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/datastream_ilm/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/datastream_ilm/install.ts index eafda9a99e4fc..90f7201112381 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/datastream_ilm/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/datastream_ilm/install.ts @@ -12,7 +12,11 @@ import { ElasticsearchAssetType, type PackageInstallContext, } from '../../../../../common/types/models'; -import type { EsAssetReference, RegistryDataStream } from '../../../../../common/types/models'; +import type { + AssetsMap, + EsAssetReference, + RegistryDataStream, +} from '../../../../../common/types/models'; import { updateEsAssetReferences } from '../../packages/es_assets_reference'; import { getAssetFromAssetsMap } from '../../archive'; @@ -40,7 +44,7 @@ export const installIlmForDataStream = async ( logger: Logger, esReferences: EsAssetReference[] ) => { - const { packageInfo: registryPackage, paths, assetsMap } = packageInstallContext; + const { packageInfo: registryPackage, paths } = packageInstallContext; const previousInstalledIlmEsAssets = esReferences.filter( ({ type }) => type === ElasticsearchAssetType.dataStreamIlmPolicy ); @@ -72,6 +76,19 @@ export const installIlmForDataStream = async ( }; const dataStreamIlmPaths = paths.filter((path) => isDataStreamIlm(path)); + + const dataStreamIlmAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + dataStreamIlmAssetsMap.set(entry.path, entry.buffer); + }, + (path) => dataStreamIlmPaths.includes(path) + ); + let installedIlms: EsAssetReference[] = []; if (dataStreamIlmPaths.length > 0) { const ilmPathDatasets = dataStreams.reduce((acc, dataStream) => { @@ -103,7 +120,7 @@ export const installIlmForDataStream = async ( const ilmInstallations: IlmInstallation[] = ilmPathDatasets.map( (ilmPathDataset: IlmPathDataset) => { const content = JSON.parse( - getAssetFromAssetsMap(assetsMap, ilmPathDataset.path).toString('utf-8') + getAssetFromAssetsMap(dataStreamIlmAssetsMap, ilmPathDataset.path).toString('utf-8') ); content.policy._meta = getESAssetMetadata({ packageName: registryPackage.name }); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ilm/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ilm/install.ts index 1203ec02ba9ac..eb40ae1590524 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ilm/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ilm/install.ts @@ -16,7 +16,7 @@ import { updateEsAssetReferences } from '../../packages/es_assets_reference'; import { getESAssetMetadata } from '../meta'; import { retryTransientEsErrors } from '../retry'; import { PackageInvalidArchiveError } from '../../../../errors'; -import type { PackageInstallContext } from '../../../../../common/types'; +import type { AssetsMap, PackageInstallContext } from '../../../../../common/types'; import { MAX_CONCURRENT_ILM_POLICIES_OPERATIONS } from '../../../../constants'; export async function installILMPolicy( @@ -30,10 +30,20 @@ export async function installILMPolicy( const ilmPaths = packageInstallContext.paths.filter((path) => isILMPolicy(path)); if (!ilmPaths.length) return esReferences; + const ilmAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + ilmAssetsMap.set(entry.path, entry.buffer); + }, + (path) => ilmPaths.includes(path) + ); + const ilmPolicies = ilmPaths.map((path) => { - const body = JSON.parse( - getAssetFromAssetsMap(packageInstallContext.assetsMap, path).toString('utf-8') - ); + const body = JSON.parse(getAssetFromAssetsMap(ilmAssetsMap, path).toString('utf-8')); body.policy._meta = getESAssetMetadata({ packageName: packageInfo.name }); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.test.ts index fb12d58878d1d..9d12bd7a5e8da 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.test.ts @@ -7,6 +7,8 @@ import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; import { loggerMock } from '@kbn/logging-mocks'; +import { createArchiveIteratorFromMap } from '../../archive/archive_iterator'; + import { prepareToInstallPipelines } from './install'; jest.mock('../../archive/cache'); @@ -25,10 +27,10 @@ describe('Install pipeline tests', () => { path: '/datasettest', }, ], - }, + } as any, paths: [], - assetsMap: new Map(), - } as any); + archiveIterator: createArchiveIteratorFromMap(new Map()), + }); expect(res.assetsToAdd).toEqual([{ id: 'logs-datasettest-1.0.0', type: 'ingest_pipeline' }]); const esClient = elasticsearchClientMock.createInternalClient(); @@ -64,16 +66,18 @@ describe('Install pipeline tests', () => { 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/default.yml', 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/standard.yml', ], - assetsMap: new Map([ - [ - 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/default.yml', - Buffer.from('description: test'), - ], - [ - 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/standard.yml', - Buffer.from('description: test'), - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/default.yml', + Buffer.from('description: test'), + ], + [ + 'packagetest-1.0.0/data_stream/datasettest/elasticsearch/ingest_pipeline/standard.yml', + Buffer.from('description: test'), + ], + ]) + ), } as any); expect(res.assetsToAdd).toEqual([ { id: 'logs-datasettest-1.0.0', type: 'ingest_pipeline' }, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.ts index 51162ac2c6335..47c060cf31caf 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ingest_pipeline/install.ts @@ -25,7 +25,7 @@ import { FLEET_EVENT_INGESTED_PIPELINE_CONTENT, } from '../../../../constants'; import { getPipelineNameForDatastream } from '../../../../../common/services'; -import type { ArchiveEntry, PackageInstallContext } from '../../../../../common/types'; +import type { ArchiveEntry, AssetsMap, PackageInstallContext } from '../../../../../common/types'; import { appendMetadataToIngestPipeline } from '../meta'; import { retryTransientEsErrors } from '../retry'; @@ -157,6 +157,19 @@ export async function installAllPipelines({ > = []; const substitutions: RewriteSubstitution[] = []; + const pipelineAssetsMap: AssetsMap = new Map(); + + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + pipelineAssetsMap.set(entry.path, entry.buffer); + }, + (path) => pipelinePaths.includes(path) + ); + let datastreamPipelineCreated = false; pipelinePaths.forEach((path) => { const { name, extension } = getNameAndExtension(path); @@ -169,7 +182,7 @@ export async function installAllPipelines({ dataStream, packageVersion: packageInstallContext.packageInfo.version, }); - const content = getAssetFromAssetsMap(packageInstallContext.assetsMap, path).toString('utf-8'); + const content = getAssetFromAssetsMap(pipelineAssetsMap, path).toString('utf-8'); pipelinesInfos.push({ nameForInstallation, shouldInstallCustomPipelines: dataStream && isMainPipeline, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ml_model/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ml_model/install.ts index e0c5d1b0088e4..16fceaaaeb0ad 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ml_model/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/ml_model/install.ts @@ -13,7 +13,7 @@ import { ElasticsearchAssetType, type PackageInstallContext, } from '../../../../../common/types/models'; -import type { EsAssetReference } from '../../../../../common/types/models'; +import type { AssetsMap, EsAssetReference } from '../../../../../common/types/models'; import { retryTransientEsErrors } from '../retry'; @@ -34,9 +34,19 @@ export const installMlModel = async ( const mlModelPath = packageInstallContext.paths.find((path) => isMlModel(path)); if (mlModelPath !== undefined) { - const content = getAssetFromAssetsMap(packageInstallContext.assetsMap, mlModelPath).toString( - 'utf-8' + const mlModelAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + mlModelAssetsMap.set(entry.path, entry.buffer); + }, + (path) => path === mlModelPath ); + + const content = getAssetFromAssetsMap(mlModelAssetsMap, mlModelPath).toString('utf-8'); const pathParts = mlModelPath.split('/'); const modelId = pathParts[pathParts.length - 1].replace('.json', ''); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.test.ts index 910d5b6cc98ca..29cea7ecc167a 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.test.ts @@ -10,6 +10,7 @@ import { createAppContextStartContractMock } from '../../../../mocks'; import { appContextService } from '../../..'; import { loadDatastreamsFieldsFromYaml } from '../../fields/field'; import type { PackageInstallContext, RegistryDataStream } from '../../../../../common/types'; +import { createArchiveIteratorFromMap } from '../../archive/archive_iterator'; import { prepareTemplate, prepareToInstallTemplates } from './install'; @@ -53,7 +54,11 @@ describe('EPM index template install', () => { const templatePriorityDatasetIsPrefixUnset = 200; const { indexTemplate: { indexTemplate }, - } = prepareTemplate({ packageInstallContext, dataStream: dataStreamDatasetIsPrefixUnset }); + } = await prepareTemplate({ + packageInstallContext, + fieldAssetsMap: new Map(), + dataStream: dataStreamDatasetIsPrefixUnset, + }); expect(indexTemplate.priority).toBe(templatePriorityDatasetIsPrefixUnset); expect(indexTemplate.index_patterns).toEqual([templateIndexPatternDatasetIsPrefixUnset]); }); @@ -74,7 +79,11 @@ describe('EPM index template install', () => { const templatePriorityDatasetIsPrefixFalse = 200; const { indexTemplate: { indexTemplate }, - } = prepareTemplate({ packageInstallContext, dataStream: dataStreamDatasetIsPrefixFalse }); + } = prepareTemplate({ + packageInstallContext, + fieldAssetsMap: new Map(), + dataStream: dataStreamDatasetIsPrefixFalse, + }); expect(indexTemplate.priority).toBe(templatePriorityDatasetIsPrefixFalse); expect(indexTemplate.index_patterns).toEqual([templateIndexPatternDatasetIsPrefixFalse]); @@ -96,7 +105,11 @@ describe('EPM index template install', () => { const templatePriorityDatasetIsPrefixTrue = 150; const { indexTemplate: { indexTemplate }, - } = prepareTemplate({ packageInstallContext, dataStream: dataStreamDatasetIsPrefixTrue }); + } = prepareTemplate({ + packageInstallContext, + fieldAssetsMap: new Map(), + dataStream: dataStreamDatasetIsPrefixTrue, + }); expect(indexTemplate.priority).toBe(templatePriorityDatasetIsPrefixTrue); expect(indexTemplate.index_patterns).toEqual([templateIndexPatternDatasetIsPrefixTrue]); @@ -119,6 +132,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream: dataStreamDatasetIsPrefixTrue, }); @@ -149,6 +163,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream: dataStreamDatasetIsPrefixTrue, }); @@ -180,6 +195,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, dataStream: dataStreamDatasetIsPrefixTrue, + fieldAssetsMap: new Map(), experimentalDataStreamFeature: { data_stream: 'metrics-package.dataset', features: { @@ -218,6 +234,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, dataStream: dataStreamDatasetIsPrefixTrue, + fieldAssetsMap: new Map(), experimentalDataStreamFeature: { data_stream: 'metrics-package.dataset', features: { @@ -255,6 +272,7 @@ describe('EPM index template install', () => { const { indexTemplate } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream: dataStreamDatasetIsPrefixTrue, }); @@ -286,6 +304,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, dataStream, + fieldAssetsMap: new Map(), }); const packageTemplate = componentTemplates['logs-package.dataset@package'].template; @@ -313,6 +332,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, dataStream, + fieldAssetsMap: new Map(), }); const packageTemplate = componentTemplates['logs-package.dataset@package'].template; @@ -350,6 +370,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream, }); @@ -388,6 +409,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream, }); @@ -428,6 +450,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream, }); @@ -472,6 +495,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream, }); @@ -507,6 +531,7 @@ describe('EPM index template install', () => { const { componentTemplates } = prepareTemplate({ packageInstallContext, + fieldAssetsMap: new Map(), dataStream, }); @@ -515,7 +540,7 @@ describe('EPM index template install', () => { expect(packageTemplate).not.toHaveProperty('lifecycle'); }); - test('test prepareToInstallTemplates does not include stack component templates in tracked assets', () => { + test('test prepareToInstallTemplates does not include stack component templates in tracked assets', async () => { const dataStreamDatasetIsPrefixUnset = { type: 'logs', dataset: 'package.dataset', @@ -526,13 +551,14 @@ describe('EPM index template install', () => { ingest_pipeline: 'default', } as RegistryDataStream; - const { assetsToAdd } = prepareToInstallTemplates( + const { assetsToAdd } = await prepareToInstallTemplates( { packageInfo: { name: 'package', version: '0.0.1', data_streams: [dataStreamDatasetIsPrefixUnset], }, + archiveIterator: createArchiveIteratorFromMap(new Map()), } as PackageInstallContext, [], [] diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.ts index cac196fcb50c2..6540f27ea182a 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/template/install.ts @@ -32,7 +32,7 @@ import type { ExperimentalDataStreamFeature, } from '../../../../types'; import type { Fields } from '../../fields/field'; -import { loadDatastreamsFieldsFromYaml, processFields } from '../../fields/field'; +import { isFields, loadDatastreamsFieldsFromYaml, processFields } from '../../fields/field'; import { getAssetFromAssetsMap, getPathParts } from '../../archive'; import { FLEET_COMPONENT_TEMPLATES, @@ -48,7 +48,7 @@ import { forEachMappings, } from '../../../experimental_datastream_features_helper'; import { appContextService } from '../../../app_context'; -import type { PackageInstallContext } from '../../../../../common/types'; +import type { AssetsMap, PackageInstallContext } from '../../../../../common/types'; import { generateMappings, @@ -62,16 +62,16 @@ import { isUserSettingsTemplate } from './utils'; const FLEET_COMPONENT_TEMPLATE_NAMES = FLEET_COMPONENT_TEMPLATES.map((tmpl) => tmpl.name); -export const prepareToInstallTemplates = ( +export const prepareToInstallTemplates = async ( packageInstallContext: PackageInstallContext, esReferences: EsAssetReference[], experimentalDataStreamFeatures: ExperimentalDataStreamFeature[] = [], onlyForDataStreams?: RegistryDataStream[] -): { +): Promise<{ assetsToAdd: EsAssetReference[]; assetsToRemove: EsAssetReference[]; install: (esClient: ElasticsearchClient, logger: Logger) => Promise; -} => { +}> => { const { packageInfo } = packageInstallContext; // remove package installation's references to index templates const assetsToRemove = esReferences.filter( @@ -80,6 +80,13 @@ export const prepareToInstallTemplates = ( type === ElasticsearchAssetType.componentTemplate ); + const fieldAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries(async (entry) => { + if (entry.buffer) { + fieldAssetsMap.set(entry.path, entry.buffer); + } + }, isFields); + // build templates per data stream from yml files const dataStreams = onlyForDataStreams || packageInfo.data_streams; if (!dataStreams) return { assetsToAdd: [], assetsToRemove, install: () => Promise.resolve([]) }; @@ -90,7 +97,12 @@ export const prepareToInstallTemplates = ( datastreamFeature.data_stream === getRegistryDataStreamAssetBaseName(dataStream) ); - return prepareTemplate({ packageInstallContext, dataStream, experimentalDataStreamFeature }); + return prepareTemplate({ + packageInstallContext, + fieldAssetsMap, + dataStream, + experimentalDataStreamFeature, + }); }); const assetsToAdd = getAllTemplateRefs(templates.map((template) => template.indexTemplate)); @@ -131,14 +143,23 @@ const installPreBuiltTemplates = async ( ) => { const templatePaths = packageInstallContext.paths.filter((path) => isTemplate(path)); try { + const templateAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + templateAssetsMap.set(entry.path, entry.buffer); + }, + (path) => templatePaths.includes(path) + ); await pMap( templatePaths, async (path) => { const { file } = getPathParts(path); const templateName = file.substr(0, file.lastIndexOf('.')); - const content = JSON.parse( - getAssetFromAssetsMap(packageInstallContext.assetsMap, path).toString('utf8') - ); + const content = JSON.parse(getAssetFromAssetsMap(templateAssetsMap, path).toString('utf8')); const esClientParams = { name: templateName, body: content }; const esClientRequestOptions = { ignore: [404] }; @@ -175,14 +196,23 @@ const installPreBuiltComponentTemplates = async ( ) => { const templatePaths = packageInstallContext.paths.filter((path) => isComponentTemplate(path)); try { + const templateAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + templateAssetsMap.set(entry.path, entry.buffer); + }, + (path) => templatePaths.includes(path) + ); await pMap( templatePaths, async (path) => { const { file } = getPathParts(path); const templateName = file.substr(0, file.lastIndexOf('.')); - const content = JSON.parse( - getAssetFromAssetsMap(packageInstallContext.assetsMap, path).toString('utf8') - ); + const content = JSON.parse(getAssetFromAssetsMap(templateAssetsMap, path).toString('utf8')); const esClientParams = { name: templateName, @@ -573,15 +603,21 @@ function countFields(fields: Fields): number { export function prepareTemplate({ packageInstallContext, + fieldAssetsMap, dataStream, experimentalDataStreamFeature, }: { packageInstallContext: PackageInstallContext; + fieldAssetsMap: AssetsMap; dataStream: RegistryDataStream; experimentalDataStreamFeature?: ExperimentalDataStreamFeature; }): { componentTemplates: TemplateMap; indexTemplate: IndexTemplateEntry } { const { name: packageName, version: packageVersion } = packageInstallContext.packageInfo; - const fields = loadDatastreamsFieldsFromYaml(packageInstallContext, dataStream.path); + const fields = loadDatastreamsFieldsFromYaml( + packageInstallContext, + fieldAssetsMap, + dataStream.path + ); const isIndexModeTimeSeries = dataStream.elasticsearch?.index_mode === 'time_series' || diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/install.ts index 4f1e3ab262772..dac4b02dbe098 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/install.ts @@ -40,6 +40,7 @@ import type { ESAssetMetadata, IndexTemplate, RegistryElasticsearch, + AssetsMap, } from '../../../../../common/types/models'; import { getInstallation } from '../../packages'; import { retryTransientEsErrors } from '../retry'; @@ -93,6 +94,17 @@ const installLegacyTransformsAssets = async ( let installedTransforms: EsAssetReference[] = []; if (transformPaths.length > 0) { + const transformAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + transformAssetsMap.set(entry.path, entry.buffer); + }, + (path) => transformPaths.includes(path) + ); const transformRefs = transformPaths.reduce((acc, path) => { acc.push({ id: getLegacyTransformNameForInstallation( @@ -117,9 +129,7 @@ const installLegacyTransformsAssets = async ( ); const transforms: TransformInstallation[] = transformPaths.map((path: string) => { - const content = JSON.parse( - getAssetFromAssetsMap(packageInstallContext.assetsMap, path).toString('utf-8') - ); + const content = JSON.parse(getAssetFromAssetsMap(transformAssetsMap, path).toString('utf-8')); content._meta = getESAssetMetadata({ packageName: packageInstallContext.packageInfo.name }); return { @@ -153,7 +163,7 @@ const installLegacyTransformsAssets = async ( return { installedTransforms, esReferences }; }; -const processTransformAssetsPerModule = ( +const processTransformAssetsPerModule = async ( packageInstallContext: PackageInstallContext, installNameSuffix: string, transformPaths: string[], @@ -161,7 +171,7 @@ const processTransformAssetsPerModule = ( force?: boolean, username?: string ) => { - const { assetsMap, packageInfo: installablePackage } = packageInstallContext; + const { packageInfo: installablePackage } = packageInstallContext; const transformsSpecifications = new Map(); const destinationIndexTemplates: DestinationIndexTemplateInstallation[] = []; const transforms: TransformInstallation[] = []; @@ -170,6 +180,17 @@ const processTransformAssetsPerModule = ( const transformsToRemoveWithDestIndex: EsAssetReference[] = []; const indicesToAddRefs: EsAssetReference[] = []; + const transformAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries( + async (entry) => { + if (!entry.buffer) { + return; + } + + transformAssetsMap.set(entry.path, entry.buffer); + }, + (path) => transformPaths.includes(path) + ); transformPaths.forEach((path: string) => { const { transformModuleId, fileName } = getTransformFolderAndFileNames( installablePackage, @@ -182,7 +203,7 @@ const processTransformAssetsPerModule = ( } const packageAssets = transformsSpecifications.get(transformModuleId); - const content = load(getAssetFromAssetsMap(assetsMap, path).toString('utf-8')); + const content = load(getAssetFromAssetsMap(transformAssetsMap, path).toString('utf-8')); // Handling fields.yml and all other files within 'fields' folder if (fileName === TRANSFORM_SPECS_TYPES.FIELDS || isFields(path)) { @@ -387,6 +408,12 @@ const processTransformAssetsPerModule = ( version: t.transformVersion, })); + const fieldAssetsMap: AssetsMap = new Map(); + await packageInstallContext.archiveIterator.traverseEntries(async (entry) => { + if (entry.buffer) { + fieldAssetsMap.set(entry.path, entry.buffer); + } + }, isFields); // Load and generate mappings for (const destinationIndexTemplate of destinationIndexTemplates) { if (!destinationIndexTemplate.transformModuleId) { @@ -397,7 +424,11 @@ const processTransformAssetsPerModule = ( .get(destinationIndexTemplate.transformModuleId) ?.set( 'mappings', - loadMappingForTransform(packageInstallContext, destinationIndexTemplate.transformModuleId) + loadMappingForTransform( + packageInstallContext, + fieldAssetsMap, + destinationIndexTemplate.transformModuleId + ) ); } @@ -441,7 +472,7 @@ const installTransformsAssets = async ( transformsSpecifications, transformsToRemove, transformsToRemoveWithDestIndex, - } = processTransformAssetsPerModule( + } = await processTransformAssetsPerModule( packageInstallContext, installNameSuffix, transformPaths, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/legacy_transforms.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/legacy_transforms.test.ts index e705b10ebe011..7ec38a4935c24 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/legacy_transforms.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/legacy_transforms.test.ts @@ -26,6 +26,7 @@ import { appContextService } from '../../../app_context'; import { getESAssetMetadata } from '../meta'; +import { createArchiveIteratorFromMap } from '../../archive/archive_iterator'; import type { PackageInstallContext } from '../../../../../common/types'; import { PACKAGES_SAVED_OBJECT_TYPE } from '../../../../constants'; @@ -167,6 +168,22 @@ describe('test transform install with legacy schema', () => { Buffer.from('{"content": "data"}'), ], ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/data_stream/policy/elasticsearch/ingest_pipeline/default.json', + Buffer.from('{"content": "data"}'), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata/default.json', + Buffer.from('{"content": "data"}'), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/default.json', + Buffer.from('{"content": "data"}'), + ], + ]) + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -339,6 +356,14 @@ describe('test transform install with legacy schema', () => { Buffer.from('{"content": "data"}'), ], ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/default.json', + Buffer.from('{"content": "data"}'), + ], + ]) + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -566,6 +591,14 @@ describe('test transform install with legacy schema', () => { Buffer.from('{"content": "data"}'), ], ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/default.json', + Buffer.from('{"content": "data"}'), + ], + ]) + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.test.ts index de962850fba8c..48cccc79f942a 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.test.ts @@ -14,10 +14,10 @@ describe('loadMappingForTransform', () => { const fields = loadMappingForTransform( { packageInfo: {} as any, - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], }, + new Map(), 'test' ); @@ -28,36 +28,37 @@ describe('loadMappingForTransform', () => { const fields = loadMappingForTransform( { packageInfo: {} as any, - assetsMap: new Map([ - [ - '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs.yml', - Buffer.from( - ` + + archiveIterator: createArchiveIteratorFromMap(new Map()), + paths: [ + '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs.yml', + '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs-extra.yml', + ], + }, + new Map([ + [ + '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs.yml', + Buffer.from( + ` - description: Description of the threat feed in a UI friendly format. name: threat.feed.description type: keyword - description: The name of the threat feed in UI friendly format. name: threat.feed.name type: keyword` - ), - ], - [ - '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs-extra.yml', - Buffer.from( - ` + ), + ], + [ + '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs-extra.yml', + Buffer.from( + ` - description: The display name indicator in an UI friendly format level: extended name: threat.indicator.name type: keyword` - ), - ], - ]), - archiveIterator: createArchiveIteratorFromMap(new Map()), - paths: [ - '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs.yml', - '/package/ti_opencti/2.1.0/elasticsearch/transform/latest_ioc/fields/ecs-extra.yml', + ), ], - }, + ]), 'latest_ioc' ); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.ts index 130dae0ecca51..6fec40c435c42 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/mappings.ts @@ -4,15 +4,20 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import { type PackageInstallContext } from '../../../../../common/types/models'; +import type { AssetsMap, PackageInstallContext } from '../../../../../common/types/models'; import { loadTransformFieldsFromYaml, processFields } from '../../fields/field'; import { generateMappings } from '../template/template'; export function loadMappingForTransform( packageInstallContext: PackageInstallContext, + fieldAssetsMap: AssetsMap, transformModuleId: string ) { - const fields = loadTransformFieldsFromYaml(packageInstallContext, transformModuleId); + const fields = loadTransformFieldsFromYaml( + packageInstallContext, + fieldAssetsMap, + transformModuleId + ); const validFields = processFields(fields); return generateMappings(validFields); } diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/transforms.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/transforms.test.ts index d5a41c1033710..09ac73d796ce0 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/transforms.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/elasticsearch/transform/transforms.test.ts @@ -22,6 +22,7 @@ import { PACKAGES_SAVED_OBJECT_TYPE } from '../../../../constants'; import { getESAssetMetadata } from '../meta'; +import { createArchiveIteratorFromMap } from '../../archive/archive_iterator'; import { createAppContextStartContractMock } from '../../../../mocks'; import type { PackageInstallContext } from '../../../../../common/types'; @@ -268,28 +269,30 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/beats.yml', - sourceData.BEATS_FIELDS, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/agent.yml', - sourceData.AGENT_FIELDS, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', - sourceData.FIELDS, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', - sourceData.MANIFEST, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - sourceData.TRANSFORM, - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/beats.yml', + Buffer.from(sourceData.BEATS_FIELDS), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/agent.yml', + Buffer.from(sourceData.AGENT_FIELDS), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', + Buffer.from(sourceData.FIELDS), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', + Buffer.from(sourceData.MANIFEST), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + Buffer.from(sourceData.TRANSFORM), + ], + ]) + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -570,20 +573,22 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', - Buffer.from(sourceData.FIELDS), - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', - Buffer.from(sourceData.MANIFEST), - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - Buffer.from(sourceData.TRANSFORM), - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', + Buffer.from(sourceData.FIELDS), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', + Buffer.from(sourceData.MANIFEST), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + Buffer.from(sourceData.TRANSFORM), + ], + ]) + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -852,16 +857,18 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', - Buffer.from(sourceData.FIELDS), - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - Buffer.from(sourceData.TRANSFORM), - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', + Buffer.from(sourceData.FIELDS), + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + Buffer.from(sourceData.TRANSFORM), + ], + ]) as any + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -1080,16 +1087,18 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', - sourceData.MANIFEST, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - sourceData.TRANSFORM, - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', + sourceData.MANIFEST, + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + sourceData.TRANSFORM, + ], + ]) as any + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -1184,16 +1193,18 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', - sourceData.MANIFEST, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - sourceData.TRANSFORM, - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', + sourceData.MANIFEST, + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + sourceData.TRANSFORM, + ], + ]) as any + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, @@ -1279,20 +1290,22 @@ _meta: 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', ], - assetsMap: new Map([ - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', - sourceData.FIELDS, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', - sourceData.MANIFEST, - ], - [ - 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', - sourceData.TRANSFORM, - ], - ]), + archiveIterator: createArchiveIteratorFromMap( + new Map([ + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/fields/fields.yml', + sourceData.FIELDS, + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/manifest.yml', + sourceData.MANIFEST, + ], + [ + 'endpoint-0.16.0-dev.0/elasticsearch/transform/metadata_current/transform.yml', + sourceData.TRANSFORM, + ], + ]) as any + ), } as unknown as PackageInstallContext, esClient, savedObjectsClient, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/fields/field.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/fields/field.ts index a3ebf58d02e3b..30482541ff7a1 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/fields/field.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/fields/field.ts @@ -7,7 +7,7 @@ import { load } from 'js-yaml'; -import type { PackageInstallContext } from '../../../../common/types'; +import type { AssetsMap, PackageInstallContext } from '../../../../common/types'; import { getAssetsDataFromAssetsMap } from '../packages/assets'; // This should become a copy of https://github.com/elastic/beats/blob/d9a4c9c240a9820fab15002592e5bb6db318543b/libbeat/mapping/field.go#L39 @@ -310,12 +310,13 @@ function combineFilter(...filters: Array<(path: string) => boolean>) { export const loadDatastreamsFieldsFromYaml = ( packageInstallContext: PackageInstallContext, + fieldAssetsMap: AssetsMap, datasetName?: string ): Field[] => { // Fetch all field definition files const fieldDefinitionFiles = getAssetsDataFromAssetsMap( packageInstallContext.packageInfo, - packageInstallContext.assetsMap, + fieldAssetsMap, isFields, datasetName ); @@ -334,12 +335,13 @@ export const loadDatastreamsFieldsFromYaml = ( export const loadTransformFieldsFromYaml = ( packageInstallContext: PackageInstallContext, + fieldAssetsMap: AssetsMap, transformName: string ): Field[] => { // Fetch all field definition files const fieldDefinitionFiles = getAssetsDataFromAssetsMap( packageInstallContext.packageInfo, - packageInstallContext.assetsMap, + fieldAssetsMap, combineFilter(isFields, filterForTransformAssets(transformName)) ); return fieldDefinitionFiles.reduce((acc, file) => { diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/install.ts index 5adf1708eb25a..0ca3e79f58616 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/install.ts @@ -17,9 +17,9 @@ import type { Logger, } from '@kbn/core/server'; import { createListStream } from '@kbn/utils'; -import { partition, chunk } from 'lodash'; +import { partition, chunk, once } from 'lodash'; -import { getAssetFromAssetsMap, getPathParts } from '../../archive'; +import { getPathParts } from '../../archive'; import { KibanaAssetType, KibanaSavedObjectType } from '../../../../types'; import type { AssetReference, Installation, PackageSpecTags } from '../../../../types'; import type { KibanaAssetReference, PackageInstallContext } from '../../../../../common/types'; @@ -28,7 +28,7 @@ import { getIndexPatternSavedObjects, makeManagedIndexPatternsGlobal, } from '../index_pattern/install'; -import { kibanaAssetsToAssetsRef, saveKibanaAssetsRefs } from '../../packages/install'; +import { saveKibanaAssetsRefs } from '../../packages/install'; import { deleteKibanaSavedObjectsAssets } from '../../packages/remove'; import { FleetError, KibanaSOReferenceError } from '../../../../errors'; import { withPackageSpan } from '../../packages/utils'; @@ -38,7 +38,7 @@ import { appContextService } from '../../..'; import { tagKibanaAssets } from './tag_assets'; import { getSpaceAwareSaveobjectsClients } from './saved_objects'; -const MAX_ASSETS_TO_INSTALL_IN_PARALLEL = 1000; +const MAX_ASSETS_TO_INSTALL_IN_PARALLEL = 200; type SavedObjectsImporterContract = Pick; const formatImportErrorsForLog = (errors: SavedObjectsImportFailure[]) => @@ -111,42 +111,51 @@ export async function installKibanaAssets(options: { savedObjectsImporter: SavedObjectsImporterContract; logger: Logger; pkgName: string; - kibanaAssets: Record; + kibanaAssetsArchiveIterator: ReturnType; }): Promise { - const { kibanaAssets, savedObjectsClient, savedObjectsImporter, logger } = options; + const { kibanaAssetsArchiveIterator, savedObjectsClient, savedObjectsImporter, logger } = options; - const assetsToInstall = Object.entries(kibanaAssets).flatMap(([assetType, assets]) => { - if (!validKibanaAssetTypes.has(assetType as KibanaAssetType)) { - return []; - } + let assetsToInstall: ArchiveAsset[] = []; + let res: SavedObjectsImportSuccess[] = []; - if (!assets.length) { - return []; - } + const installManagedIndexPatternOnce = once(() => + installManagedIndexPattern({ + savedObjectsClient, + savedObjectsImporter, + }) + ); + + async function flushAssetsToInstall() { + await installManagedIndexPatternOnce(); + + const installedAssets = await installKibanaSavedObjects({ + logger, + savedObjectsImporter, + kibanaAssets: assetsToInstall, + assetsChunkSize: MAX_ASSETS_TO_INSTALL_IN_PARALLEL, + }); + assetsToInstall = []; + res = [...res, ...installedAssets]; + } + await kibanaAssetsArchiveIterator(async ({ assetType, asset }) => { const assetFilter = AssetFilters[assetType]; if (assetFilter) { - return assetFilter(assets); + assetsToInstall = [...assetsToInstall, ...assetFilter([asset])]; + } else { + assetsToInstall.push(asset); } - return assets; + if (assetsToInstall.length >= MAX_ASSETS_TO_INSTALL_IN_PARALLEL) { + await flushAssetsToInstall(); + } }); - if (!assetsToInstall.length) { - return []; + if (assetsToInstall.length) { + await flushAssetsToInstall(); } - await installManagedIndexPattern({ - savedObjectsClient, - savedObjectsImporter, - }); - - return await installKibanaSavedObjects({ - logger, - savedObjectsImporter, - kibanaAssets: assetsToInstall, - assetsChunkSize: MAX_ASSETS_TO_INSTALL_IN_PARALLEL, - }); + return res; } export async function installManagedIndexPattern({ @@ -267,50 +276,41 @@ export async function installKibanaAssetsAndReferences({ const { savedObjectsImporter, savedObjectTagAssignmentService, savedObjectTagClient } = getSpaceAwareSaveobjectsClients(spaceId); // This is where the memory consumption is rising up in the first place - const kibanaAssets = getKibanaAssets(packageInstallContext); + const kibanaAssetsArchiveIterator = getKibanaAssetsArchiveIterator(packageInstallContext); + if (installedPkg) { await deleteKibanaSavedObjectsAssets({ installedPkg, spaceId }); } let installedKibanaAssetsRefs: KibanaAssetReference[] = []; - if (!installAsAdditionalSpace) { - // save new kibana refs before installing the assets - installedKibanaAssetsRefs = await saveKibanaAssetsRefs( - savedObjectsClient, - pkgName, - kibanaAssetsToAssetsRef(kibanaAssets) - ); - } const importedAssets = await installKibanaAssets({ savedObjectsClient, logger, savedObjectsImporter, pkgName, - kibanaAssets, + kibanaAssetsArchiveIterator, }); - if (installAsAdditionalSpace) { - const assets = importedAssets.map( - ({ id, type, destinationId }) => - ({ - id: destinationId ?? id, - originId: id, - type, - } as KibanaAssetReference) - ); - installedKibanaAssetsRefs = await saveKibanaAssetsRefs( - savedObjectsClient, - pkgName, - assets, - installedPkg && installedPkg.attributes.installed_kibana_space_id === spaceId - ? false - : installAsAdditionalSpace - ); - } + const assets = importedAssets.map( + ({ id, type, destinationId }) => + ({ + id: destinationId ?? id, + ...(destinationId ? { originId: id } : {}), + type, + } as KibanaAssetReference) + ); + installedKibanaAssetsRefs = await saveKibanaAssetsRefs( + savedObjectsClient, + pkgName, + assets, + installedPkg && installedPkg.attributes.installed_kibana_space_id === spaceId + ? false + : installAsAdditionalSpace + ); + await withPackageSpan('Create and assign package tags', () => tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle, pkgName, spaceId, @@ -355,25 +355,32 @@ export const isKibanaAssetType = (path: string) => { return parts.service === 'kibana' && (kibanaAssetTypes as string[]).includes(parts.type); }; -export function getKibanaAssets( - packageInstallContext: PackageInstallContext -): Record { - const result = Object.fromEntries( - kibanaAssetTypes.map((type) => [type, []]) - ) as Record; - - packageInstallContext.paths.filter(isKibanaAssetType).forEach((path) => { - const buffer = getAssetFromAssetsMap(packageInstallContext.assetsMap, path); - const asset = JSON.parse(buffer.toString('utf8')); - - const assetType = getPathParts(path).type as KibanaAssetType; - const soType = KibanaSavedObjectTypeMapping[assetType]; - if (asset.type === soType) { - result[assetType].push(asset); - } - }); +function getKibanaAssetsArchiveIterator(packageInstallContext: PackageInstallContext) { + return ( + onEntry: (entry: { + path: string; + asset: ArchiveAsset; + assetType: KibanaAssetType; + }) => Promise + ) => { + return packageInstallContext.archiveIterator.traverseEntries(async (entry) => { + if (!entry.buffer) { + return; + } - return result; + const asset = JSON.parse(entry.buffer.toString('utf8')); + + const assetType = getPathParts(entry.path).type as KibanaAssetType; + const soType = KibanaSavedObjectTypeMapping[assetType]; + if (!validKibanaAssetTypes.has(assetType)) { + return; + } + + if (asset.type === soType) { + await onEntry({ path: entry.path, assetType, asset }); + } + }, isKibanaAssetType); + }; } const isImportConflictError = (e: SavedObjectsImportFailure) => e?.error?.type === 'conflict'; diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.test.ts index 82989715be0e1..e5d646c3b85d6 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.test.ts @@ -42,16 +42,15 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; + const importedAssets = [{ id: 'dashboard1', type: 'dashboard' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], + importedAssets, }); expect(savedObjectTagClient.create).toHaveBeenCalledWith( @@ -72,7 +71,7 @@ describe('tagKibanaAssets', () => { ); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['fleet-managed-default', 'fleet-pkg-system-default'], - assign: kibanaAssets.dashboard, + assign: importedAssets, unassign: [], refresh: false, }); @@ -80,22 +79,21 @@ describe('tagKibanaAssets', () => { it('should only assign Managed and System tags when tags already exist', async () => { savedObjectTagClient.get.mockResolvedValue({ name: '', color: '', description: '' }); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; + const importedAssets = [{ id: 'dashboard1', type: 'dashboard' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], }); expect(savedObjectTagClient.create).not.toHaveBeenCalled(); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['fleet-managed-default', 'fleet-pkg-system-default'], - assign: kibanaAssets.dashboard, + assign: importedAssets, unassign: [], refresh: false, }); @@ -103,16 +101,16 @@ describe('tagKibanaAssets', () => { it('should use destinationId instead of original SO id if imported asset has it', async () => { savedObjectTagClient.get.mockResolvedValue({ name: '', color: '', description: '' }); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [{ id: 'dashboard1', destinationId: 'destination1' } as any], + importedAssets: [ + { id: 'dashboard1', destinationId: 'destination1', type: 'dashboard' } as any, + ], }); expect(savedObjectTagClient.create).not.toHaveBeenCalled(); @@ -129,33 +127,32 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [{ id: 'dashboard1', type: 'dashboard' }], - search: [{ id: 's1', type: 'search' }], - config: [{ id: 'c1', type: 'config' }], - visualization: [{ id: 'v1', type: 'visualization' }], - osquery_pack_asset: [{ id: 'osquery-pack-asset1', type: 'osquery-pack-asset' }], - osquery_saved_query: [{ id: 'osquery_saved_query1', type: 'osquery_saved_query' }], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 's1', type: 'search' }, + { id: 'c1', type: 'config' }, + { id: 'v1', type: 'visualization' }, + { id: 'osquery-pack-asset1', type: 'osquery-pack-asset' }, + { id: 'osquery_saved_query1', type: 'osquery-saved-query' }, + ] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], }); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['fleet-managed-default', 'fleet-pkg-system-default'], assign: [ - ...kibanaAssets.dashboard, - ...kibanaAssets.search, - ...kibanaAssets.visualization, - ...kibanaAssets.osquery_pack_asset, - ...kibanaAssets.osquery_saved_query, + { id: 'dashboard1', type: 'dashboard' }, + { id: 's1', type: 'search' }, + { id: 'v1', type: 'visualization' }, + { id: 'osquery-pack-asset1', type: 'osquery-pack-asset' }, + { id: 'osquery_saved_query1', type: 'osquery-saved-query' }, ], unassign: [], refresh: false, @@ -163,16 +160,15 @@ describe('tagKibanaAssets', () => { }); it('should do nothing if no taggable assets', async () => { - const kibanaAssets = { config: [{ id: 'c1', type: 'config' }] } as any; + const importedAssets = [{ id: 'c1', type: 'config' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], + importedAssets, }); expect(savedObjectTagAssignmentService.updateTagAssignments).not.toHaveBeenCalled(); @@ -188,16 +184,15 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; + const importedAssets = [{ id: 'dashboard1', type: 'dashboard' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], }); expect(savedObjectTagClient.create).not.toHaveBeenCalledWith( @@ -219,7 +214,7 @@ describe('tagKibanaAssets', () => { ); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['managed', 'fleet-pkg-system-default'], - assign: kibanaAssets.dashboard, + assign: importedAssets, unassign: [], refresh: false, }); @@ -235,16 +230,15 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; + const importedAssets = [{ id: 'dashboard1', type: 'dashboard' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], + importedAssets, }); expect(savedObjectTagClient.create).toHaveBeenCalledWith( @@ -266,7 +260,7 @@ describe('tagKibanaAssets', () => { ); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['fleet-managed-default', 'system'], - assign: kibanaAssets.dashboard, + assign: importedAssets, unassign: [], refresh: false, }); @@ -283,22 +277,21 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { dashboard: [{ id: 'dashboard1', type: 'dashboard' }] } as any; + const importedAssets = [{ id: 'dashboard1', type: 'dashboard' }] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'System', pkgName: 'system', spaceId: 'default', - importedAssets: [], }); expect(savedObjectTagClient.create).not.toHaveBeenCalled(); expect(savedObjectTagAssignmentService.updateTagAssignments).toHaveBeenCalledWith({ tags: ['managed', 'system'], - assign: kibanaAssets.dashboard, + assign: importedAssets, unassign: [], refresh: false, }); @@ -309,14 +302,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'Foo', @@ -326,11 +317,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], + importedAssets, assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(3); @@ -412,23 +402,20 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [{ text: 'Bar', asset_ids: ['dashboard1', 'search_id1'] }]; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(3); @@ -510,13 +497,11 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + ] as any; const assetTags = [ { text: 'myCustomTag', @@ -527,11 +512,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(3); @@ -610,14 +594,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'Foo', @@ -633,11 +615,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], + importedAssets, assetTags, }); expect(savedObjectTagClient.create).not.toHaveBeenCalled(); @@ -657,14 +638,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'foo', @@ -674,11 +653,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(2); @@ -702,14 +680,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'Security Solution', @@ -719,11 +695,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledWith( @@ -753,14 +728,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'Security Solution', @@ -770,11 +743,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'my-secondary-space', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledWith(managedTagPayloadArg1, { @@ -815,23 +787,20 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags: [], }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(2); @@ -843,14 +812,12 @@ describe('tagKibanaAssets', () => { savedObjectTagClient.create.mockImplementation(({ name }: { name: string }) => Promise.resolve({ id: name.toLowerCase(), name }) ); - const kibanaAssets = { - dashboard: [ - { id: 'dashboard1', type: 'dashboard' }, - { id: 'dashboard2', type: 'dashboard' }, - { id: 'search_id1', type: 'search' }, - { id: 'search_id2', type: 'search' }, - ], - } as any; + const importedAssets = [ + { id: 'dashboard1', type: 'dashboard' }, + { id: 'dashboard2', type: 'dashboard' }, + { id: 'search_id1', type: 'search' }, + { id: 'search_id2', type: 'search' }, + ] as any; const assetTags = [ { text: 'Foo', @@ -861,11 +828,10 @@ describe('tagKibanaAssets', () => { await tagKibanaAssets({ savedObjectTagAssignmentService, savedObjectTagClient, - kibanaAssets, + importedAssets, pkgTitle: 'TestPackage', pkgName: 'test-pkg', spaceId: 'default', - importedAssets: [], assetTags, }); expect(savedObjectTagClient.create).toHaveBeenCalledTimes(3); diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.ts index 27c825f039e4a..bebc0ba9dc810 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/tag_assets.ts @@ -6,7 +6,7 @@ */ import { v5 as uuidv5 } from 'uuid'; -import { uniqBy } from 'lodash'; +import { omit, uniqBy } from 'lodash'; import pMap from 'p-map'; import type { SavedObjectsImportSuccess } from '@kbn/core-saved-objects-common'; import { taggableTypes } from '@kbn/saved-objects-tagging-plugin/common/constants'; @@ -14,14 +14,11 @@ import type { IAssignmentService } from '@kbn/saved-objects-tagging-plugin/serve import type { ITagsClient } from '@kbn/saved-objects-tagging-plugin/common/types'; import { MAX_CONCURRENT_PACKAGE_ASSETS } from '../../../../constants'; -import type { KibanaAssetType } from '../../../../../common'; + import type { PackageSpecTags } from '../../../../types'; import { appContextService } from '../../../app_context'; -import type { ArchiveAsset } from './install'; -import { KibanaSavedObjectTypeMapping } from './install'; - interface ObjectReference { type: string; id: string; @@ -85,7 +82,6 @@ const getRandomColor = () => { interface TagAssetsParams { savedObjectTagAssignmentService: IAssignmentService; savedObjectTagClient: ITagsClient; - kibanaAssets: Record; pkgTitle: string; pkgName: string; spaceId: string; @@ -93,14 +89,15 @@ interface TagAssetsParams { assetTags?: PackageSpecTags[]; } +const getNewId = (asset: SavedObjectsImportSuccess) => + asset?.destinationId ? asset.destinationId : asset.id; + export async function tagKibanaAssets(opts: TagAssetsParams) { - const { savedObjectTagAssignmentService, kibanaAssets, importedAssets } = opts; + const { savedObjectTagAssignmentService, importedAssets } = opts; - const getNewId = (assetId: string) => - importedAssets.find((imported) => imported.id === assetId)?.destinationId ?? assetId; - const taggableAssets = getTaggableAssets(kibanaAssets).map((asset) => ({ - ...asset, - id: getNewId(asset.id), + const taggableAssets = getTaggableAssets(importedAssets).map((asset) => ({ + ...omit(asset, 'destinationId'), + id: getNewId(asset), })); if (taggableAssets.length > 0) { const [managedTagId, packageTagId] = await Promise.all([ @@ -150,18 +147,8 @@ export async function tagKibanaAssets(opts: TagAssetsParams) { } } -function getTaggableAssets(kibanaAssets: TagAssetsParams['kibanaAssets']) { - return Object.entries(kibanaAssets).flatMap(([assetType, assets]) => { - if (!taggableTypes.includes(KibanaSavedObjectTypeMapping[assetType as KibanaAssetType])) { - return []; - } - - if (!assets.length) { - return []; - } - - return assets; - }); +function getTaggableAssets(importedAssets: SavedObjectsImportSuccess[]) { + return importedAssets.filter((asset) => taggableTypes.includes(asset.type)); } async function ensureManagedTag( @@ -219,7 +206,7 @@ async function ensurePackageTag( // Ensure that asset tags coming from the kibana/tags.yml file are correctly parsed and created async function getPackageSpecTags( - taggableAssets: ArchiveAsset[], + taggableAssets: SavedObjectsImportSuccess[], opts: Pick ): Promise { const { spaceId, savedObjectTagClient, pkgName, assetTags } = opts; @@ -254,7 +241,7 @@ async function getPackageSpecTags( // Get all the assets of types defined in tag.asset_types from taggable kibanaAssets const getAssetTypesObjectReferences = ( assetTypes: string[] | undefined, - taggableAssets: ArchiveAsset[] + taggableAssets: SavedObjectsImportSuccess[] ): ObjectReference[] => { if (!assetTypes || assetTypes.length === 0) return []; @@ -268,7 +255,7 @@ const getAssetTypesObjectReferences = ( // Get the references to ids defined in tag.asset_ids from taggable kibanaAssets const getAssetIdsObjectReferences = ( assetIds: string[] | undefined, - taggableAssets: ArchiveAsset[] + taggableAssets: SavedObjectsImportSuccess[] ): ObjectReference[] => { if (!assetIds || assetIds.length === 0) return []; diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/package_service.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/package_service.ts index af03016524f17..e87c50f3b263c 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/package_service.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/package_service.ts @@ -414,7 +414,6 @@ class PackageClientImpl implements PackageClient { const { installedTransforms } = await installTransforms({ packageInstallContext: { - assetsMap, packageInfo, paths, archiveIterator, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install.ts index a71805be76b18..92fd4bb7a1a08 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install.ts @@ -492,17 +492,16 @@ async function installPackageFromRegistry({ } // get latest package version and requested version in parallel for performance - const [latestPackage, { paths, packageInfo, assetsMap, archiveIterator, verificationResult }] = + const [latestPackage, { paths, packageInfo, archiveIterator, verificationResult }] = await Promise.all([ latestPkg ? Promise.resolve(latestPkg) : queryLatest(), Registry.getPackage(pkgName, pkgVersion, { ignoreUnverified: force && !neverIgnoreVerificationError, - useStreaming, + useStreaming: true, }), ]); const packageInstallContext: PackageInstallContext = { packageInfo, - assetsMap, paths, archiveIterator, }; @@ -831,7 +830,7 @@ async function installPackageByUpload({ packageInfo, }); - const { paths, assetsMap, archiveIterator } = await unpackBufferToAssetsMap({ + const { paths, archiveIterator } = await unpackBufferToAssetsMap({ archiveBuffer, contentType, useStreaming, @@ -839,7 +838,6 @@ async function installPackageByUpload({ const packageInstallContext: PackageInstallContext = { packageInfo: { ...packageInfo, version: pkgVersion }, - assetsMap, paths, archiveIterator, }; @@ -1036,7 +1034,6 @@ export async function installCustomPackage( const archiveIterator = createArchiveIteratorFromMap(assetsMap); const packageInstallContext: PackageInstallContext = { - assetsMap, paths, packageInfo, archiveIterator, @@ -1374,7 +1371,6 @@ export async function installAssetsForInputPackagePolicy(opts: { const archiveIterator = createArchiveIteratorFromMap(pkg.assetsMap); packageInstallContext = { - assetsMap: pkg.assetsMap, packageInfo: pkg.packageInfo, paths: pkg.paths, archiveIterator, @@ -1382,7 +1378,6 @@ export async function installAssetsForInputPackagePolicy(opts: { } else { const archiveIterator = createArchiveIteratorFromMap(installedPkgWithAssets.assetsMap); packageInstallContext = { - assetsMap: installedPkgWithAssets.assetsMap, packageInfo: installedPkgWithAssets.packageInfo, paths: installedPkgWithAssets.paths, archiveIterator, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_index_template_pipeline.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_index_template_pipeline.ts index ad738d1710fcc..b079f8ee86dbf 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_index_template_pipeline.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_index_template_pipeline.ts @@ -55,7 +55,7 @@ export async function installIndexTemplatesAndPipelines({ packageInstallContext, onlyForDataStreams ); - const preparedIndexTemplates = prepareToInstallTemplates( + const preparedIndexTemplates = await prepareToInstallTemplates( packageInstallContext, esReferences, experimentalDataStreamFeatures, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/_state_machine_package_install.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/_state_machine_package_install.test.ts index 73b78a6cc4aa0..dfd6cce6b91bc 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/_state_machine_package_install.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/_state_machine_package_install.test.ts @@ -111,7 +111,6 @@ describe('_stateMachineInstallPackage', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -174,7 +173,6 @@ describe('_stateMachineInstallPackage', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -211,7 +209,6 @@ describe('_stateMachineInstallPackage', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -261,7 +258,6 @@ describe('_stateMachineInstallPackage', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -341,7 +337,6 @@ describe('_stateMachineInstallPackage', () => { conditions: { kibana: { version: 'x.y.z' } }, owner: { github: 'elastic/fleet' }, } as any, - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], }, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_create_restart_installation.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_create_restart_installation.test.ts index e5a7fed55fe87..6cec95cd7c53c 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_create_restart_installation.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_create_restart_installation.test.ts @@ -85,7 +85,6 @@ describe('stepCreateRestartInstallation', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -122,7 +121,6 @@ describe('stepCreateRestartInstallation', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -167,7 +165,6 @@ describe('stepCreateRestartInstallation', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -212,7 +209,6 @@ describe('stepCreateRestartInstallation', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_ilm_policies.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_ilm_policies.test.ts index 4c106a0c68f15..7b78b7967abc7 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_ilm_policies.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_ilm_policies.test.ts @@ -241,7 +241,6 @@ describe('stepInstallILMPolicies', () => { conditions: { kibana: { version: 'x.y.z' } }, owner: { github: 'elastic/fleet' }, } as any, - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], }, diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_kibana_assets.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_kibana_assets.test.ts index cf9d953868b6a..dcf7862fcfa53 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_kibana_assets.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_install_kibana_assets.test.ts @@ -74,7 +74,6 @@ const packageInstallContext = { owner: { github: 'elastic/fleet' }, } as any, paths: ['some/path/1', 'some/path/2'], - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), }; @@ -99,7 +98,6 @@ describe('stepInstallKibanaAssets', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -139,7 +137,6 @@ describe('stepInstallKibanaAssets', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -186,7 +183,6 @@ describe('stepInstallKibanaAssetsWithStreaming', () => { esClient, logger: loggerMock.create(), packageInstallContext: { - assetsMap, archiveIterator, paths: [], packageInfo: { diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts index f081d9a93e633..54221723cb28b 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_archive_entries.ts @@ -6,7 +6,7 @@ */ import { ASSETS_SAVED_OBJECT_TYPE } from '../../../../../constants'; -import type { PackageAssetReference } from '../../../../../types'; +import type { AssetsMap, KibanaAssetType, PackageAssetReference } from '../../../../../types'; import { removeArchiveEntries, saveArchiveEntriesFromAssetsMap } from '../../../archive/storage'; @@ -14,43 +14,52 @@ import { withPackageSpan } from '../../utils'; import type { InstallContext } from '../_state_machine_package_install'; import { INSTALL_STATES } from '../../../../../../common/types'; -import { isKibanaAssetType } from '../../../kibana/assets/install'; +import { getPathParts } from '../../../archive'; export async function stepSaveArchiveEntries(context: InstallContext) { const { packageInstallContext, savedObjectsClient, installSource, useStreaming } = context; const { packageInfo, archiveIterator } = packageInstallContext; - let assetsMap = packageInstallContext?.assetsMap; - let paths = packageInstallContext?.paths; - // For stream based installations, we don't want to save any assets but - // manifest.yaml due to the large number of assets in the package. - if (useStreaming) { - assetsMap = new Map(); - await archiveIterator.traverseEntries(async (entry) => { - // Skip only kibana assets type - if (!isKibanaAssetType(entry.path)) { - assetsMap.set(entry.path, entry.buffer); - } - }); - paths = Array.from(assetsMap.keys()); + let assetsToSaveMap: AssetsMap = new Map(); + + let packageAssetRefs: PackageAssetReference[] = []; + + async function flushAssets() { + const paths = Array.from(assetsToSaveMap.keys()); + const packageAssetResults = await withPackageSpan('Update archive entries', () => + saveArchiveEntriesFromAssetsMap({ + savedObjectsClient, + assetsMap: assetsToSaveMap, + paths, + packageInfo, + installSource, + }) + ); + packageAssetRefs = [ + ...packageAssetRefs, + ...packageAssetResults.saved_objects.map((result) => ({ + id: result.id, + type: ASSETS_SAVED_OBJECT_TYPE as typeof ASSETS_SAVED_OBJECT_TYPE, + })), + ]; + + assetsToSaveMap = new Map(); } - const packageAssetResults = await withPackageSpan('Update archive entries', () => - saveArchiveEntriesFromAssetsMap({ - savedObjectsClient, - assetsMap, - paths, - packageInfo, - installSource, - }) - ); - const packageAssetRefs: PackageAssetReference[] = packageAssetResults.saved_objects.map( - (result) => ({ - id: result.id, - type: ASSETS_SAVED_OBJECT_TYPE, - }) - ); + await archiveIterator.traverseEntries(async (entry) => { + const assetType = getPathParts(entry.path).type as KibanaAssetType; + if (assetType === 'security_rule' && useStreaming) { + // Skip security rules to avoid storing to many things + } else { + assetsToSaveMap.set(entry.path, entry.buffer); + } + if (assetsToSaveMap.size > 100) { + await flushAssets(); + } + }); + + await flushAssets(); return { packageAssetRefs }; } diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_system_object.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_system_object.test.ts index c4ae211c58fc5..62414ad0e8089 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_system_object.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/step_save_system_object.test.ts @@ -66,7 +66,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -133,7 +132,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { diff --git a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/update_latest_executed_state.test.ts b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/update_latest_executed_state.test.ts index aea879aba5479..02ace384e760b 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/update_latest_executed_state.test.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/epm/packages/install_state_machine/steps/update_latest_executed_state.test.ts @@ -62,7 +62,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -118,7 +117,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -156,7 +154,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { @@ -202,7 +199,6 @@ describe('updateLatestExecutedState', () => { esClient, logger, packageInstallContext: { - assetsMap: new Map(), archiveIterator: createArchiveIteratorFromMap(new Map()), paths: [], packageInfo: { diff --git a/x-pack/platform/plugins/shared/fleet/server/services/package_policies/experimental_datastream_features.ts b/x-pack/platform/plugins/shared/fleet/server/services/package_policies/experimental_datastream_features.ts index af8352b639d75..5e9630f625230 100644 --- a/x-pack/platform/plugins/shared/fleet/server/services/package_policies/experimental_datastream_features.ts +++ b/x-pack/platform/plugins/shared/fleet/server/services/package_policies/experimental_datastream_features.ts @@ -75,11 +75,11 @@ export async function handleExperimentalDatastreamFeatureOptIn({ ); return prepareTemplate({ packageInstallContext: { - assetsMap, archiveIterator: createArchiveIteratorFromMap(assetsMap), packageInfo, paths, }, + fieldAssetsMap: assetsMap, dataStream, experimentalDataStreamFeature, }); diff --git a/x-pack/test/fleet_api_integration/apis/epm/__snapshots__/bulk_get_assets.snap b/x-pack/test/fleet_api_integration/apis/epm/__snapshots__/bulk_get_assets.snap index 37c7882fff203..3ff5a36d9147b 100644 --- a/x-pack/test/fleet_api_integration/apis/epm/__snapshots__/bulk_get_assets.snap +++ b/x-pack/test/fleet_api_integration/apis/epm/__snapshots__/bulk_get_assets.snap @@ -98,6 +98,12 @@ Array [ "id": "metrics-all_assets.test_metrics@custom", "type": "component_template", }, + Object { + "appLink": "", + "attributes": Object {}, + "id": "sample_csp_rule_template", + "type": "csp-rule-template", + }, Object { "appLink": "/app/dashboards#/view/sample_dashboard", "attributes": Object { @@ -117,39 +123,21 @@ Array [ "type": "dashboard", }, Object { - "appLink": "/app/lens#/edit/sample_lens", - "attributes": Object { - "description": "", - "title": "sample-lens", - }, - "id": "sample_lens", - "type": "lens", - }, - Object { - "appLink": "/app/visualize#/edit/sample_visualization", + "appLink": "/app/management/kibana/dataViews/dataView/test-*", "attributes": Object { - "description": "sample visualization update", - "title": "sample vis title", + "title": "test-*", }, - "id": "sample_visualization", - "type": "visualization", + "id": "test-*", + "type": "index-pattern", }, Object { - "appLink": "/app/discover#/view/sample_search", + "appLink": "/app/lens#/edit/sample_lens", "attributes": Object { "description": "", - "title": "All logs [Logs Kafka] ECS", - }, - "id": "sample_search", - "type": "search", - }, - Object { - "appLink": "/app/management/kibana/dataViews/dataView/test-*", - "attributes": Object { - "title": "test-*", + "title": "sample-lens", }, - "id": "test-*", - "type": "index-pattern", + "id": "sample_lens", + "type": "lens", }, Object { "appLink": "/app/ml/supplied_configurations/?_a=(supplied_configurations%3A(queryText%3A'Nginx%20access%20logs'))", @@ -160,20 +148,6 @@ Array [ "id": "sample_ml_module", "type": "ml-module", }, - Object { - "appLink": "", - "attributes": Object { - "description": "Identifies a suspicious parent child process relationship with cmd.exe descending from svchost.exe", - }, - "id": "sample_security_rule", - "type": "security-rule", - }, - Object { - "appLink": "", - "attributes": Object {}, - "id": "sample_csp_rule_template", - "type": "csp-rule-template", - }, Object { "appLink": "", "attributes": Object {}, @@ -188,6 +162,23 @@ Array [ "id": "sample_osquery_saved_query", "type": "osquery-saved-query", }, + Object { + "appLink": "/app/discover#/view/sample_search", + "attributes": Object { + "description": "", + "title": "All logs [Logs Kafka] ECS", + }, + "id": "sample_search", + "type": "search", + }, + Object { + "appLink": "", + "attributes": Object { + "description": "Identifies a suspicious parent child process relationship with cmd.exe descending from svchost.exe", + }, + "id": "sample_security_rule", + "type": "security-rule", + }, Object { "appLink": "", "attributes": Object { @@ -196,5 +187,14 @@ Array [ "id": "sample_tag", "type": "tag", }, + Object { + "appLink": "/app/visualize#/edit/sample_visualization", + "attributes": Object { + "description": "sample visualization update", + "title": "sample vis title", + }, + "id": "sample_visualization", + "type": "visualization", + }, ] `; diff --git a/x-pack/test/fleet_api_integration/apis/epm/update_assets.ts b/x-pack/test/fleet_api_integration/apis/epm/update_assets.ts index 17d54786245af..8611bcd0932e9 100644 --- a/x-pack/test/fleet_api_integration/apis/epm/update_assets.ts +++ b/x-pack/test/fleet_api_integration/apis/epm/update_assets.ts @@ -8,6 +8,7 @@ import expect from '@kbn/expect'; import { FLEET_INSTALL_FORMAT_VERSION } from '@kbn/fleet-plugin/server/constants'; +import { sortBy } from 'lodash'; import { FtrProviderContext } from '../../../api_integration/ftr_provider_context'; import { skipIfNoDockerRegistry } from '../../helpers'; @@ -336,50 +337,56 @@ export default function (providerContext: FtrProviderContext) { id: 'all_assets', }); - expect(res.attributes).eql({ + expect({ + ...res.attributes, + installed_kibana: sortBy(res.attributes.installed_kibana, ['id']), + }).eql({ installed_kibana_space_id: 'default', - installed_kibana: [ - { - id: 'sample_dashboard', - type: 'dashboard', - }, - { - id: 'sample_lens', - type: 'lens', - }, - { - id: 'sample_visualization', - type: 'visualization', - }, - { - id: 'sample_search2', - type: 'search', - }, - { - id: 'sample_ml_module', - type: 'ml-module', - }, - { - id: 'sample_security_rule', - type: 'security-rule', - }, - { - id: 'sample_csp_rule_template2', - type: 'csp-rule-template', - }, - { - id: 'sample_osquery_pack_asset', - type: 'osquery-pack-asset', - }, - { - id: 'sample_osquery_saved_query', - type: 'osquery-saved-query', - }, - { - id: 'sample_tag', - type: 'tag', - }, - ], + installed_kibana: sortBy( + [ + { + id: 'sample_dashboard', + type: 'dashboard', + }, + { + id: 'sample_lens', + type: 'lens', + }, + { + id: 'sample_visualization', + type: 'visualization', + }, + { + id: 'sample_search2', + type: 'search', + }, + { + id: 'sample_ml_module', + type: 'ml-module', + }, + { + id: 'sample_security_rule', + type: 'security-rule', + }, + { + id: 'sample_csp_rule_template2', + type: 'csp-rule-template', + }, + { + id: 'sample_osquery_pack_asset', + type: 'osquery-pack-asset', + }, + { + id: 'sample_osquery_saved_query', + type: 'osquery-saved-query', + }, + { + id: 'sample_tag', + type: 'tag', + }, + ], + 'id' + ), installed_es: [ { id: 'all_assets', From 647a183e3b79cf1544b8aea8dd21d4e1a780b010 Mon Sep 17 00:00:00 2001 From: "Joey F. Poon" Date: Wed, 5 Mar 2025 03:43:28 +0900 Subject: [PATCH 6/9] [Security Solution] defend insights langgraph upgrade (#211038) ## Summary This is intended to be a "minimal" migration for Defend Insights to langgraph + output chunking. Other than the increased events due to the context increase from output chunking, the functionality is unchanged. * migrates defend insights to langgraph * adds output chunking / refinement ### Checklist - [x] [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 --- .github/CODEOWNERS | 2 + .../kbn-elastic-assistant-common/constants.ts | 2 +- .../helpers/get_raw_data_or_default/index.ts | 2 +- .../__mocks__/defend_insights_schema.mock.ts | 2 +- .../server/__mocks__/raw_defend_insights.ts | 108 ++ .../server/__mocks__/request_context.ts | 2 +- .../server/ai_assistant_service/index.ts | 4 +- .../server/lib}/defend_insights/errors.ts | 4 +- .../constants.ts | 20 + .../index.test.ts | 22 + .../get_generate_or_end_decision/index.ts | 9 + .../edges/generate_or_end/index.test.ts | 72 ++ .../edges/generate_or_end/index.ts | 37 + .../index.test.ts | 43 + .../index.ts | 28 + .../helpers/get_should_end/index.test.ts | 60 ++ .../helpers/get_should_end/index.ts | 16 + .../generate_or_refine_or_end/index.test.ts | 120 +++ .../edges/generate_or_refine_or_end/index.ts | 65 ++ .../helpers/get_has_results/index.test.ts | 37 + .../edges/helpers/get_has_results/index.ts | 10 + .../get_refine_or_end_decision/index.test.ts | 62 ++ .../get_refine_or_end_decision/index.ts | 25 + .../helpers/get_should_end/index.test.ts | 60 ++ .../helpers/get_should_end/index.ts | 16 + .../edges/refine_or_end/index.test.ts | 122 +++ .../edges/refine_or_end/index.ts | 61 ++ .../get_retrieve_or_generate/index.test.ts | 19 + .../get_retrieve_or_generate/index.ts | 13 + .../index.test.ts | 70 ++ .../index.ts | 36 + .../index.test.ts | 28 + .../index.ts | 14 + .../get_max_retries_reached/index.test.ts | 22 + .../helpers/get_max_retries_reached/index.ts | 14 + .../default_defend_insights_graph/index.ts | 128 +++ .../mock/mock_anonymization_fields.ts | 911 ++++++++++++++++ .../mock/mock_anonymized_events.ts | 26 + .../mock/mock_defend_insights.ts | 31 + .../mock/mock_file_events_query_results.ts | 984 ++++++++++++++++++ .../index.test.ts | 90 ++ .../discard_previous_generations/index.ts | 30 + .../index.test.ts | 40 + .../get_anonymized_events_from_state/index.ts | 11 + .../get_events_context_prompt/index.test.ts | 37 + .../get_events_context_prompt/index.ts | 22 + .../get_use_unrefined_results/index.test.ts | 51 + .../get_use_unrefined_results/index.ts | 27 + .../nodes/generate/index.test.ts | 366 +++++++ .../nodes/generate/index.ts | 167 +++ .../schema}/incompatible_antivirus.ts | 13 +- .../nodes/generate/schema}/index.ts | 8 +- .../index.test.ts | 46 + .../index.ts | 20 + .../nodes/helpers/extract_json/index.test.ts | 85 ++ .../nodes/helpers/extract_json/index.ts | 21 + .../generations_are_repeating/index.test.ts | 90 ++ .../generations_are_repeating/index.ts | 25 + .../index.test.ts | 45 + .../index.ts | 36 + .../nodes/helpers/get_combined/index.test.ts | 37 + .../nodes/helpers/get_combined/index.ts | 14 + .../helpers/get_combined_prompt/index.test.ts | 62 ++ .../helpers/get_combined_prompt/index.ts | 43 + .../helpers/get_continue_prompt/index.test.ts | 22 + .../helpers/get_continue_prompt/index.ts | 15 + .../helpers/get_output_parser/index.test.ts | 31 + .../nodes/helpers/get_output_parser/index.ts | 16 + .../helpers/parse_combined_or_throw/index.ts | 55 + .../helpers/prompts/incompatible_antivirus.ts | 23 + .../nodes/helpers}/prompts/index.ts | 6 +- .../response_is_hallucinated/index.test.ts | 26 + .../helpers/response_is_hallucinated/index.ts | 9 + .../index.test.ts | 102 ++ .../discard_previous_refinements/index.ts | 30 + .../get_combined_refine_prompt/index.test.ts | 83 ++ .../get_combined_refine_prompt/index.ts | 48 + .../get_default_refine_prompt/index.ts | 11 + .../get_use_unrefined_results/index.test.ts | 46 + .../get_use_unrefined_results/index.ts | 17 + .../nodes/refine/index.ts | 173 +++ .../anonymized_events_retriever/index.test.ts | 101 ++ .../anonymized_events_retriever/index.ts | 90 ++ .../get_events/get_file_events_query.ts | 26 +- .../get_anonymized_events/get_events/index.ts | 29 + .../get_anonymized_events/index.test.ts | 271 +++++ .../helpers/get_anonymized_events}/index.ts | 73 +- .../nodes/retriever/index.test.ts | 118 +++ .../nodes/retriever/index.ts | 71 ++ .../state/index.test.ts | 120 +++ .../state/index.ts | 105 ++ .../default_defend_insights_graph/types.ts | 31 + .../persistence}/field_maps_configuration.ts | 0 .../persistence}/get_defend_insight.test.ts | 3 +- .../persistence}/get_defend_insight.ts | 0 .../persistence}/helpers.test.ts | 0 .../defend_insights/persistence}/helpers.ts | 0 .../persistence}/index.test.ts | 6 +- .../defend_insights/persistence}/index.ts | 8 +- .../defend_insights/persistence}/types.ts | 2 +- .../server/lib/langchain/graphs/index.ts | 24 +- .../elastic_assistant/server/plugin.ts | 5 +- .../get_defend_insight.test.ts | 6 +- .../get_defend_insights.test.ts | 6 +- .../routes/defend_insights/helpers.test.ts | 45 +- .../server/routes/defend_insights/helpers.ts | 249 ++++- .../post_defend_insights.test.ts | 47 +- .../defend_insights/post_defend_insights.ts | 80 +- .../routes/defend_insights/translations.ts | 28 + .../evaluate/get_graphs_from_names/index.ts | 9 +- .../server/routes/helpers.ts | 4 +- .../server/services/app_context.ts | 23 + .../plugins/elastic_assistant/server/types.ts | 10 +- .../defend_insights/get_events/index.test.ts | 183 ---- .../tools/defend_insights/index.test.ts | 69 -- .../assistant/tools/defend_insights/index.ts | 127 --- .../prompts/incompatible_antivirus.ts | 16 - .../server/assistant/tools/index.ts | 2 - .../workflow_insights/builders/index.ts | 4 +- .../services/workflow_insights/index.test.ts | 40 +- .../services/workflow_insights/index.ts | 22 +- .../security_solution/server/plugin.ts | 1 + 122 files changed, 6783 insertions(+), 677 deletions(-) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/raw_defend_insights.ts rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools => elastic_assistant/server/lib}/defend_insights/errors.ts (70%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/constants.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymization_fields.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymized_events.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_defend_insights.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_file_events_query_results.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.ts rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools/defend_insights/output_parsers => elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema}/incompatible_antivirus.ts (77%) rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools/defend_insights/output_parsers => elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema}/index.ts (61%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/parse_combined_or_throw/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/incompatible_antivirus.ts rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools/defend_insights => elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers}/prompts/index.ts (81%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_default_refine_prompt/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.ts rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools/defend_insights => elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events}/get_events/get_file_events_query.ts (72%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.test.ts rename x-pack/solutions/security/plugins/{security_solution/server/assistant/tools/defend_insights/get_events => elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events}/index.ts (56%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.test.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.ts create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/types.ts rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/field_maps_configuration.ts (100%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/get_defend_insight.test.ts (95%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/get_defend_insight.ts (100%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/helpers.test.ts (100%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/helpers.ts (100%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/index.test.ts (98%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/index.ts (97%) rename x-pack/solutions/security/plugins/elastic_assistant/server/{ai_assistant_data_clients/defend_insights => lib/defend_insights/persistence}/types.ts (95%) create mode 100644 x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/translations.ts delete mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts delete mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts delete mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.ts delete mode 100644 x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 20485606602dd..afb07a34d3bc2 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -2507,6 +2507,8 @@ x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/defend x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts @elastic/security-defend-workflows x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights @elastic/security-defend-workflows x-pack/plugins/elastic_assistant/server/routes/defend_insights @elastic/security-defend-workflows +x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights @elastic/security-defend-workflows +x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights @elastic/security-defend-workflows /x-pack/solutions/security/plugins/security_solution/public/common/components/response_actions @elastic/security-defend-workflows /x-pack/solutions/security/plugins/security_solution_serverless/public/upselling/pages/osquery_automated_response_actions.tsx @elastic/security-defend-workflows diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts index 5042b66facd70..aba070040a833 100755 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts @@ -57,6 +57,6 @@ export const ELASTIC_AI_ASSISTANT_EVALUATE_URL = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/evaluate` as const; // Defend insights -export const DEFEND_INSIGHTS_TOOL_ID = 'defend-insights'; +export const DEFEND_INSIGHTS_ID = 'defend-insights'; export const DEFEND_INSIGHTS = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/defend_insights`; export const DEFEND_INSIGHTS_BY_ID = `${DEFEND_INSIGHTS}/{id}`; diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts index edbe320c95305..49db36d51bc9d 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts @@ -5,8 +5,8 @@ * 2.0. */ -import { isRawDataValid } from '../is_raw_data_valid'; import type { MaybeRawData } from '../types'; +import { isRawDataValid } from '../is_raw_data_valid'; /** Returns the raw data if it valid, or a default if it's not */ export const getRawDataOrDefault = (rawData: MaybeRawData): Record => diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts index 474f456d4e8fa..3a4c7d14a45f2 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts @@ -8,7 +8,7 @@ import type { estypes } from '@elastic/elasticsearch'; import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; -import type { EsDefendInsightSchema } from '../ai_assistant_data_clients/defend_insights/types'; +import type { EsDefendInsightSchema } from '../lib/defend_insights/persistence/types'; export const getDefendInsightsSearchEsMock = () => { const searchResponse: estypes.SearchResponse = { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/raw_defend_insights.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/raw_defend_insights.ts new file mode 100644 index 0000000000000..8bfa03bdeb480 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/raw_defend_insights.ts @@ -0,0 +1,108 @@ +/* + * 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 { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { EsDefendInsightSchema } from '../lib/defend_insights/persistence/types'; + +export const getParsedDefendInsightsMock = (timestamp: string): EsDefendInsightSchema[] => [ + { + '@timestamp': timestamp, + created_at: timestamp, + updated_at: timestamp, + last_viewed_at: timestamp, + users: [ + { + id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0', + name: 'elastic', + }, + ], + status: DefendInsightStatus.Enum.succeeded, + api_config: { + action_type_id: '.bedrock', + connector_id: 'ac4e19d1-e2e2-49af-bf4b-59428473101c', + model: 'anthropic.claude-3-5-sonnet-20240620-v1:0', + }, + endpoint_ids: ['6e09ec1c-644c-4148-a02d-be451c35400d'], + insight_type: DefendInsightType.Enum.incompatible_antivirus, + insights: [ + { + group: 'windows_defenders', + events: [], + }, + ], + namespace: 'default', + id: '655c52ec-49ee-4d20-87e5-7edd6d8f84e8', + generation_intervals: [ + { + date: timestamp, + duration_ms: 13113, + }, + ], + average_interval_ms: 13113, + events_context_count: 100, + replacements: [ + { + uuid: '2009c67b-89b8-43d9-b502-2c32f71875a0', + value: 'root', + }, + { + uuid: '9f7f91b6-6853-48b7-bfb8-403f5efb2364', + value: 'joey-dev-default-3539', + }, + ], + }, + { + '@timestamp': timestamp, + created_at: timestamp, + updated_at: timestamp, + last_viewed_at: timestamp, + users: [ + { + id: '00468e82-e37f-4224-80c1-c62e594c74b1', + name: 'ubuntu', + }, + ], + status: DefendInsightStatus.Enum.succeeded, + api_config: { + action_type_id: '.bedrock', + connector_id: 'bc5e19d1-e2e2-49af-bf4b-59428473101d', + model: 'anthropic.claude-3-5-sonnet-20240620-v1:0', + }, + endpoint_ids: ['b557bb12-8206-44b6-b2a5-dbcce5b1e65e'], + insight_type: DefendInsightType.Enum.noisy_process_tree, + insights: [ + { + group: 'linux_security', + events: [], + }, + ], + namespace: 'default', + id: '7a1b52ec-49ee-4d20-87e5-7edd6d8f84e9', + generation_intervals: [ + { + date: timestamp, + duration_ms: 13113, + }, + ], + average_interval_ms: 13113, + events_context_count: 100, + replacements: [ + { + uuid: '3119c67b-89b8-43d9-b502-2c32f71875b1', + value: 'ubuntu', + }, + { + uuid: '8e7f91b6-6853-48b7-bfb8-403f5efb2365', + value: 'ubuntu-dev-default-3540', + }, + ], + }, +]; + +export const getRawDefendInsightsMock = (timestamp: string) => + JSON.stringify(getParsedDefendInsightsMock(timestamp)); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request_context.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request_context.ts index 8866e23a9816b..e93e27a7b8c8a 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request_context.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request_context.ts @@ -29,7 +29,7 @@ import { } from '../ai_assistant_data_clients/knowledge_base'; import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; -import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights'; +import { DefendInsightsDataClient } from '../lib/defend_insights/persistence'; import { authenticatedUser } from './user'; export const createMockClients = () => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts index 554fa9cb84ff8..8200684710cca 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts @@ -21,7 +21,7 @@ import { omit } from 'lodash'; import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status'; import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers'; import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration'; -import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration'; +import { defendInsightsFieldMap } from '../lib/defend_insights/persistence/field_maps_configuration'; import { getDefaultAnonymizationFields } from '../../common/anonymization'; import { AssistantResourceNames, GetElser } from '../types'; import { @@ -49,7 +49,7 @@ import { GetAIAssistantKnowledgeBaseDataClientParams, } from '../ai_assistant_data_clients/knowledge_base'; import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence'; -import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights'; +import { DefendInsightsDataClient } from '../lib/defend_insights/persistence'; import { createGetElserId, ensureProductDocumentationInstalled } from './helpers'; import { hasAIAssistantLicense } from '../routes/helpers'; diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/errors.ts similarity index 70% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/errors.ts index 03633d2ae1eed..701808f1df6d7 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/errors.ts @@ -5,9 +5,7 @@ * 2.0. */ -import { EndpointError } from '../../../../common/endpoint/errors'; - -export class InvalidDefendInsightTypeError extends EndpointError { +export class InvalidDefendInsightTypeError extends Error { constructor() { super('invalid defend insight type'); } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/constants.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/constants.ts new file mode 100644 index 0000000000000..d581c9d65c79e --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/constants.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. + */ + +// LangGraph metadata +export const DEFEND_INSIGHTS_GRAPH_RUN_NAME = 'Defend insights'; + +// Limits +export const DEFAULT_MAX_GENERATION_ATTEMPTS = 10; +export const DEFAULT_MAX_HALLUCINATION_FAILURES = 5; +export const DEFAULT_MAX_REPEATED_GENERATIONS = 3; + +export const NodeType = { + GENERATE_NODE: 'generate', + REFINE_NODE: 'refine', + RETRIEVE_ANONYMIZED_EVENTS_NODE: 'retrieve_anonymized_events', +} as const; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.test.ts new file mode 100644 index 0000000000000..adc95decacf31 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.test.ts @@ -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. + */ + +import { getGenerateOrEndDecision } from '.'; + +describe('getGenerateOrEndDecision', () => { + it('returns "end" when hasZeroEvents is true', () => { + const result = getGenerateOrEndDecision(true); + + expect(result).toEqual('end'); + }); + + it('returns "generate" when hasZeroEvents is false', () => { + const result = getGenerateOrEndDecision(false); + + expect(result).toEqual('generate'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.ts new file mode 100644 index 0000000000000..2205438bd2da1 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/helpers/get_generate_or_end_decision/index.ts @@ -0,0 +1,9 @@ +/* + * 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. + */ + +export const getGenerateOrEndDecision = (hasZeroEvents: boolean): 'end' | 'generate' => + hasZeroEvents ? 'end' : 'generate'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.test.ts new file mode 100644 index 0000000000000..a1f48bc223eb5 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.test.ts @@ -0,0 +1,72 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; + +import { getGenerateOrEndEdge } from '.'; +import type { GraphState } from '../../types'; + +const logger = loggerMock.create(); + +const graphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: [ + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + ], + combinedGenerations: 'generations', + combinedRefinements: 'refinements', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 10, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getGenerateOrEndEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it("returns 'end' when there are zero events", () => { + const state: GraphState = { + ...graphState, + anonymizedEvents: [], // <-- zero events + }; + + const edge = getGenerateOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'generate' when there are events", () => { + const edge = getGenerateOrEndEdge(logger); + const result = edge(graphState); + + expect(result).toEqual('generate'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.ts new file mode 100644 index 0000000000000..d79cc8e766b9f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_end/index.ts @@ -0,0 +1,37 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; + +import type { GraphState } from '../../types'; +import { getGenerateOrEndDecision } from './helpers/get_generate_or_end_decision'; + +export const getGenerateOrEndEdge = (logger?: Logger) => { + const edge = (state: GraphState): 'end' | 'generate' => { + logger?.debug(() => '---GENERATE OR END---'); + const { anonymizedEvents } = state; + + const hasZeroEvents = !anonymizedEvents.length; + + const decision = getGenerateOrEndDecision(hasZeroEvents); + + logger?.debug( + () => `generatOrEndEdge evaluated the following (derived) state:\n${JSON.stringify( + { + anonymizedEvents: anonymizedEvents.length, + hasZeroEvents, + }, + null, + 2 + )} +\n---GENERATE OR END: ${decision}---` + ); + return decision; + }; + + return edge; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.test.ts new file mode 100644 index 0000000000000..8ceb75e5ba55b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.test.ts @@ -0,0 +1,43 @@ +/* + * 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 { getGenerateOrRefineOrEndDecision } from '.'; + +describe('getGenerateOrRefineOrEndDecision', () => { + it("returns 'end' if getShouldEnd returns true", () => { + const result = getGenerateOrRefineOrEndDecision({ + hasUnrefinedResults: false, + hasZeroEvents: true, + maxHallucinationFailuresReached: true, + maxRetriesReached: true, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'refine' if hasUnrefinedResults is true and getShouldEnd returns false", () => { + const result = getGenerateOrRefineOrEndDecision({ + hasUnrefinedResults: true, + hasZeroEvents: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('refine'); + }); + + it("returns 'generate' if hasUnrefinedResults is false and getShouldEnd returns false", () => { + const result = getGenerateOrRefineOrEndDecision({ + hasUnrefinedResults: false, + hasZeroEvents: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('generate'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.ts new file mode 100644 index 0000000000000..07d279315c126 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_generate_or_refine_or_end_decision/index.ts @@ -0,0 +1,28 @@ +/* + * 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 { getShouldEnd } from '../get_should_end'; + +export const getGenerateOrRefineOrEndDecision = ({ + hasUnrefinedResults, + hasZeroEvents, + maxHallucinationFailuresReached, + maxRetriesReached, +}: { + hasUnrefinedResults: boolean; + hasZeroEvents: boolean; + maxHallucinationFailuresReached: boolean; + maxRetriesReached: boolean; +}): 'end' | 'generate' | 'refine' => { + if (getShouldEnd({ hasZeroEvents, maxHallucinationFailuresReached, maxRetriesReached })) { + return 'end'; + } else if (hasUnrefinedResults) { + return 'refine'; + } else { + return 'generate'; + } +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.test.ts new file mode 100644 index 0000000000000..5a659e9db3c98 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.test.ts @@ -0,0 +1,60 @@ +/* + * 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 { getShouldEnd } from '.'; + +describe('getShouldEnd', () => { + it('returns true if hasZeroEvents is true', () => { + const result = getShouldEnd({ + hasZeroEvents: true, // <-- true + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxHallucinationFailuresReached is true', () => { + const result = getShouldEnd({ + hasZeroEvents: false, + maxHallucinationFailuresReached: true, // <-- true + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxRetriesReached is true', () => { + const result = getShouldEnd({ + hasZeroEvents: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, // <-- true + }); + + expect(result).toBe(true); + }); + + it('returns false if all conditions are false', () => { + const result = getShouldEnd({ + hasZeroEvents: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(false); + }); + + it('returns true if all conditions are true', () => { + const result = getShouldEnd({ + hasZeroEvents: true, + maxHallucinationFailuresReached: true, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.ts new file mode 100644 index 0000000000000..e6d6e6f663b98 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/helpers/get_should_end/index.ts @@ -0,0 +1,16 @@ +/* + * 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. + */ + +export const getShouldEnd = ({ + hasZeroEvents, + maxHallucinationFailuresReached, + maxRetriesReached, +}: { + hasZeroEvents: boolean; + maxHallucinationFailuresReached: boolean; + maxRetriesReached: boolean; +}): boolean => hasZeroEvents || maxRetriesReached || maxHallucinationFailuresReached; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.test.ts new file mode 100644 index 0000000000000..d082868ae122b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.test.ts @@ -0,0 +1,120 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; + +import { getGenerateOrRefineOrEndEdge } from '.'; +import type { GraphState } from '../../types'; +import type { DefendInsight } from '@kbn/elastic-assistant-common'; + +const logger = loggerMock.create(); + +const mockDefendInsight: DefendInsight = { + group: 'test-group', + events: [ + { + id: 'test-id', + endpointId: 'test-endpoint', + value: 'test-value', + }, + ], +}; + +const graphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: [ + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,c675d7eb6ee181d788b474117bae8d3ed4bdc2168605c330a93dd342534fb02b\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', + }, + ], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getGenerateOrRefineOrEndEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it('returns "end" when there are zero events', () => { + const withZeroEvents: GraphState = { + ...graphState, + anonymizedEvents: [], // <-- zero events + }; + + const edge = getGenerateOrRefineOrEndEdge(logger); + const result = edge(withZeroEvents); + + expect(result).toEqual('end'); + }); + + it('returns "end" when max hallucination failures are reached', () => { + const withMaxHallucinationFailures: GraphState = { + ...graphState, + hallucinationFailures: 5, + }; + + const edge = getGenerateOrRefineOrEndEdge(logger); + const result = edge(withMaxHallucinationFailures); + + expect(result).toEqual('end'); + }); + + it('returns "end" when max retries are reached', () => { + const withMaxRetries: GraphState = { + ...graphState, + generationAttempts: 10, + }; + + const edge = getGenerateOrRefineOrEndEdge(logger); + const result = edge(withMaxRetries); + + expect(result).toEqual('end'); + }); + + it('returns refine when there are unrefined results', () => { + const withUnrefinedResults: GraphState = { + ...graphState, + unrefinedResults: [mockDefendInsight], + }; + + const edge = getGenerateOrRefineOrEndEdge(logger); + const result = edge(withUnrefinedResults); + + expect(result).toEqual('refine'); + }); + + it('return generate when there are no unrefined results', () => { + const edge = getGenerateOrRefineOrEndEdge(logger); + const result = edge(graphState); + + expect(result).toEqual('generate'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.ts new file mode 100644 index 0000000000000..f0305e5b07c1f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/generate_or_refine_or_end/index.ts @@ -0,0 +1,65 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; + +import type { GraphState } from '../../types'; +import { getGenerateOrRefineOrEndDecision } from './helpers/get_generate_or_refine_or_end_decision'; +import { getHasResults } from '../helpers/get_has_results'; +import { getMaxHallucinationFailuresReached } from '../../helpers/get_max_hallucination_failures_reached'; +import { getMaxRetriesReached } from '../../helpers/get_max_retries_reached'; + +export const getGenerateOrRefineOrEndEdge = (logger?: Logger) => { + const edge = (state: GraphState): 'end' | 'generate' | 'refine' => { + logger?.debug(() => '---GENERATE OR REFINE OR END---'); + const { + anonymizedEvents, + generationAttempts, + hallucinationFailures, + maxGenerationAttempts, + maxHallucinationFailures, + unrefinedResults, + } = state; + + const hasZeroEvents = !anonymizedEvents.length; + const hasUnrefinedResults = getHasResults(unrefinedResults); + const maxRetriesReached = getMaxRetriesReached({ generationAttempts, maxGenerationAttempts }); + const maxHallucinationFailuresReached = getMaxHallucinationFailuresReached({ + hallucinationFailures, + maxHallucinationFailures, + }); + + const decision = getGenerateOrRefineOrEndDecision({ + hasUnrefinedResults, + hasZeroEvents, + maxHallucinationFailuresReached, + maxRetriesReached, + }); + + logger?.debug( + () => + `generatOrRefineOrEndEdge evaluated the following (derived) state:\n${JSON.stringify( + { + anonymizedEvents: anonymizedEvents.length, + generationAttempts, + hallucinationFailures, + hasUnrefinedResults, + hasZeroEvents, + maxHallucinationFailuresReached, + maxRetriesReached, + unrefinedResults: unrefinedResults?.length ?? 0, + }, + null, + 2 + )} + \n---GENERATE OR REFINE OR END: ${decision}---` + ); + return decision; + }; + + return edge; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.test.ts new file mode 100644 index 0000000000000..e73d1200817aa --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.test.ts @@ -0,0 +1,37 @@ +/* + * 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 type { DefendInsight } from '@kbn/elastic-assistant-common'; + +import { getHasResults } from '.'; + +const mockDefendInsight: DefendInsight = { + group: 'test-group', + events: [ + { + id: 'test-id', + endpointId: 'test-endpoint', + value: 'test-value', + }, + ], +}; + +describe('getHasResults', () => { + it('returns true when insights is non-null array with items', () => { + const insights: DefendInsight[] = [mockDefendInsight]; + expect(getHasResults(insights)).toBe(true); + }); + + it('returns true when insights is empty array', () => { + const insights: DefendInsight[] = []; + expect(getHasResults(insights)).toBe(true); + }); + + it('returns false when insights is null', () => { + expect(getHasResults(null)).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.ts new file mode 100644 index 0000000000000..58c0fd567536f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/helpers/get_has_results/index.ts @@ -0,0 +1,10 @@ +/* + * 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 type { DefendInsight } from '@kbn/elastic-assistant-common'; + +export const getHasResults = (insights: DefendInsight[] | null): boolean => insights !== null; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts new file mode 100644 index 0000000000000..3c44114ed0f89 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.test.ts @@ -0,0 +1,62 @@ +/* + * 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 { getRefineOrEndDecision } from '.'; + +describe('getRefineOrEndDecision', () => { + it("returns 'end' when there are final results", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: true, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + describe('limits shared by both the generate and refine steps', () => { + it("returns 'end' when the max hallucinations limit is reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when the max generation attempts limit is reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when multiple limits are reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: true, + }); + + expect(result).toEqual('end'); + }); + }); + + it("returns 'refine' when no final results and no limits reached", () => { + const result = getRefineOrEndDecision({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toEqual('refine'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.ts new file mode 100644 index 0000000000000..7168aa08aeef2 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_refine_or_end_decision/index.ts @@ -0,0 +1,25 @@ +/* + * 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 { getShouldEnd } from '../get_should_end'; + +export const getRefineOrEndDecision = ({ + hasFinalResults, + maxHallucinationFailuresReached, + maxRetriesReached, +}: { + hasFinalResults: boolean; + maxHallucinationFailuresReached: boolean; + maxRetriesReached: boolean; +}): 'refine' | 'end' => + getShouldEnd({ + hasFinalResults, + maxHallucinationFailuresReached, + maxRetriesReached, + }) + ? 'end' + : 'refine'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts new file mode 100644 index 0000000000000..8c35773f8bea2 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.test.ts @@ -0,0 +1,60 @@ +/* + * 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 { getShouldEnd } from '.'; + +describe('getShouldEnd', () => { + it('returns true when hasFinalResults is true', () => { + const result = getShouldEnd({ + hasFinalResults: true, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true when maxHallucinationFailuresReached is true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true when maxRetriesReached is true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns true when both maxHallucinationFailuresReached and maxRetriesReached are true', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: true, // <-- limit reached + maxRetriesReached: true, // <-- another limit reached + }); + + expect(result).toBe(true); + }); + + it('returns false when all conditions are false', () => { + const result = getShouldEnd({ + hasFinalResults: false, + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.ts new file mode 100644 index 0000000000000..697f93dd3a02f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/helpers/get_should_end/index.ts @@ -0,0 +1,16 @@ +/* + * 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. + */ + +export const getShouldEnd = ({ + hasFinalResults, + maxHallucinationFailuresReached, + maxRetriesReached, +}: { + hasFinalResults: boolean; + maxHallucinationFailuresReached: boolean; + maxRetriesReached: boolean; +}): boolean => hasFinalResults || maxRetriesReached || maxHallucinationFailuresReached; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.test.ts new file mode 100644 index 0000000000000..288728344d20b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.test.ts @@ -0,0 +1,122 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; +import type { Document } from '@langchain/core/documents'; +import type { DefendInsight } from '@kbn/elastic-assistant-common'; + +import { getRefineOrEndEdge } from '.'; +import type { GraphState } from '../../types'; + +const logger = loggerMock.create(); + +const mockDocument: Document = { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', +}; + +const mockDefendInsight: DefendInsight = { + group: 'test-group', + events: [ + { + id: 'test-id', + endpointId: 'test-endpoint', + value: 'test-value', + }, + ], +}; + +const initialGraphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: [mockDocument], + combinedGenerations: 'generations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['gen', 'erations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: [mockDefendInsight], +}; + +describe('getRefineOrEndEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it("returns 'end' when there are final results", () => { + const state: GraphState = { + ...initialGraphState, + insights: [mockDefendInsight], // Has final results + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'refine' when there are no final results and no limits reached", () => { + const state: GraphState = { + ...initialGraphState, + insights: null, // No final results + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('refine'); + }); + + it("returns 'end' when max generation attempts limit is reached", () => { + const state: GraphState = { + ...initialGraphState, + insights: null, + generationAttempts: initialGraphState.maxGenerationAttempts, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when max hallucination failures limit is reached", () => { + const state: GraphState = { + ...initialGraphState, + insights: null, + hallucinationFailures: initialGraphState.maxHallucinationFailures, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); + + it("returns 'end' when multiple limits are reached", () => { + const state: GraphState = { + ...initialGraphState, + insights: null, + generationAttempts: initialGraphState.maxGenerationAttempts, + hallucinationFailures: initialGraphState.maxHallucinationFailures, + }; + + const edge = getRefineOrEndEdge(logger); + const result = edge(state); + + expect(result).toEqual('end'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.ts new file mode 100644 index 0000000000000..1ebdd2f2c08c5 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/refine_or_end/index.ts @@ -0,0 +1,61 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; + +import type { GraphState } from '../../types'; +import { getRefineOrEndDecision } from './helpers/get_refine_or_end_decision'; +import { getHasResults } from '../helpers/get_has_results'; +import { getMaxHallucinationFailuresReached } from '../../helpers/get_max_hallucination_failures_reached'; +import { getMaxRetriesReached } from '../../helpers/get_max_retries_reached'; + +export const getRefineOrEndEdge = (logger?: Logger) => { + const edge = (state: GraphState): 'end' | 'refine' => { + logger?.debug(() => '---REFINE OR END---'); + const { + insights, + generationAttempts, + hallucinationFailures, + maxGenerationAttempts, + maxHallucinationFailures, + } = state; + + const hasFinalResults = getHasResults(insights); + const maxRetriesReached = getMaxRetriesReached({ generationAttempts, maxGenerationAttempts }); + const maxHallucinationFailuresReached = getMaxHallucinationFailuresReached({ + hallucinationFailures, + maxHallucinationFailures, + }); + + const decision = getRefineOrEndDecision({ + hasFinalResults, + maxHallucinationFailuresReached, + maxRetriesReached, + }); + + logger?.debug( + () => + `refineOrEndEdge evaluated the following (derived) state:\n${JSON.stringify( + { + insights: insights?.length ?? 0, + generationAttempts, + hallucinationFailures, + hasFinalResults, + maxHallucinationFailuresReached, + maxRetriesReached, + }, + null, + 2 + )} + \n---REFINE OR END: ${decision}---` + ); + + return decision; + }; + + return edge; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.test.ts new file mode 100644 index 0000000000000..24e5201d12c3c --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.test.ts @@ -0,0 +1,19 @@ +/* + * 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 { mockAnonymizedEvents } from '../../../mock/mock_anonymized_events'; +import { getRetrieveOrGenerate } from '.'; + +describe('getRetrieveOrGenerate', () => { + it("returns 'retrieve_anonymized_events' when anonymizedEvents is empty", () => { + expect(getRetrieveOrGenerate([])).toBe('retrieve_anonymized_events'); + }); + + it("returns 'generate' when anonymizedEvents is not empty", () => { + expect(getRetrieveOrGenerate(mockAnonymizedEvents)).toBe('generate'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.ts new file mode 100644 index 0000000000000..8d6bff518aefd --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/get_retrieve_or_generate/index.ts @@ -0,0 +1,13 @@ +/* + * 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 type { Document } from '@langchain/core/documents'; + +export const getRetrieveOrGenerate = ( + anonymizedEvents: Document[] +): 'retrieve_anonymized_events' | 'generate' => + anonymizedEvents.length === 0 ? 'retrieve_anonymized_events' : 'generate'; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.test.ts new file mode 100644 index 0000000000000..6087d1c2a3ffc --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.test.ts @@ -0,0 +1,70 @@ +/* + * 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 { loggerMock } from '@kbn/logging-mocks'; +import type { Document } from '@langchain/core/documents'; + +import { getRetrieveAnonymizedEventsOrGenerateEdge } from '.'; +import type { GraphState } from '../../types'; + +const logger = loggerMock.create(); + +const mockDocument: Document = { + metadata: {}, + pageContent: + '@timestamp,2024-10-10T21:01:24.148Z\n' + + '_id,e809ffc5e0c2e731c1f146e0f74250078136a87574534bf8e9ee55445894f7fc\n' + + 'host.name,e1cb3cf0-30f3-4f99-a9c8-518b955c6f90\n' + + 'user.name,039c15c5-3964-43e7-a891-42fe2ceeb9ff', +}; + +const initialGraphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: [], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getRetrieveAnonymizedEventsOrGenerateEdge', () => { + beforeEach(() => jest.clearAllMocks()); + + it('returns "generate" when anonymizedEvents is NOT empty', () => { + const state: GraphState = { + ...initialGraphState, + anonymizedEvents: [mockDocument], + }; + + const edge = getRetrieveAnonymizedEventsOrGenerateEdge(logger); + const result = edge(state); + + expect(result).toEqual('generate'); + }); + + it('returns "retrieve_anonymized_events" when anonymizedEvents is empty', () => { + const state: GraphState = { + ...initialGraphState, + anonymizedEvents: [], // <-- empty + }; + + const edge = getRetrieveAnonymizedEventsOrGenerateEdge(logger); + const result = edge(state); + + expect(result).toEqual('retrieve_anonymized_events'); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.ts new file mode 100644 index 0000000000000..d6ef9a92203af --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/edges/retrieve_anonymized_events_or_generate/index.ts @@ -0,0 +1,36 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; + +import type { GraphState } from '../../types'; +import { getRetrieveOrGenerate } from './get_retrieve_or_generate'; + +export const getRetrieveAnonymizedEventsOrGenerateEdge = (logger?: Logger) => { + const edge = (state: GraphState): 'retrieve_anonymized_events' | 'generate' => { + logger?.debug(() => '---RETRIEVE ANONYMIZED EVENTS OR GENERATE---'); + const { anonymizedEvents } = state; + + const decision = getRetrieveOrGenerate(anonymizedEvents); + + logger?.debug( + () => + `retrieveAnonymizedEventsOrGenerateEdge evaluated the following (derived) state:\n${JSON.stringify( + { + anonymizedEvents: anonymizedEvents.length, + }, + null, + 2 + )} + \n---RETRIEVE ANONYMIZED EVENTS OR GENERATE: ${decision}---` + ); + + return decision; + }; + + return edge; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.test.ts new file mode 100644 index 0000000000000..138179109708e --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.test.ts @@ -0,0 +1,28 @@ +/* + * 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 { getMaxHallucinationFailuresReached } from '.'; + +describe('getMaxHallucinationFailuresReached', () => { + it('return true when hallucination failures is equal to the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 2, maxHallucinationFailures: 2 }) + ).toBe(true); + }); + + it('returns true when hallucination failures is greater than the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 3, maxHallucinationFailures: 2 }) + ).toBe(true); + }); + + it('returns false when hallucination failures is less than the max hallucination failures', () => { + expect( + getMaxHallucinationFailuresReached({ hallucinationFailures: 1, maxHallucinationFailures: 2 }) + ).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.ts new file mode 100644 index 0000000000000..07985381afa73 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_hallucination_failures_reached/index.ts @@ -0,0 +1,14 @@ +/* + * 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. + */ + +export const getMaxHallucinationFailuresReached = ({ + hallucinationFailures, + maxHallucinationFailures, +}: { + hallucinationFailures: number; + maxHallucinationFailures: number; +}): boolean => hallucinationFailures >= maxHallucinationFailures; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.test.ts new file mode 100644 index 0000000000000..47f49a75415c9 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.test.ts @@ -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. + */ + +import { getMaxRetriesReached } from '.'; + +describe('getMaxRetriesReached', () => { + it('returns true when generation attempts is equal to the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 2, maxGenerationAttempts: 2 })).toBe(true); + }); + + it('returns true when generation attempts is greater than the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 3, maxGenerationAttempts: 2 })).toBe(true); + }); + + it('returns false when generation attempts is less than the max generation attempts', () => { + expect(getMaxRetriesReached({ generationAttempts: 1, maxGenerationAttempts: 2 })).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.ts new file mode 100644 index 0000000000000..c1e36917b45cf --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/helpers/get_max_retries_reached/index.ts @@ -0,0 +1,14 @@ +/* + * 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. + */ + +export const getMaxRetriesReached = ({ + generationAttempts, + maxGenerationAttempts, +}: { + generationAttempts: number; + maxGenerationAttempts: number; +}): boolean => generationAttempts >= maxGenerationAttempts; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/index.ts new file mode 100644 index 0000000000000..7ce8c5701d212 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/index.ts @@ -0,0 +1,128 @@ +/* + * 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 type { CompiledStateGraph } from '@langchain/langgraph'; +import type { ElasticsearchClient, Logger } from '@kbn/core/server'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; +import { END, START, StateGraph } from '@langchain/langgraph'; +import { DefendInsightType, Replacements } from '@kbn/elastic-assistant-common'; +import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +import type { GraphState } from './types'; +import { getRetrieveAnonymizedEventsOrGenerateEdge } from './edges/retrieve_anonymized_events_or_generate'; +import { getGenerateNode } from './nodes/generate'; +import { getGenerateOrEndEdge } from './edges/generate_or_end'; +import { getGenerateOrRefineOrEndEdge } from './edges/generate_or_refine_or_end'; +import { getRefineNode } from './nodes/refine'; +import { getRefineOrEndEdge } from './edges/refine_or_end'; +import { getRetrieveAnonymizedEventsNode } from './nodes/retriever'; +import { getDefaultGraphState } from './state'; +import { NodeType } from './constants'; + +export interface GetDefaultDefendInsightsGraphParams { + insightType: DefendInsightType; + endpointIds: string[]; + anonymizationFields: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + llm: ActionsClientLlm; + logger?: Logger; + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; + size?: number; + start?: string; + end?: string; +} + +export type DefaultDefendInsightsGraph = ReturnType; + +/** + * This function returns a compiled state graph that represents the default + * Defend Insights graph. + */ +export const getDefaultDefendInsightsGraph = ({ + insightType, + endpointIds, + anonymizationFields, + esClient, + llm, + logger, + onNewReplacements, + replacements, + size, + start, + end, +}: GetDefaultDefendInsightsGraphParams): CompiledStateGraph< + GraphState, + Partial, + 'generate' | 'refine' | 'retrieve_anonymized_events' | '__start__' +> => { + try { + const graphState = getDefaultGraphState({ insightType, start, end }); + + // get nodes: + const retrieveAnonymizedEventsNode = getRetrieveAnonymizedEventsNode({ + insightType, + endpointIds, + anonymizationFields, + esClient, + logger, + onNewReplacements, + replacements, + size, + }); + + const generateNode = getGenerateNode({ + insightType, + llm, + logger, + }); + + const refineNode = getRefineNode({ + insightType, + llm, + logger, + }); + + // get edges: + const generateOrEndEdge = getGenerateOrEndEdge(logger); + + const generatOrRefineOrEndEdge = getGenerateOrRefineOrEndEdge(logger); + + const refineOrEndEdge = getRefineOrEndEdge(logger); + + const retrieveAnonymizedEventsOrGenerateEdge = + getRetrieveAnonymizedEventsOrGenerateEdge(logger); + + // create the graph: + const graph = new StateGraph({ channels: graphState }) + .addNode(NodeType.RETRIEVE_ANONYMIZED_EVENTS_NODE, retrieveAnonymizedEventsNode) + .addNode(NodeType.GENERATE_NODE, generateNode) + .addNode(NodeType.REFINE_NODE, refineNode) + .addConditionalEdges(START, retrieveAnonymizedEventsOrGenerateEdge, { + generate: NodeType.GENERATE_NODE, + retrieve_anonymized_events: NodeType.RETRIEVE_ANONYMIZED_EVENTS_NODE, + }) + .addConditionalEdges(NodeType.RETRIEVE_ANONYMIZED_EVENTS_NODE, generateOrEndEdge, { + end: END, + generate: NodeType.GENERATE_NODE, + }) + .addConditionalEdges(NodeType.GENERATE_NODE, generatOrRefineOrEndEdge, { + end: END, + generate: NodeType.GENERATE_NODE, + refine: NodeType.REFINE_NODE, + }) + .addConditionalEdges(NodeType.REFINE_NODE, refineOrEndEdge, { + end: END, + refine: NodeType.REFINE_NODE, + }); + + // compile the graph: + return graph.compile(); + } catch (e) { + throw new Error(`Unable to compile DefendInsightsGraph\n${e}`); + } +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymization_fields.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymization_fields.ts new file mode 100644 index 0000000000000..541fb78b17c5d --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymization_fields.ts @@ -0,0 +1,911 @@ +/* + * 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 type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +export const mockAnonymizationFields: AnonymizationFieldResponse[] = [ + { + id: '9f95b649-f20e-4edf-bd76-1d21ab6f8e2e', + timestamp: '2024-05-06T22:16:48.489Z', + field: '_id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '13aae62e-e8d1-42a5-b369-38406de9de27', + timestamp: '2024-05-06T22:16:48.489Z', + field: '@timestamp', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'ecf8f8f0-955a-4fd1-b11a-e997c3f70c60', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'cloud.availability_zone', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5ea31dc8-a43f-4b79-8cb7-5eddef99e52e', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'cloud.provider', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '85f18a84-ea74-47ac-89e1-a25d78122229', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'cloud.region', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '0059af85-f6de-4500-aca7-196aa5e9b4e8', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'destination.ip', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '55d5f500-cd79-4809-ac40-507756f2188b', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'dns.question.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '71d4b104-277d-4c83-bb8a-26833cbcb620', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'dns.question.type', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '706b6fe4-0834-4d37-b9da-351a17683a80', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'event.category', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '6b8f7793-77e8-4ef6-bd92-573eafc71385', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'event.dataset', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '0851d53d-da3e-4c5a-8b9e-ed9c8fdf990b', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'event.module', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '373d9cf5-77b4-4dea-a2b6-82ea90bcf1a7', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'event.outcome', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e50e603c-e8f2-43cb-86a5-2b01cde43fa9', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'file.Ext.original.path', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '63b56dc8-c395-4ad6-b92e-f2a5d16de84d', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'file.hash.sha256', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '8e2ca725-a8a4-4e18-8ed7-8e7815400217', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'file.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '9d5ead5e-7929-4923-a37a-c763127e189f', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'file.path', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '8d5c4c4d-6af6-4784-8046-20531e633bab', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'group.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5f89fd96-f133-4b93-b54c-7bb91f9bbeb0', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'group.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '199e4a4e-9224-4e20-a397-63fec38ecb0b', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.asset.criticality', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '22f23471-4f6a-4cec-9b2a-cf270ffb53d5', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.name', + allowed: true, + anonymized: true, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '671f9bd2-d2c8-4637-baee-6a64cfebd518', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.os.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '660f9d45-9050-405d-adce-2dd597abb2ea', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.os.version', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cb971f55-0f80-407a-9924-56751301e884', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.risk.calculated_level', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '0a8a078c-8812-4f2e-83c5-804bff42519a', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'host.risk.calculated_score_norm', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '34cf8706-366e-49e6-81a8-6d27175b1776', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.original_time', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '10f1bcf5-c7f2-4570-a799-7d865be761d6', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.risk_score', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e06f542c-f9a0-4916-a426-142775cdfb6e', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.description', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5bb110d1-b652-4c92-b7ec-9f0bf32532af', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '4166edab-6f9a-4de0-9a03-4be8f41c1902', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.references', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '39d8cdc3-4e30-4a5c-81f9-fc0172b7571f', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.framework', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '9aa16534-a81b-41e6-808e-73b771774630', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.tactic.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'fbd77d15-0fe0-4601-b84a-c44341102d3f', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.tactic.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'ba18a341-1603-4cc0-adfe-3c2507a71b72', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.tactic.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '7ccbee93-0b1b-4d3b-8706-f4fd0f6695f1', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'b4b56b1f-ea5b-4f57-ab28-038d65dc1153', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '06e18204-4b84-43bb-b20e-f2a109afafd8', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '0cb32f1a-532f-40d1-8ca5-92debeaee618', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.subtechnique.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '53f75a04-5535-4f0c-8591-379a7391c636', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.subtechnique.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'ccbd069a-c640-497f-9de6-6d8ad5a1f55d', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.rule.threat.technique.subtechnique.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cefb5f12-6de2-403c-bc0a-518ab25ae354', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.severity', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'a8c7e1cd-5837-4254-b151-a1fca47aba20', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'kibana.alert.workflow_status', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e3fb3a83-da33-4ffb-b111-992b8ac070f7', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'message', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5177f9a4-672d-4aaf-afcd-cbc8003de954', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'network.protocol', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '01c5b596-8164-49cf-ae7d-df14822b0cfd', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.args', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'aa460567-f3f2-460e-8917-7272c1d01bcb', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.code_signature.exists', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '566bd237-5b9d-4ecd-b9fb-8d11c50dbe7c', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.code_signature.signing_id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '2e572530-943c-4800-9150-ae80eb0751fc', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.code_signature.status', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '08a36abe-f2f6-4fab-9ac8-3194f84ffc8f', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.code_signature.subject_name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'c6697286-21e5-41a1-8572-4fd85a8893d5', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.code_signature.trusted', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '48ae58e0-196b-4fe5-907f-30322f9d4bfb', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.command_line', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'f69a4487-770f-4c8b-8c01-b0c8e44e5d8c', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.executable', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '7eb80322-eacc-499f-a0e5-4e2eedadd669', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.exit_code', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'c4ca5831-820c-474e-bfcb-f9bb887f7f2c', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.Ext.memory_region.bytes_compressed_present', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'c84ede38-f079-4ce5-a1de-33ce676e05a0', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.Ext.memory_region.malware_signature.all_names', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '7208f10d-7d02-4edd-9e03-26252b8ea17d', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.Ext.memory_region.malware_signature.primary.matches', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '62669a22-b51c-4f8b-95af-62ac81e092d8', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.Ext.memory_region.malware_signature.primary.signature.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '35c3fb06-5224-4616-82a8-20a56440e3a9', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.Ext.token.integrity_level_name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '74c88612-0565-4015-8b41-866b78077e04', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.hash.md5', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '4e7d5890-0d28-4768-8f84-d2317b3d846a', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.hash.sha1', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '419a59ab-bf92-45fa-a627-87662b22c624', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.hash.sha256', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '36d85f0b-aada-4df8-979a-5fb7fbd94d90', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '12e00c31-df09-4136-ad40-663c32f8fff4', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.args', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'dbbe7fba-efda-4e0e-8fdb-754d6793ab1a', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.args_count', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e4b2b245-0505-44a3-90dd-d508098a17b7', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.code_signature.exists', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cc808e37-cbe7-44c2-bf25-127b7f7a19b1', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.code_signature.status', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '69002247-abdf-4f08-b5cd-d29527b69eaf', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.code_signature.subject_name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'fce068a6-1110-47d7-9f85-4677aed52d67', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.code_signature.trusted', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '1085b328-9c7d-44a1-b0b1-637d37a48cac', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.command_line', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'dae6a0cd-e07e-42da-8302-b03c626e6ca1', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.executable', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '4acfb8c2-5e98-468f-b4b6-1be365a911e4', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.parent.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'b8bfba60-46f3-4997-9b8a-5914623bd9df', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.pe.original_file_name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '37e14145-addc-4873-9d54-e5fda145eedf', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.pid', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '10691388-3b28-49b3-9f8d-fab62ecd054a', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'process.working_directory', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '46273311-2ada-46df-945b-e227d2412301', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.feature', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '0d2d0b0b-2aa5-4ec1-80f8-1eb178884c41', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.data', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5fa6395f-e4db-44f4-a7a3-d41ce765b4ce', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.entropy', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '5838f1ca-e077-49a8-af9d-f315e0c11012', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.extension', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cf69671a-a1e9-4824-a666-d480e74f8ada', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.metrics', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '2558d47b-a9fe-46d0-9997-6c1eb63604e4', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.operation', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'a9555faf-8870-4dd2-9580-f34fb8ee4e31', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.path', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '360cd306-e37c-4fec-8491-f8d4fed20d1b', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.files.score', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'ea2d8983-59ee-46a6-9959-bb2fb3f6de77', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'Ransomware.version', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'bfcf60dc-2aa8-4ab0-a8fe-00b03c3ad804', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'rule.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cc26349e-d9d0-44b2-92b9-dad614ab16fa', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'rule.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '867fe8ec-5e9c-4956-b183-86034b7381ee', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'source.ip', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'd6a3024d-aa82-4202-a783-0edab7cec7cf', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.framework', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e5cdbc2a-0c7d-401b-ade1-ac5469f7c3b1', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.tactic.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '1252bbb4-3cc3-4e44-944b-750ad32e2425', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.tactic.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '28111355-baca-4bb2-a75a-c3f67c211a24', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.tactic.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'e16b78b8-87e3-488a-b4df-881f799003b2', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'fb444fcb-453d-42a0-a3d7-185e9f5e0b97', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '6d972ebc-015a-43a0-bab0-ca35d8780a88', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'a609aa63-bbb5-4ef8-ae4a-61c6eb958663', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.subtechnique.id', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'edf1957a-675a-4c47-8237-d6f460683858', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.subtechnique.name', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'a729bd68-f9a2-409a-88d9-e9b704ab5c1f', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'threat.technique.subtechnique.reference', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: 'cfb37a11-0d68-49ec-95b0-d095a97703fe', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'user.asset.criticality', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '3147ffe8-a7ae-480e-b8b2-ad6e4dd4c949', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'user.domain', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '72b2e896-61b9-46c7-b1bc-66eec416a6e0', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'user.name', + allowed: true, + anonymized: true, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '35f6df7f-b6f0-4971-a0c2-110411dbb9db', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'user.risk.calculated_level', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, + { + id: '3db78c6d-4d42-4aa1-bc61-44dea34a1615', + timestamp: '2024-05-06T22:16:48.489Z', + field: 'user.risk.calculated_score_norm', + allowed: true, + anonymized: false, + createdAt: '2024-05-06T22:16:48.489Z', + namespace: 'default', + }, +]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymized_events.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymized_events.ts new file mode 100644 index 0000000000000..54a126d35151d --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_anonymized_events.ts @@ -0,0 +1,26 @@ +/* + * 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 type { Document } from '@langchain/core/documents'; + +export const mockAnonymizedEvents: Document[] = [ + { + pageContent: + '_id,87c42d26897490ee02ba42ec4e872910b29f3c69bda357b8faf197b533b8528a\nagent.id,f5b69281-3e7e-4b52-9225-e5c30dc29c78\nprocess.executable,C:\\Windows\\System32\\foobar.exe', + metadata: {}, + }, + { + pageContent: + '_id,be6d293f9a71ba209adbcacc3ba04adfd8e9456260f6af342b7cb0478a7a144a\nagent.id,f5b69281-3e7e-4b52-9225-e5c30dc29c78\nprocess.executable,C:\\Program Files\\Some Antivirus\\foobar.exe', + metadata: {}, + }, +]; + +export const mockAnonymizedEventsReplacements: Record = { + '42c4e419-c859-47a5-b1cb-f069d48fa509': 'Administrator', + 'f5b69281-3e7e-4b52-9225-e5c30dc29c78': 'SRVWIN07', +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_defend_insights.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_defend_insights.ts new file mode 100644 index 0000000000000..d7860179c107b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_defend_insights.ts @@ -0,0 +1,31 @@ +/* + * 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 type { DefendInsight } from '@kbn/elastic-assistant-common'; + +export const mockDefendInsights: DefendInsight[] = [ + { + group: 'Windows Defender', + events: [ + { + id: '123', + endpointId: 'endpoint-1', + value: 'some/file/path.exe', + }, + ], + }, + { + group: 'AVG Antivirus', + events: [ + { + id: '456', + endpointId: 'endpoint-2', + value: 'another/file/path.exe', + }, + ], + }, +]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_file_events_query_results.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_file_events_query_results.ts new file mode 100644 index 0000000000000..03b06a61ac47e --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/mock/mock_file_events_query_results.ts @@ -0,0 +1,984 @@ +/* + * 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. + */ + +export const mockFileEventsQueryResults = { + took: 10, + timed_out: false, + _shards: { + total: 5, + successful: 5, + skipped: 0, + failed: 0, + }, + hits: { + total: { + value: 10, + relation: 'gte', + }, + max_score: 1, + hits: [ + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: '1WgClZIBomhdMPLEGc-s', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + name: 'systemd', + pid: 1, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + executable: '/usr/lib/systemd/systemd', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T11:06:12.1806296Z', + file: { + Ext: { + original: { + path: '/run/systemd/units/.#invocation:gce-workload-cert-refresh.service13e32c882e5449c0', + extension: 'service13e32c882e5449c0', + name: '.#invocation:gce-workload-cert-refresh.service13e32c882e5449c0', + }, + }, + path: '/run/systemd/units/invocation:gce-workload-cert-refresh.service', + extension: 'service', + name: 'invocation:gce-workload-cert-refresh.service', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524785, + ingested: '2024-10-16T11:06:35Z', + created: '2024-10-16T11:06:12.1806296Z', + kind: 'event', + module: 'endpoint', + action: 'rename', + id: 'NiUEE6P7Pv00Sp1S++++9q6D', + category: ['file'], + type: ['change'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: '2GgClZIBomhdMPLEGc-s', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T11:06:12.2007223Z', + file: { + path: '/run/systemd/journal/streams/.#8:5203253Z9Q47q', + name: '.#8:5203253Z9Q47q', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524790, + ingested: '2024-10-16T11:06:35Z', + created: '2024-10-16T11:06:12.2007223Z', + kind: 'event', + module: 'endpoint', + action: 'creation', + id: 'NiUEE6P7Pv00Sp1S++++9q6L', + category: ['file'], + type: ['creation'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: '2WgClZIBomhdMPLEGc-s', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T11:06:12.20079Z', + file: { + Ext: { + original: { + path: '/run/systemd/journal/streams/.#8:5203253Z9Q47q', + name: '.#8:5203253Z9Q47q', + }, + }, + path: '/run/systemd/journal/streams/8:5203253', + name: '8:5203253', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524791, + ingested: '2024-10-16T11:06:35Z', + created: '2024-10-16T11:06:12.20079Z', + kind: 'event', + module: 'endpoint', + action: 'rename', + id: 'NiUEE6P7Pv00Sp1S++++9q6M', + category: ['file'], + type: ['change'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: '3WgClZIBomhdMPLEGc-s', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + name: 'systemd', + pid: 1, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + executable: '/usr/lib/systemd/systemd', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T11:06:12.2180429Z', + file: { + path: '/run/systemd/units/invocation:gce-workload-cert-refresh.service', + extension: 'service', + name: 'invocation:gce-workload-cert-refresh.service', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524798, + ingested: '2024-10-16T11:06:35Z', + created: '2024-10-16T11:06:12.2180429Z', + kind: 'event', + module: 'endpoint', + action: 'deletion', + id: 'NiUEE6P7Pv00Sp1S++++9q6W', + category: ['file'], + type: ['deletion'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: '3mgClZIBomhdMPLEGc-s', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T11:06:12.2195128Z', + file: { + path: '/run/systemd/journal/streams/8:5203253', + name: '8:5203253', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524799, + ingested: '2024-10-16T11:06:35Z', + created: '2024-10-16T11:06:12.2195128Z', + kind: 'event', + module: 'endpoint', + action: 'deletion', + id: 'NiUEE6P7Pv00Sp1S++++9q6Y', + category: ['file'], + type: ['deletion'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: 'jGj4lJIBomhdMPLEbM2r', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + name: 'systemd', + pid: 1, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + executable: '/usr/lib/systemd/systemd', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T10:55:54.5240803Z', + file: { + Ext: { + original: { + path: '/run/systemd/units/.#invocation:gce-workload-cert-refresh.servicee6f43c6db49deb4c', + extension: 'servicee6f43c6db49deb4c', + name: '.#invocation:gce-workload-cert-refresh.servicee6f43c6db49deb4c', + }, + }, + path: '/run/systemd/units/invocation:gce-workload-cert-refresh.service', + extension: 'service', + name: 'invocation:gce-workload-cert-refresh.service', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524637, + ingested: '2024-10-16T10:56:01Z', + created: '2024-10-16T10:55:54.5240803Z', + kind: 'event', + module: 'endpoint', + action: 'rename', + id: 'NiUEE6P7Pv00Sp1S++++9q35', + category: ['file'], + type: ['change'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: 'j2j4lJIBomhdMPLEbM2r', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T10:55:54.5526557Z', + file: { + path: '/run/systemd/journal/streams/.#8:5203091bcpcW7', + name: '.#8:5203091bcpcW7', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524642, + ingested: '2024-10-16T10:56:01Z', + created: '2024-10-16T10:55:54.5526557Z', + kind: 'event', + module: 'endpoint', + action: 'creation', + id: 'NiUEE6P7Pv00Sp1S++++9q3D', + category: ['file'], + type: ['creation'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: 'kGj4lJIBomhdMPLEbM2r', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T10:55:54.5527354Z', + file: { + Ext: { + original: { + path: '/run/systemd/journal/streams/.#8:5203091bcpcW7', + name: '.#8:5203091bcpcW7', + }, + }, + path: '/run/systemd/journal/streams/8:5203091', + name: '8:5203091', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524643, + ingested: '2024-10-16T10:56:01Z', + created: '2024-10-16T10:55:54.5527354Z', + kind: 'event', + module: 'endpoint', + action: 'rename', + id: 'NiUEE6P7Pv00Sp1S++++9q3E', + category: ['file'], + type: ['change'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: 'lGj4lJIBomhdMPLEbM2r', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + name: 'systemd', + pid: 1, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + executable: '/usr/lib/systemd/systemd', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T10:55:54.5694977Z', + file: { + path: '/run/systemd/units/invocation:gce-workload-cert-refresh.service', + extension: 'service', + name: 'invocation:gce-workload-cert-refresh.service', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524650, + ingested: '2024-10-16T10:56:01Z', + created: '2024-10-16T10:55:54.5694977Z', + kind: 'event', + module: 'endpoint', + action: 'deletion', + id: 'NiUEE6P7Pv00Sp1S++++9q3N', + category: ['file'], + type: ['deletion'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + { + _index: '.ds-logs-endpoint.events.file-default-2024.09.16-000001', + _id: 'lWj4lJIBomhdMPLEbM2r', + _score: 1, + _source: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + type: 'endpoint', + version: '8.15.1', + }, + process: { + parent: { + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTEtMTcyNjQ4NDM0NQ==', + }, + name: 'systemd-journald', + pid: 217, + entity_id: 'ZjA1MzFmYTYtZjI1Yy00NmQ5LTgxNWEtOTMzM2IyNWJhNzI4LTIxNy0xNzI2NDg0MzQ4', + executable: '/usr/lib/systemd/systemd-journald', + }, + message: 'Endpoint file event', + '@timestamp': '2024-10-16T10:55:54.5710352Z', + file: { + path: '/run/systemd/journal/streams/8:5203091', + name: '8:5203091', + }, + ecs: { + version: '8.10.0', + }, + data_stream: { + namespace: 'default', + type: 'logs', + dataset: 'endpoint.events.file', + }, + elastic: { + agent: { + id: 'f0531fa6-f25c-46d9-815a-9333b25ba728', + }, + }, + host: { + hostname: 'endpoint-complete-deb-gd-1', + os: { + Ext: { + variant: 'Debian', + }, + kernel: '6.1.0-25-cloud-amd64 #1 SMP PREEMPT_DYNAMIC Debian 6.1.106-3 (2024-08-26)', + name: 'Linux', + family: 'debian', + type: 'linux', + version: '12.7', + platform: 'debian', + full: 'Debian 12.7', + }, + ip: ['127.0.0.1', '::1', '10.142.0.11', 'fe80::4001:aff:fe8e:b'], + name: 'endpoint-complete-deb-gd-1', + id: '18a77774a51b4024813da8493f10c056', + mac: ['42-01-0a-8e-00-0b'], + architecture: 'x86_64', + }, + event: { + agent_id_status: 'verified', + sequence: 2524651, + ingested: '2024-10-16T10:56:01Z', + created: '2024-10-16T10:55:54.5710352Z', + kind: 'event', + module: 'endpoint', + action: 'deletion', + id: 'NiUEE6P7Pv00Sp1S++++9q3P', + category: ['file'], + type: ['deletion'], + dataset: 'endpoint.events.file', + outcome: 'unknown', + }, + user: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + group: { + Ext: { + real: { + name: 'root', + id: 0, + }, + }, + name: 'root', + id: 0, + }, + }, + }, + ], + }, +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts new file mode 100644 index 0000000000000..2c60956811601 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.test.ts @@ -0,0 +1,90 @@ +/* + * 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 { discardPreviousGenerations } from '.'; +import { GraphState } from '../../../../types'; + +const graphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: [ + { + metadata: {}, + pageContent: `_id,event-id1 +agent.id,agent-id1 +process.executable,some/file/path.exe`, + }, + { + metadata: {}, + pageContent: `_id,event-id2 +agent.id,agent-id2 +process.executable,another/file/path.exe`, + }, + ], + combinedGenerations: 'combinedGenerations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['combined', 'Generations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('discardPreviousGenerations', () => { + describe('common state updates', () => { + let result: GraphState; + + beforeEach(() => { + result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: false, + state: graphState, + }); + }); + + it('resets the combined generations', () => { + expect(result.combinedGenerations).toBe(''); + }); + + it('increments the generation attempts', () => { + expect(result.generationAttempts).toBe(graphState.generationAttempts + 1); + }); + + it('resets the collection of generations', () => { + expect(result.generations).toEqual([]); + }); + }); + + it('increments hallucinationFailures when a hallucination is detected', () => { + const result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: true, // <-- hallucination detected + state: graphState, + }); + + expect(result.hallucinationFailures).toBe(graphState.hallucinationFailures + 1); + }); + + it('does NOT increment hallucinationFailures when a hallucination is NOT detected', () => { + const result = discardPreviousGenerations({ + generationAttempts: graphState.generationAttempts, + hallucinationFailures: graphState.hallucinationFailures, + isHallucinationDetected: false, // <-- no hallucination detected + state: graphState, + }); + + expect(result.hallucinationFailures).toBe(graphState.hallucinationFailures); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.ts new file mode 100644 index 0000000000000..a40dde44f8d67 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/discard_previous_generations/index.ts @@ -0,0 +1,30 @@ +/* + * 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 { GraphState } from '../../../../types'; + +export const discardPreviousGenerations = ({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected, + state, +}: { + generationAttempts: number; + hallucinationFailures: number; + isHallucinationDetected: boolean; + state: GraphState; +}): GraphState => { + return { + ...state, + combinedGenerations: '', // <-- reset the combined generations + generationAttempts: generationAttempts + 1, + generations: [], // <-- reset the generations + hallucinationFailures: isHallucinationDetected + ? hallucinationFailures + 1 + : hallucinationFailures, + }; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.test.ts new file mode 100644 index 0000000000000..ae9f4a7622931 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.test.ts @@ -0,0 +1,40 @@ +/* + * 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 type { GraphState } from '../../../../types'; +import { mockAnonymizedEvents } from '../../../../mock/mock_anonymized_events'; +import { getAnonymizedEventsFromState } from '.'; + +const graphState: GraphState = { + insights: null, + prompt: 'prompt', + anonymizedEvents: mockAnonymizedEvents, + combinedGenerations: 'combinedGenerations', + combinedRefinements: '', + errors: [], + generationAttempts: 2, + generations: ['combined', 'Generations'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: null, +}; + +describe('getAnonymizedEventsFromState', () => { + it('returns the anonymized events from the state', () => { + const result = getAnonymizedEventsFromState(graphState); + + expect(result).toEqual([ + mockAnonymizedEvents[0].pageContent, + mockAnonymizedEvents[1].pageContent, + ]); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.ts new file mode 100644 index 0000000000000..9b8aaeb55f1d3 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_anonymized_events_from_state/index.ts @@ -0,0 +1,11 @@ +/* + * 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 { GraphState } from '../../../../types'; + +export const getAnonymizedEventsFromState = (state: GraphState): string[] => + state.anonymizedEvents.map((doc) => doc.pageContent); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.test.ts new file mode 100644 index 0000000000000..c85f1844c9420 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.test.ts @@ -0,0 +1,37 @@ +/* + * 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 { getDefendInsightsPrompt } from '../../../helpers/prompts'; +import { getEventsContextPrompt } from '.'; + +const insightType = 'incompatible_antivirus'; + +describe('getEventsContextPrompt', () => { + it('generates the correct prompt', () => { + const anonymizedEvents = ['event 1', 'event 2', 'event 3']; + + const expected = `${getDefendInsightsPrompt({ type: insightType })} + +Use context from the following events to provide insights: + +""" +event 1 + +event 2 + +event 3 +""" +`; + + const prompt = getEventsContextPrompt({ + anonymizedEvents, + prompt: getDefendInsightsPrompt({ type: insightType }), + }); + + expect(prompt).toEqual(expected); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.ts new file mode 100644 index 0000000000000..9967c40b71149 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_events_context_prompt/index.ts @@ -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. + */ + +// NOTE: we ask the LLM to `provide insights`. We do NOT use the feature name, `DefendInsights`, in the prompt. +export const getEventsContextPrompt = ({ + anonymizedEvents, + prompt, +}: { + anonymizedEvents: string[]; + prompt: string; +}) => `${prompt} + +Use context from the following events to provide insights: + +""" +${anonymizedEvents.join('\n\n')} +""" +`; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts new file mode 100644 index 0000000000000..02629f78f3b29 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.test.ts @@ -0,0 +1,51 @@ +/* + * 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 { mockDefendInsights } from '../../../../mock/mock_defend_insights'; +import { getUseUnrefinedResults } from '.'; + +describe('getUseUnrefinedResults', () => { + it('returns true when the next attempt would exceed the limit, and we have unrefined results', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: mockDefendInsights, + }) + ).toBe(true); + }); + + it('returns false when the next attempt would NOT exceed the limit', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 1, + maxGenerationAttempts: 3, + unrefinedResults: mockDefendInsights, + }) + ).toBe(false); + }); + + it('returns false when unrefined results is null', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: null, + }) + ).toBe(false); + }); + + it('returns false when unrefined results is empty', () => { + expect( + getUseUnrefinedResults({ + generationAttempts: 2, + maxGenerationAttempts: 3, + unrefinedResults: [], + }) + ).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.ts new file mode 100644 index 0000000000000..2e18439be6e53 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/helpers/get_use_unrefined_results/index.ts @@ -0,0 +1,27 @@ +/* + * 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 type { DefendInsight } from '@kbn/elastic-assistant-common'; + +import { getMaxRetriesReached } from '../../../../helpers/get_max_retries_reached'; + +export const getUseUnrefinedResults = ({ + generationAttempts, + maxGenerationAttempts, + unrefinedResults, +}: { + generationAttempts: number; + maxGenerationAttempts: number; + unrefinedResults: DefendInsight[] | null; +}): boolean => { + const nextAttemptWouldExcedLimit = getMaxRetriesReached({ + generationAttempts: generationAttempts + 1, // + 1, because we just used an attempt + maxGenerationAttempts, + }); + + return nextAttemptWouldExcedLimit && unrefinedResults != null && unrefinedResults.length > 0; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.test.ts new file mode 100644 index 0000000000000..cc49a1cc59ff7 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.test.ts @@ -0,0 +1,366 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; +import { FakeLLM } from '@langchain/core/utils/testing'; + +import type { GraphState } from '../../types'; +import { + mockAnonymizedEvents, + mockAnonymizedEventsReplacements, +} from '../../mock/mock_anonymized_events'; +import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions'; +import { getDefaultRefinePrompt } from '../refine/helpers/get_default_refine_prompt'; +import { getAnonymizedEventsFromState } from './helpers/get_anonymized_events_from_state'; +import { getGenerateNode } from '.'; + +const insightTimestamp = new Date().toISOString(); + +jest.mock('../helpers/get_chain_with_format_instructions', () => { + const mockInvoke = jest.fn().mockResolvedValue(''); + + return { + getChainWithFormatInstructions: jest.fn().mockReturnValue({ + chain: { + invoke: mockInvoke, + }, + formatInstructions: ['mock format instructions'], + llmType: 'openai', + mockInvoke, + }), + }; +}); + +const mockLogger = { + debug: (x: Function) => x(), +} as unknown as Logger; + +let mockLlm: ActionsClientLlm; + +const initialGraphState: GraphState = { + prompt: 'test prompt', + anonymizedEvents: [...mockAnonymizedEvents], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + insights: null, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: getDefaultRefinePrompt(), + replacements: mockAnonymizedEventsReplacements, + unrefinedResults: null, +}; + +describe('getGenerateNode', () => { + beforeEach(() => { + jest.clearAllMocks(); + + jest.useFakeTimers(); + jest.setSystemTime(new Date(insightTimestamp)); + + mockLlm = new FakeLLM({ + response: '', + }) as unknown as ActionsClientLlm; + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('returns a function', () => { + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlm, + logger: mockLogger, + }); + + expect(typeof generateNode).toBe('function'); + }); + + it('invokes the chain with the expected events from state and formatting instructions', async () => { + const mockInvoke = getChainWithFormatInstructions('incompatible_antivirus', mockLlm).chain + .invoke as jest.Mock; + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlm, + logger: mockLogger, + }); + + await generateNode(initialGraphState); + + expect(mockInvoke).toHaveBeenCalledWith({ + format_instructions: ['mock format instructions'], + query: `${initialGraphState.prompt} + +Use context from the following events to provide insights: + +""" +${getAnonymizedEventsFromState(initialGraphState).join('\n\n')} +""" +`, + }); + }); + + it('removes the surrounding json from the response', async () => { + const response = + 'You asked for some JSON, here it is:\n```json\n{"key": "value"}\n```\nI hope that works for you.'; + + const mockLlmWithResponse = new FakeLLM({ response }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions('incompatible_antivirus', mockLlmWithResponse) + .chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const state = await generateNode(initialGraphState); + + expect(state).toEqual({ + ...initialGraphState, + combinedGenerations: '{"key": "value"}', + errors: [ + 'generate node is unable to parse (fake) response from attempt 0; (this may be an incomplete response from the model): [\n {\n "code": "invalid_type",\n "expected": "array",\n "received": "undefined",\n "path": [\n "insights"\n ],\n "message": "Required"\n }\n]', + ], + generationAttempts: 1, + generations: ['{"key": "value"}'], + }); + }); + + it('handles hallucinations', async () => { + const hallucinatedResponse = + 'tactics like **Credential Access**, **Command and Control**, and **Persistence**.",\n "entitySummaryMarkdown": "Events detected on host **{{ host.name hostNameValue }}**'; + + const mockLlmWithHallucination = new FakeLLM({ + response: hallucinatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions( + 'incompatible_antivirus', + mockLlmWithHallucination + ).chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(hallucinatedResponse); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithHallucination, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: '{"key": "value"}', + generationAttempts: 1, + generations: ['{"key": "value"}'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: '', // <-- reset + generationAttempts: 2, // <-- incremented + generations: [], // <-- reset + hallucinationFailures: 1, // <-- incremented + }); + }); + + it('discards previous generations and starts over when the maxRepeatedGenerations limit is reached', async () => { + const repeatedResponse = 'gen1'; + + const mockLlmWithRepeatedGenerations = new FakeLLM({ + response: repeatedResponse, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions( + 'incompatible_antivirus', + mockLlmWithRepeatedGenerations + ).chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(repeatedResponse); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithRepeatedGenerations, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: 'gen1gen1', + generationAttempts: 2, + generations: ['gen1', 'gen1'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: '', + generationAttempts: 3, // <-- incremented + generations: [], + }); + }); + + it('combines the response with the previous generations', async () => { + const response = 'gen1'; + + const mockLlmWithResponse = new FakeLLM({ + response, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions('incompatible_antivirus', mockLlmWithResponse) + .chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(response); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: 'gen0', + generationAttempts: 1, + generations: ['gen0'], + }; + + const state = await generateNode(withPreviousGenerations); + + expect(state).toEqual({ + ...withPreviousGenerations, + combinedGenerations: 'gen0gen1', + errors: [ + 'generate node is unable to parse (fake) response from attempt 1; (this may be an incomplete response from the model): SyntaxError: Unexpected token \'g\', "gen0gen1" is not valid JSON', + ], + generationAttempts: 2, + generations: ['gen0', 'gen1'], + }); + }); + + it('returns unrefined results when combined responses pass validation', async () => { + const rawInsights = JSON.stringify({ + '@timestamp': insightTimestamp, + insight_type: 'incompatible_antivirus', + insights: [ + { + group: 'test_group', + events: [], + }, + ], + }); + + const mockLlmWithResponse = new FakeLLM({ + response: rawInsights, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions('incompatible_antivirus', mockLlmWithResponse) + .chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(rawInsights); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: '', + generationAttempts: 0, + generations: [], + }; + + const state = await generateNode(withPreviousGenerations); + const expectedResults = [ + { + group: 'test_group', + events: [], + timestamp: insightTimestamp, + }, + ]; + + expect(state).toEqual({ + ...withPreviousGenerations, + insights: null, + combinedGenerations: rawInsights, + errors: [], + generationAttempts: 1, + generations: [rawInsights], + unrefinedResults: expectedResults, + hallucinationFailures: 0, + }); + }); + + it('skips the refinements step if the max number of retries has already been reached', async () => { + const rawInsights = JSON.stringify({ + '@timestamp': insightTimestamp, + insight_type: 'incompatible_antivirus', + insights: [ + { + group: 'test_group', + events: [], + }, + ], + }); + + const mockLlmWithResponse = new FakeLLM({ + response: rawInsights, + }) as unknown as ActionsClientLlm; + const mockInvoke = getChainWithFormatInstructions('incompatible_antivirus', mockLlmWithResponse) + .chain.invoke as jest.Mock; + + mockInvoke.mockResolvedValue(rawInsights); + + const generateNode = getGenerateNode({ + insightType: 'incompatible_antivirus', + llm: mockLlmWithResponse, + logger: mockLogger, + }); + + const withPreviousGenerations = { + ...initialGraphState, + combinedGenerations: '', + generationAttempts: 9, // One away from max + generations: [], + hallucinationFailures: 0, + insights: null, + unrefinedResults: null, + }; + + const state = await generateNode(withPreviousGenerations); + + const expectedResults = [ + { + group: 'test_group', + events: [], + timestamp: insightTimestamp, + }, + ]; + + expect(state).toEqual({ + ...withPreviousGenerations, + insights: expectedResults, + combinedGenerations: rawInsights, + errors: [], + generationAttempts: 10, + generations: [rawInsights], + unrefinedResults: expectedResults, + hallucinationFailures: 0, + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.ts new file mode 100644 index 0000000000000..4aa5413877e30 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/index.ts @@ -0,0 +1,167 @@ +/* + * 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 type { ActionsClientLlm } from '@kbn/langchain/server'; +import type { Logger } from '@kbn/core/server'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { GraphState } from '../../types'; +import { discardPreviousGenerations } from './helpers/discard_previous_generations'; +import { extractJson } from '../helpers/extract_json'; +import { getAnonymizedEventsFromState } from './helpers/get_anonymized_events_from_state'; +import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions'; +import { getCombined } from '../helpers/get_combined'; +import { getCombinedDefendInsightsPrompt } from '../helpers/get_combined_prompt'; +import { generationsAreRepeating } from '../helpers/generations_are_repeating'; +import { getUseUnrefinedResults } from './helpers/get_use_unrefined_results'; +import { parseCombinedOrThrow } from '../helpers/parse_combined_or_throw'; +import { responseIsHallucinated } from '../helpers/response_is_hallucinated'; + +export const getGenerateNode = ({ + insightType, + llm, + logger, +}: { + insightType: DefendInsightType; + llm: ActionsClientLlm; + logger?: Logger; +}): ((state: GraphState) => Promise) => { + const generate = async (state: GraphState): Promise => { + logger?.debug(() => `---GENERATE---`); + + const anonymizedEvents: string[] = getAnonymizedEventsFromState(state); + + const { + prompt, + combinedGenerations, + generationAttempts, + generations, + hallucinationFailures, + maxGenerationAttempts, + maxRepeatedGenerations, + } = state; + + let combinedResponse = ''; // mutable, because it must be accessed in the catch block + let partialResponse = ''; // mutable, because it must be accessed in the catch block + + try { + const query = getCombinedDefendInsightsPrompt({ + anonymizedEvents, + prompt, + combinedMaybePartialResults: combinedGenerations, + }); + + const { chain, formatInstructions, llmType } = getChainWithFormatInstructions( + insightType, + llm + ); + + logger?.debug( + () => `generate node is invoking the chain (${llmType}), attempt ${generationAttempts}` + ); + + const rawResponse = await chain.invoke({ + format_instructions: formatInstructions, + query, + }); + + // LOCAL MUTATION: + partialResponse = extractJson(rawResponse); // remove the surrounding ```json``` + + // if the response is hallucinated, discard previous generations and start over: + if (responseIsHallucinated(partialResponse)) { + logger?.debug( + () => + `generate node detected a hallucination (${llmType}), on attempt ${generationAttempts}; discarding the accumulated generations and starting over` + ); + + return discardPreviousGenerations({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected: true, + state, + }); + } + + // if the generations are repeating, discard previous generations and start over: + if ( + generationsAreRepeating({ + currentGeneration: partialResponse, + previousGenerations: generations, + sampleLastNGenerations: maxRepeatedGenerations - 1, + }) + ) { + logger?.debug( + () => + `generate node detected (${llmType}), detected ${maxRepeatedGenerations} repeated generations on attempt ${generationAttempts}; discarding the accumulated results and starting over` + ); + + // discard the accumulated results and start over: + return discardPreviousGenerations({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected: false, + state, + }); + } + + // LOCAL MUTATION: + combinedResponse = getCombined({ combinedGenerations, partialResponse }); // combine the new response with the previous ones + + const parsedResults = parseCombinedOrThrow({ + insightType, + combinedResponse, + generationAttempts, + llmType, + logger, + nodeName: 'generate', + }); + + // use the unrefined results if we already reached the max number of retries: + // Ensure each insight has required fields + const timestampedResults = parsedResults.map((result) => ({ + ...result, + timestamp: new Date().toISOString(), + })); + + const useUnrefinedResults = getUseUnrefinedResults({ + generationAttempts, + maxGenerationAttempts, + unrefinedResults: timestampedResults, + }); + + if (useUnrefinedResults) { + logger?.debug( + () => + `generate node is using unrefined results response (${llm._llmType()}) from attempt ${generationAttempts}, because all attempts have been used` + ); + } + + return { + ...state, + insights: useUnrefinedResults ? timestampedResults : null, + combinedGenerations: combinedResponse, + generationAttempts: generationAttempts + 1, + generations: [...generations, partialResponse], + unrefinedResults: timestampedResults, + }; + } catch (error) { + const parsingError = `generate node is unable to parse (${llm._llmType()}) response from attempt ${generationAttempts}; (this may be an incomplete response from the model): ${error}`; + logger?.debug(() => parsingError); // logged at debug level because the error is expected when the model returns an incomplete response + + return { + ...state, + combinedGenerations: combinedResponse, + errors: [...state.errors, parsingError], + generationAttempts: generationAttempts + 1, + generations: [...generations, partialResponse], + }; + } + }; + + return generate; +}; diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/incompatible_antivirus.ts similarity index 77% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/incompatible_antivirus.ts index b6430e4408355..8e21a70da0f83 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/incompatible_antivirus.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/incompatible_antivirus.ts @@ -5,13 +5,14 @@ * 2.0. */ -import { StructuredOutputParser } from 'langchain/output_parsers'; +// import { StructuredOutputParser } from 'langchain/output_parsers'; +// StructuredOutputParser.fromZodSchema import { z } from '@kbn/zod'; -export function getIncompatibleVirusOutputParser() { - return StructuredOutputParser.fromZodSchema( - z.array( +export function getIncompatibleVirusSchema() { + return z.object({ + insights: z.array( z.object({ group: z.string().describe('The program which is triggering the events'), events: z @@ -23,6 +24,6 @@ export function getIncompatibleVirusOutputParser() { .array() .describe('The events that the insight is based on'), }) - ) - ); + ), + }); } diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/index.ts similarity index 61% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/index.ts index 78933b72702bf..6c51e6c813a31 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/output_parsers/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/generate/schema/index.ts @@ -7,12 +7,12 @@ import { DefendInsightType } from '@kbn/elastic-assistant-common'; -import { InvalidDefendInsightTypeError } from '../errors'; -import { getIncompatibleVirusOutputParser } from './incompatible_antivirus'; +import { InvalidDefendInsightTypeError } from '../../../../../errors'; +import { getIncompatibleVirusSchema } from './incompatible_antivirus'; -export function getDefendInsightsOutputParser({ type }: { type: DefendInsightType }) { +export function getSchema({ type }: { type: DefendInsightType }) { if (type === DefendInsightType.Enum.incompatible_antivirus) { - return getIncompatibleVirusOutputParser(); + return getIncompatibleVirusSchema(); } throw new InvalidDefendInsightTypeError(); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts new file mode 100644 index 0000000000000..4c95cb05faae0 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.test.ts @@ -0,0 +1,46 @@ +/* + * 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 { addTrailingBackticksIfNecessary } from '.'; + +describe('addTrailingBackticksIfNecessary', () => { + it('adds trailing backticks when necessary', () => { + const input = '```json\n{\n "key": "value"\n}'; + const expected = '```json\n{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(expected); + }); + + it('does NOT add trailing backticks when they are already present', () => { + const input = '```json\n{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it("does NOT add trailing backticks when there's no leading JSON wrapper", () => { + const input = '{\n "key": "value"\n}'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it('handles empty string input', () => { + const input = ''; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); + + it('handles input without a JSON wrapper, but with trailing backticks', () => { + const input = '{\n "key": "value"\n}\n```'; + const result = addTrailingBackticksIfNecessary(input); + + expect(result).toEqual(input); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/index.ts new file mode 100644 index 0000000000000..fd824709f5fcf --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/add_trailing_backticks_if_necessary/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. + */ + +export const addTrailingBackticksIfNecessary = (text: string): string => { + const leadingJSONpattern = /^\w*```json(.*?)/s; + const trailingBackticksPattern = /(.*?)```\w*$/s; + + const hasLeadingJSONWrapper = leadingJSONpattern.test(text); + const hasTrailingBackticks = trailingBackticksPattern.test(text); + + if (hasLeadingJSONWrapper && !hasTrailingBackticks) { + return `${text}\n\`\`\``; + } + + return text; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.test.ts new file mode 100644 index 0000000000000..78230514238e7 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.test.ts @@ -0,0 +1,85 @@ +/* + * 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 { extractJson } from '.'; + +describe('extractJson', () => { + it('returns an empty string if input is undefined', () => { + const input = undefined; + + expect(extractJson(input)).toBe(''); + }); + + it('returns an empty string if input an array', () => { + const input = ['some', 'array']; + + expect(extractJson(input)).toBe(''); + }); + + it('returns an empty string if input is an object', () => { + const input = {}; + + expect(extractJson(input)).toBe(''); + }); + + it('returns the JSON text surrounded by ```json and ``` with no whitespace or additional text', () => { + const input = '```json{"key": "value"}```'; + + const expected = '{"key": "value"}'; + + expect(extractJson(input)).toBe(expected); + }); + + it('returns the JSON block when surrounded by additional text and whitespace', () => { + const input = + 'You asked for some JSON, here it is:\n```json\n{"key": "value"}\n```\nI hope that works for you.'; + + const expected = '{"key": "value"}'; + + expect(extractJson(input)).toBe(expected); + }); + + it('returns the original text if no JSON block is found', () => { + const input = "There's no JSON here, just some text."; + + expect(extractJson(input)).toBe(input); + }); + + it('trims leading and trailing whitespace from the extracted JSON', () => { + const input = 'Text before\n```json\n {"key": "value"} \n```\nText after'; + + const expected = '{"key": "value"}'; + + expect(extractJson(input)).toBe(expected); + }); + + it('handles incomplete JSON blocks with no trailing ```', () => { + const input = 'Text before\n```json\n{"key": "value"'; // <-- no closing ```, because incomplete generation + + expect(extractJson(input)).toBe('{"key": "value"'); + }); + + it('handles multiline defend insight json (real world edge case)', () => { + const input = + '```json\n{\n "@timestamp": "2024-02-08T02:54:40.000Z",\n "created_at": "2024-02-08T02:54:40.000Z",\n "updated_at": "2024-02-08T02:54:40.000Z",\n "last_viewed_at": "2024-02-08T02:54:40.000Z",\n "status": "succeeded",\n "events_context_count": 1,\n "endpoint_ids": ["endpoint-1", "endpoint-2"],\n "insight_type": "incompatible_antivirus",\n "insights": [\n {\n "group": "Windows Defender",\n "events": [\n {\n "id": "event-1",\n "endpoint_id": "endpoint-1",\n "value": "/path/to/executable"\n }\n ]\n }\n ],\n "api_config": {\n "connector_id": "connector-1",\n "action_type_id": "action-1",\n "provider": "openai",\n "model": "gpt-4"\n },\n "replacements": [\n {\n "type": "text",\n "value": "some-replacement"\n }\n ]\n}```'; + + const expected = + '{\n "@timestamp": "2024-02-08T02:54:40.000Z",\n "created_at": "2024-02-08T02:54:40.000Z",\n "updated_at": "2024-02-08T02:54:40.000Z",\n "last_viewed_at": "2024-02-08T02:54:40.000Z",\n "status": "succeeded",\n "events_context_count": 1,\n "endpoint_ids": ["endpoint-1", "endpoint-2"],\n "insight_type": "incompatible_antivirus",\n "insights": [\n {\n "group": "Windows Defender",\n "events": [\n {\n "id": "event-1",\n "endpoint_id": "endpoint-1",\n "value": "/path/to/executable"\n }\n ]\n }\n ],\n "api_config": {\n "connector_id": "connector-1",\n "action_type_id": "action-1",\n "provider": "openai",\n "model": "gpt-4"\n },\n "replacements": [\n {\n "type": "text",\n "value": "some-replacement"\n }\n ]\n}'; + + expect(extractJson(input)).toBe(expected); + }); + + it('handles "Here is my analysis" with defend insight json (real world edge case)', () => { + const input = + 'Here is my analysis in JSON format:\n\n```json\n{\n "@timestamp": "2024-02-08T02:54:40.000Z",\n "created_at": "2024-02-08T02:54:40.000Z",\n "updated_at": "2024-02-08T02:54:40.000Z",\n "last_viewed_at": "2024-02-08T02:54:40.000Z",\n "status": "succeeded",\n "events_context_count": 2,\n "endpoint_ids": ["endpoint-3", "endpoint-4"],\n "insight_type": "incompatible_antivirus",\n "insights": [\n {\n "group": "McAfee",\n "events": [\n {\n "id": "event-2",\n "endpoint_id": "endpoint-3",\n "value": "/usr/local/bin/mcafee"\n },\n {\n "id": "event-3", \n "endpoint_id": "endpoint-4",\n "value": "/usr/local/bin/mcafee"\n }\n ]\n }\n ],\n "api_config": {\n "connector_id": "connector-2",\n "action_type_id": "action-2",\n "provider": "openai",\n "model": "gpt-4"\n },\n "replacements": [\n {\n "type": "text",\n "value": "another-replacement"\n }\n ]\n}```'; + + const expected = + '{\n "@timestamp": "2024-02-08T02:54:40.000Z",\n "created_at": "2024-02-08T02:54:40.000Z",\n "updated_at": "2024-02-08T02:54:40.000Z",\n "last_viewed_at": "2024-02-08T02:54:40.000Z",\n "status": "succeeded",\n "events_context_count": 2,\n "endpoint_ids": ["endpoint-3", "endpoint-4"],\n "insight_type": "incompatible_antivirus",\n "insights": [\n {\n "group": "McAfee",\n "events": [\n {\n "id": "event-2",\n "endpoint_id": "endpoint-3",\n "value": "/usr/local/bin/mcafee"\n },\n {\n "id": "event-3", \n "endpoint_id": "endpoint-4",\n "value": "/usr/local/bin/mcafee"\n }\n ]\n }\n ],\n "api_config": {\n "connector_id": "connector-2",\n "action_type_id": "action-2",\n "provider": "openai",\n "model": "gpt-4"\n },\n "replacements": [\n {\n "type": "text",\n "value": "another-replacement"\n }\n ]\n}'; + + expect(extractJson(input)).toBe(expected); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.ts new file mode 100644 index 0000000000000..089756840e568 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/extract_json/index.ts @@ -0,0 +1,21 @@ +/* + * 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. + */ + +export const extractJson = (input: unknown): string => { + if (typeof input !== 'string') { + return ''; + } + + const regex = /```json\s*([\s\S]*?)(?:\s*```|$)/; + const match = input.match(regex); + + if (match && match[1]) { + return match[1].trim(); + } + + return input; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.test.ts new file mode 100644 index 0000000000000..7d6db4dd72dfd --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.test.ts @@ -0,0 +1,90 @@ +/* + * 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 { generationsAreRepeating } from '.'; + +describe('getIsGenerationRepeating', () => { + it('returns true when all previous generations are the same as the current generation', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: ['gen1', 'gen1', 'gen1'], // <-- all the same, length 3 + sampleLastNGenerations: 3, + }); + + expect(result).toBe(true); + }); + + it('returns false when some of the previous generations are NOT the same as the current generation', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: ['gen1', 'gen2', 'gen1'], // <-- some are different, length 3 + sampleLastNGenerations: 3, + }); + + expect(result).toBe(false); + }); + + it('returns true when all *sampled* generations are the same as the current generation, and there are older samples past the last N', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: [ + 'gen2', // <-- older sample will be ignored + 'gen1', + 'gen1', + 'gen1', + ], + sampleLastNGenerations: 3, + }); + + expect(result).toBe(true); + }); + + it('returns false when some of the *sampled* generations are NOT the same as the current generation, and there are additional samples past the last N', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: [ + 'gen1', // <-- older sample will be ignored + 'gen1', + 'gen1', + 'gen2', + ], + sampleLastNGenerations: 3, + }); + + expect(result).toBe(false); + }); + + it('returns false when sampling fewer generations than sampleLastNGenerations, and all are the same as the current generation', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: ['gen1', 'gen1'], // <-- same, but only 2 generations + sampleLastNGenerations: 3, + }); + + expect(result).toBe(false); + }); + + it('returns false when sampling fewer generations than sampleLastNGenerations, and some are different from the current generation', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: ['gen1', 'gen2'], // <-- different, but only 2 generations + sampleLastNGenerations: 3, + }); + + expect(result).toBe(false); + }); + + it('returns false when there are no previous generations to sample', () => { + const result = generationsAreRepeating({ + currentGeneration: 'gen1', + previousGenerations: [], + sampleLastNGenerations: 3, + }); + + expect(result).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.ts new file mode 100644 index 0000000000000..6cc9cd86c9d2f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/generations_are_repeating/index.ts @@ -0,0 +1,25 @@ +/* + * 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. + */ + +/** Returns true if the last n generations are repeating the same output */ +export const generationsAreRepeating = ({ + currentGeneration, + previousGenerations, + sampleLastNGenerations, +}: { + currentGeneration: string; + previousGenerations: string[]; + sampleLastNGenerations: number; +}): boolean => { + const generationsToSample = previousGenerations.slice(-sampleLastNGenerations); + + if (generationsToSample.length < sampleLastNGenerations) { + return false; // Not enough generations to sample + } + + return generationsToSample.every((generation) => generation === currentGeneration); +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts new file mode 100644 index 0000000000000..8e2ab6715ec45 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.test.ts @@ -0,0 +1,45 @@ +/* + * 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 { FakeLLM } from '@langchain/core/utils/testing'; +import type { ActionsClientLlm } from '@kbn/langchain/server'; + +import { getChainWithFormatInstructions } from '.'; + +describe('getChainWithFormatInstructions', () => { + const mockLlm = new FakeLLM({ + response: JSON.stringify({}, null, 2), + }) as unknown as ActionsClientLlm; + + it('returns the chain with format instructions', () => { + const expectedFormatInstructions = `You must format your output as a JSON value that adheres to a given "JSON Schema" instance. + +"JSON Schema" is a declarative language that allows you to annotate and validate JSON documents. + +For example, the example "JSON Schema" instance {{"properties": {{"foo": {{"description": "a list of test words", "type": "array", "items": {{"type": "string"}}}}}}, "required": ["foo"]}}}} +would match an object with one required property, "foo". The "type" property specifies "foo" must be an "array", and the "description" property semantically describes it as "a list of test words". The items within "foo" must be strings. +Thus, the object {{"foo": ["bar", "baz"]}} is a well-formatted instance of this example "JSON Schema". The object {{"properties": {{"foo": ["bar", "baz"]}}}} is not well-formatted. + +Your output will be parsed and type-checked according to the provided schema instance, so make sure all fields in your output match the schema exactly and there are no trailing commas! + +Here is the JSON Schema instance your output must adhere to. Include the enclosing markdown codeblock: +\`\`\`json +{"type":"object","properties":{"insights":{"type":"array","items":{"type":"object","properties":{"group":{"type":"string","description":"The program which is triggering the events"},"events":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string","description":"The event ID"},"endpointId":{"type":"string","description":"The endpoint ID"},"value":{"type":"string","description":"The process.executable value of the event"}},"required":["id","endpointId","value"],"additionalProperties":false},"description":"The events that the insight is based on"}},"required":["group","events"],"additionalProperties":false}}},"required":["insights"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"} +\`\`\` +`; + + const chainWithFormatInstructions = getChainWithFormatInstructions( + 'incompatible_antivirus', + mockLlm + ); + expect(chainWithFormatInstructions).toEqual({ + chain: expect.any(Object), + formatInstructions: expectedFormatInstructions, + llmType: 'fake', + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.ts new file mode 100644 index 0000000000000..4ffb1c30581c4 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_chain_with_format_instructions/index.ts @@ -0,0 +1,36 @@ +/* + * 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 type { ActionsClientLlm } from '@kbn/langchain/server'; +import { ChatPromptTemplate } from '@langchain/core/prompts'; +import { Runnable } from '@langchain/core/runnables'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { getOutputParser } from '../get_output_parser'; + +interface GetChainWithFormatInstructions { + chain: Runnable; + formatInstructions: string; + llmType: string; +} + +export const getChainWithFormatInstructions = ( + insightType: DefendInsightType, + llm: ActionsClientLlm +): GetChainWithFormatInstructions => { + const outputParser = getOutputParser({ type: insightType }); + const formatInstructions = outputParser.getFormatInstructions(); + + const prompt = ChatPromptTemplate.fromTemplate( + `Answer the user's question as best you can:\n{format_instructions}\n{query}` + ); + + const chain = prompt.pipe(llm); + const llmType = llm._llmType(); + + return { chain, formatInstructions, llmType }; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.test.ts new file mode 100644 index 0000000000000..75d7d83db3e92 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.test.ts @@ -0,0 +1,37 @@ +/* + * 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 { getCombined } from '.'; + +describe('getCombined', () => { + it('combines two strings correctly', () => { + const combinedGenerations = 'generation1'; + const partialResponse = 'response1'; + const expected = 'generation1response1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); + + it('handles empty combinedGenerations', () => { + const combinedGenerations = ''; + const partialResponse = 'response1'; + const expected = 'response1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); + + it('handles an empty partialResponse', () => { + const combinedGenerations = 'generation1'; + const partialResponse = ''; + const expected = 'generation1'; + const result = getCombined({ combinedGenerations, partialResponse }); + + expect(result).toEqual(expected); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.ts new file mode 100644 index 0000000000000..10b5c323891a1 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined/index.ts @@ -0,0 +1,14 @@ +/* + * 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. + */ + +export const getCombined = ({ + combinedGenerations, + partialResponse, +}: { + combinedGenerations: string; + partialResponse: string; +}): string => `${combinedGenerations}${partialResponse}`; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.test.ts new file mode 100644 index 0000000000000..ee1c3aa61dcdd --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.test.ts @@ -0,0 +1,62 @@ +/* + * 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 { getCombinedDefendInsightsPrompt } from '.'; + +describe('getCombinedDefendInsightsPrompt', () => { + it('returns the initial query when there are no partial results', () => { + const result = getCombinedDefendInsightsPrompt({ + anonymizedEvents: ['event1', 'event2'], + prompt: 'defendInsightsPrompt', + combinedMaybePartialResults: '', + }); + + expect(result).toBe(`defendInsightsPrompt + +Use context from the following events to provide insights: + +""" +event1 + +event2 +""" +`); + }); + + it('returns the initial query combined with a continuation prompt and partial results', () => { + const result = getCombinedDefendInsightsPrompt({ + anonymizedEvents: ['event1', 'event2'], + prompt: 'defendInsightsPrompt', + combinedMaybePartialResults: 'partialResults', + }); + + expect(result).toBe(`defendInsightsPrompt + +Use context from the following events to provide insights: + +""" +event1 + +event2 +""" + + +Continue exactly where you left off in the JSON output below, generating only the additional JSON output when it's required to complete your work. The additional JSON output MUST ALWAYS follow these rules: +1) it MUST conform to the schema above, because it will be checked against the JSON schema +2) it MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds), because it will be parsed as JSON +3) it MUST NOT repeat any the previous output, because that would prevent partial results from being combined +4) it MUST NOT restart from the beginning, because that would prevent partial results from being combined +5) it MUST NOT be prefixed or suffixed with additional text outside of the JSON, because that would prevent it from being combined and parsed as JSON: + + +""" +partialResults +""" + +`); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.ts new file mode 100644 index 0000000000000..2ccc26b1a0d7b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_combined_prompt/index.ts @@ -0,0 +1,43 @@ +/* + * 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 { isEmpty } from 'lodash/fp'; + +import { getEventsContextPrompt } from '../../generate/helpers/get_events_context_prompt'; +import { getContinuePrompt } from '../get_continue_prompt'; + +/** + * Returns the the initial query, or the initial query combined with a + * continuation prompt and partial results + */ +export const getCombinedDefendInsightsPrompt = ({ + anonymizedEvents, + prompt, + combinedMaybePartialResults, +}: { + anonymizedEvents: string[]; + prompt: string; + /** combined results that may contain incomplete JSON */ + combinedMaybePartialResults: string; +}): string => { + const eventsContextPrompt = getEventsContextPrompt({ + anonymizedEvents, + prompt, + }); + + return isEmpty(combinedMaybePartialResults) + ? eventsContextPrompt // no partial results yet + : `${eventsContextPrompt} + +${getContinuePrompt()} + +""" +${combinedMaybePartialResults} +""" + +`; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.test.ts new file mode 100644 index 0000000000000..35dae31a3ae6a --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.test.ts @@ -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. + */ + +import { getContinuePrompt } from '.'; + +describe('getContinuePrompt', () => { + it('returns the expected prompt string', () => { + const expectedPrompt = `Continue exactly where you left off in the JSON output below, generating only the additional JSON output when it's required to complete your work. The additional JSON output MUST ALWAYS follow these rules: +1) it MUST conform to the schema above, because it will be checked against the JSON schema +2) it MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds), because it will be parsed as JSON +3) it MUST NOT repeat any the previous output, because that would prevent partial results from being combined +4) it MUST NOT restart from the beginning, because that would prevent partial results from being combined +5) it MUST NOT be prefixed or suffixed with additional text outside of the JSON, because that would prevent it from being combined and parsed as JSON: +`; + + expect(getContinuePrompt()).toBe(expectedPrompt); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.ts new file mode 100644 index 0000000000000..628ba0531332c --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_continue_prompt/index.ts @@ -0,0 +1,15 @@ +/* + * 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. + */ + +export const getContinuePrompt = + (): string => `Continue exactly where you left off in the JSON output below, generating only the additional JSON output when it's required to complete your work. The additional JSON output MUST ALWAYS follow these rules: +1) it MUST conform to the schema above, because it will be checked against the JSON schema +2) it MUST escape all JSON special characters (i.e. backslashes, double quotes, newlines, tabs, carriage returns, backspaces, and form feeds), because it will be parsed as JSON +3) it MUST NOT repeat any the previous output, because that would prevent partial results from being combined +4) it MUST NOT restart from the beginning, because that would prevent partial results from being combined +5) it MUST NOT be prefixed or suffixed with additional text outside of the JSON, because that would prevent it from being combined and parsed as JSON: +`; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.test.ts new file mode 100644 index 0000000000000..df3229817c9ed --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.test.ts @@ -0,0 +1,31 @@ +/* + * 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 { getOutputParser } from '.'; + +describe('getOutputParser', () => { + it('returns a structured output parser with the expected format instructions', () => { + const outputParser = getOutputParser({ type: 'incompatible_antivirus' }); + + const expected = `You must format your output as a JSON value that adheres to a given \"JSON Schema\" instance. + +\"JSON Schema\" is a declarative language that allows you to annotate and validate JSON documents. + +For example, the example \"JSON Schema\" instance {{\"properties\": {{\"foo\": {{\"description\": \"a list of test words\", \"type\": \"array\", \"items\": {{\"type\": \"string\"}}}}}}, \"required\": [\"foo\"]}}}} +would match an object with one required property, \"foo\". The \"type\" property specifies \"foo\" must be an \"array\", and the \"description\" property semantically describes it as \"a list of test words\". The items within \"foo\" must be strings. +Thus, the object {{\"foo\": [\"bar\", \"baz\"]}} is a well-formatted instance of this example \"JSON Schema\". The object {{\"properties\": {{\"foo\": [\"bar\", \"baz\"]}}}} is not well-formatted. + +Your output will be parsed and type-checked according to the provided schema instance, so make sure all fields in your output match the schema exactly and there are no trailing commas! + +Here is the JSON Schema instance your output must adhere to. Include the enclosing markdown codeblock: +\`\`\`json +{"type":"object","properties":{"insights":{"type":"array","items":{"type":"object","properties":{"group":{"type":"string","description":"The program which is triggering the events"},"events":{"type":"array","items":{"type":"object","properties":{"id":{"type":"string","description":"The event ID"},"endpointId":{"type":"string","description":"The endpoint ID"},"value":{"type":"string","description":"The process.executable value of the event"}},"required":["id","endpointId","value"],"additionalProperties":false},"description":"The events that the insight is based on"}},"required":["group","events"],"additionalProperties":false}}},"required":["insights"],"additionalProperties":false,"$schema":"http://json-schema.org/draft-07/schema#"} +\`\`\` +`; + + expect(outputParser.getFormatInstructions()).toEqual(expected); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.ts new file mode 100644 index 0000000000000..dd8a5216304fe --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/get_output_parser/index.ts @@ -0,0 +1,16 @@ +/* + * 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 { StructuredOutputParser } from 'langchain/output_parsers'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { getSchema } from '../../generate/schema'; + +export function getOutputParser({ type }: { type: DefendInsightType }) { + const schema = getSchema({ type }); + return StructuredOutputParser.fromZodSchema(schema); +} diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/parse_combined_or_throw/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/parse_combined_or_throw/index.ts new file mode 100644 index 0000000000000..10d4db2e80e33 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/parse_combined_or_throw/index.ts @@ -0,0 +1,55 @@ +/* + * 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 type { Logger } from '@kbn/core/server'; +import type { DefendInsight, DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { getSchema } from '../../generate/schema'; +import { addTrailingBackticksIfNecessary } from '../add_trailing_backticks_if_necessary'; +import { extractJson } from '../extract_json'; + +export const parseCombinedOrThrow = ({ + insightType, + combinedResponse, + generationAttempts, + llmType, + logger, + nodeName, +}: { + /** combined responses that maybe valid JSON */ + insightType: DefendInsightType; + combinedResponse: string; + generationAttempts: number; + nodeName: string; + llmType: string; + logger?: Logger; +}): DefendInsight[] => { + const timestamp = new Date().toISOString(); + + const extractedJson = extractJson(addTrailingBackticksIfNecessary(combinedResponse)); + + logger?.debug( + () => + `${nodeName} node is parsing extractedJson (${llmType}) from attempt ${generationAttempts}` + ); + + const unvalidatedParsed = JSON.parse(extractedJson); + + logger?.debug( + () => + `${nodeName} node is validating combined response (${llmType}) from attempt ${generationAttempts}` + ); + + const validatedResponse = getSchema({ type: insightType }).parse(unvalidatedParsed); + + logger?.debug( + () => + `${nodeName} node successfully validated Defend insights response (${llmType}) from attempt ${generationAttempts}` + ); + + return [...validatedResponse.insights.map((insight) => ({ ...insight, timestamp }))]; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/incompatible_antivirus.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/incompatible_antivirus.ts new file mode 100644 index 0000000000000..14a6472c12e9d --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/incompatible_antivirus.ts @@ -0,0 +1,23 @@ +/* + * 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. + */ + +export function getIncompatibleAntivirusPrompt(events?: string[]): string { + const defaultPrompt = + 'You are an Elastic Security user tasked with analyzing file events from Elastic Security to identify antivirus processes. Only focus on detecting antivirus processes. Ignore processes that belong to Elastic Agent or Elastic Defend, that are not antivirus processes, or are typical processes built into the operating system. Accuracy is of the utmost importance, try to minimize false positives. Group the processes by the antivirus program, keeping track of the agent.id and _id associated to each of the individual events as endpointId and eventId respectively. If there are no events, ignore the group field. Escape backslashes to respect JSON validation. New lines must always be escaped with double backslashes, i.e. \\\\n to ensure valid JSON. Only return JSON output, as described above. Do not add any additional text to describe your output.'; + + if (!events) { + return defaultPrompt; + } + + return `${defaultPrompt} + + Use context from the following process events to provide insights: + """ + ${events.join('\n\n')} + """ + `; +} diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/index.ts similarity index 81% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/index.ts index d58778c3c544b..f36117963406b 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/prompts/index.ts @@ -7,7 +7,7 @@ import { DefendInsightType } from '@kbn/elastic-assistant-common'; -import { InvalidDefendInsightTypeError } from '../errors'; +import { InvalidDefendInsightTypeError } from '../../../../../errors'; import { getIncompatibleAntivirusPrompt } from './incompatible_antivirus'; export function getDefendInsightsPrompt({ @@ -15,10 +15,10 @@ export function getDefendInsightsPrompt({ events, }: { type: DefendInsightType; - events: string[]; + events?: string[]; }): string { if (type === DefendInsightType.Enum.incompatible_antivirus) { - return getIncompatibleAntivirusPrompt({ events }); + return getIncompatibleAntivirusPrompt(events); } throw new InvalidDefendInsightTypeError(); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.test.ts new file mode 100644 index 0000000000000..3730d6a7c4b96 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.test.ts @@ -0,0 +1,26 @@ +/* + * 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 { responseIsHallucinated } from '.'; + +describe('responseIsHallucinated', () => { + it('returns true when the response is hallucinated', () => { + expect( + responseIsHallucinated( + 'tactics like **Credential Access**, **Command and Control**, and **Persistence**.",\n "entitySummaryMarkdown": "Malware detected on host **{{ host.name hostNameValue }}**' + ) + ).toBe(true); + }); + + it('returns false when the response is not hallucinated', () => { + expect( + responseIsHallucinated( + 'A malicious file {{ file.name WsmpRExIFs.dll }} was detected on {{ host.name 082a86fa-b87d-45ce-813e-eed6b36ef0a9 }}\\n- The file was executed by' + ) + ).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.ts new file mode 100644 index 0000000000000..f938f6436db98 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/helpers/response_is_hallucinated/index.ts @@ -0,0 +1,9 @@ +/* + * 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. + */ + +export const responseIsHallucinated = (result: string): boolean => + result.includes('{{ host.name hostNameValue }}'); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts new file mode 100644 index 0000000000000..deec24149b59a --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.test.ts @@ -0,0 +1,102 @@ +/* + * 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 type { GraphState } from '../../../../types'; +import { discardPreviousRefinements } from '.'; + +const mockUnrefinedResults = [ + { + group: 'test-group-1', + events: [ + { + id: 'event-1', + endpointId: 'endpoint-1', + value: 'event value 1', + }, + ], + }, + { + group: 'test-group-2', + events: [ + { + id: 'event-2', + endpointId: 'endpoint-2', + value: 'event value 2', + }, + ], + }, +]; + +const initialState: GraphState = { + insights: null, + prompt: 'initial prompt', + anonymizedEvents: [], + combinedGenerations: 'generation1generation2', + combinedRefinements: 'refinement1', + errors: [], + generationAttempts: 3, + generations: ['generation1', 'generation2'], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: ['refinement1'], + refinePrompt: 'refinePrompt', + replacements: {}, + unrefinedResults: mockUnrefinedResults, +}; + +describe('discardPreviousRefinements', () => { + let result: GraphState; + + beforeEach(() => { + result = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: false, + state: initialState, + }); + }); + + it('resets the combined refinements', () => { + expect(result.combinedRefinements).toBe(''); + }); + + it('increments the generation attempts', () => { + expect(result.generationAttempts).toBe(initialState.generationAttempts + 1); + }); + + it('resets the refinements', () => { + expect(result.refinements).toEqual([]); + }); + + describe('hallucination scenarios', () => { + it('increments hallucination failures when hallucinations are detected', () => { + const hallucinationResult = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: true, + state: initialState, + }); + + expect(hallucinationResult.hallucinationFailures).toBe( + initialState.hallucinationFailures + 1 + ); + }); + + it('does NOT increment hallucination failures when hallucinations are NOT detected', () => { + const noHallucinationResult = discardPreviousRefinements({ + generationAttempts: initialState.generationAttempts, + hallucinationFailures: initialState.hallucinationFailures, + isHallucinationDetected: false, + state: initialState, + }); + + expect(noHallucinationResult.hallucinationFailures).toBe(initialState.hallucinationFailures); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.ts new file mode 100644 index 0000000000000..e642e598e73f0 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/discard_previous_refinements/index.ts @@ -0,0 +1,30 @@ +/* + * 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 { GraphState } from '../../../../types'; + +export const discardPreviousRefinements = ({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected, + state, +}: { + generationAttempts: number; + hallucinationFailures: number; + isHallucinationDetected: boolean; + state: GraphState; +}): GraphState => { + return { + ...state, + combinedRefinements: '', // <-- reset the combined refinements + generationAttempts: generationAttempts + 1, + refinements: [], // <-- reset the refinements + hallucinationFailures: isHallucinationDetected + ? hallucinationFailures + 1 + : hallucinationFailures, + }; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts new file mode 100644 index 0000000000000..3152c5f7aff29 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.test.ts @@ -0,0 +1,83 @@ +/* + * 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 { mockDefendInsights } from '../../../../mock/mock_defend_insights'; +import { getContinuePrompt } from '../../../helpers/get_continue_prompt'; +import { getCombinedRefinePrompt } from '.'; + +describe('getCombinedRefinePrompt', () => { + const mockPrompt = 'Initial prompt text'; + const mockRefinePrompt = 'Please refine these results'; + + it('returns base query when no combined refinements exist', () => { + const result = getCombinedRefinePrompt({ + prompt: mockPrompt, + combinedRefinements: '', + refinePrompt: mockRefinePrompt, + unrefinedResults: mockDefendInsights, + }); + + expect(result).toBe(`${mockPrompt} + +${mockRefinePrompt} + +""" +${JSON.stringify(mockDefendInsights, null, 2)} +""" + +`); + }); + + it('includes combined refinements and continue prompt when refinements exist', () => { + const mockRefinements = 'Previous refinement results'; + const result = getCombinedRefinePrompt({ + prompt: mockPrompt, + combinedRefinements: mockRefinements, + refinePrompt: mockRefinePrompt, + unrefinedResults: mockDefendInsights, + }); + + const baseQuery = `${mockPrompt} + +${mockRefinePrompt} + +""" +${JSON.stringify(mockDefendInsights, null, 2)} +""" + +`; + + expect(result).toBe(`${baseQuery} + +${getContinuePrompt()} + +""" +${mockRefinements} +""" + +`); + }); + + it('handles null unrefined results', () => { + const result = getCombinedRefinePrompt({ + prompt: mockPrompt, + combinedRefinements: '', + refinePrompt: mockRefinePrompt, + unrefinedResults: null, + }); + + expect(result).toBe(`${mockPrompt} + +${mockRefinePrompt} + +""" +null +""" + +`); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.ts new file mode 100644 index 0000000000000..9f6e7a1bbbfc7 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_combined_refine_prompt/index.ts @@ -0,0 +1,48 @@ +/* + * 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 type { DefendInsight } from '@kbn/elastic-assistant-common'; +import { isEmpty } from 'lodash/fp'; + +import { getContinuePrompt } from '../../../helpers/get_continue_prompt'; + +/** + * Returns a prompt that combines the initial query, a refine prompt, and partial results + */ +export const getCombinedRefinePrompt = ({ + prompt, + combinedRefinements, + refinePrompt, + unrefinedResults, +}: { + prompt: string; + combinedRefinements: string; + refinePrompt: string; + unrefinedResults: DefendInsight[] | null; +}): string => { + const baseQuery = `${prompt} + +${refinePrompt} + +""" +${JSON.stringify(unrefinedResults, null, 2)} +""" + +`; + + return isEmpty(combinedRefinements) + ? baseQuery // no partial results yet + : `${baseQuery} + +${getContinuePrompt()} + +""" +${combinedRefinements} +""" + +`; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_default_refine_prompt/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_default_refine_prompt/index.ts new file mode 100644 index 0000000000000..20967d3ab0f25 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_default_refine_prompt/index.ts @@ -0,0 +1,11 @@ +/* + * 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. + */ + +export const getDefaultRefinePrompt = + (): string => `You previously generated the following insights, but sometimes they include events that aren't from an antivirus program or are not grouped correctly by the same antivirus program. + +Review the insights below and remove any that are not from an antivirus program and combine duplicates into the same 'group'; leave any other insights unchanged:`; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts new file mode 100644 index 0000000000000..3b9aa160b4918 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.test.ts @@ -0,0 +1,46 @@ +/* + * 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 { getUseUnrefinedResults } from '.'; + +describe('getUseUnrefinedResults', () => { + it('returns true if both maxHallucinationFailuresReached and maxRetriesReached are true', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: true, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxHallucinationFailuresReached is true and maxRetriesReached is false', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: true, + maxRetriesReached: false, + }); + + expect(result).toBe(true); + }); + + it('returns true if maxHallucinationFailuresReached is false and maxRetriesReached is true', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: false, + maxRetriesReached: true, + }); + + expect(result).toBe(true); + }); + + it('returns false if both maxHallucinationFailuresReached and maxRetriesReached are false', () => { + const result = getUseUnrefinedResults({ + maxHallucinationFailuresReached: false, + maxRetriesReached: false, + }); + + expect(result).toBe(false); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.ts new file mode 100644 index 0000000000000..13d0a2228a3ee --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/helpers/get_use_unrefined_results/index.ts @@ -0,0 +1,17 @@ +/* + * 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. + */ + +/** + * Note: the conditions tested here are different than the generate node + */ +export const getUseUnrefinedResults = ({ + maxHallucinationFailuresReached, + maxRetriesReached, +}: { + maxHallucinationFailuresReached: boolean; + maxRetriesReached: boolean; +}): boolean => maxRetriesReached || maxHallucinationFailuresReached; // we may have reached max halucination failures, but we still want to use the unrefined results diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/index.ts new file mode 100644 index 0000000000000..e9d0dbe78088c --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/refine/index.ts @@ -0,0 +1,173 @@ +/* + * 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 type { ActionsClientLlm } from '@kbn/langchain/server'; +import type { Logger } from '@kbn/core/server'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import type { GraphState } from '../../types'; +import { getMaxHallucinationFailuresReached } from '../../helpers/get_max_hallucination_failures_reached'; +import { getMaxRetriesReached } from '../../helpers/get_max_retries_reached'; +import { getChainWithFormatInstructions } from '../helpers/get_chain_with_format_instructions'; +import { extractJson } from '../helpers/extract_json'; +import { getCombined } from '../helpers/get_combined'; +import { generationsAreRepeating } from '../helpers/generations_are_repeating'; +import { parseCombinedOrThrow } from '../helpers/parse_combined_or_throw'; +import { responseIsHallucinated } from '../helpers/response_is_hallucinated'; +import { discardPreviousRefinements } from './helpers/discard_previous_refinements'; +import { getCombinedRefinePrompt } from './helpers/get_combined_refine_prompt'; +import { getUseUnrefinedResults } from './helpers/get_use_unrefined_results'; + +export const getRefineNode = ({ + insightType, + llm, + logger, +}: { + insightType: DefendInsightType; + llm: ActionsClientLlm; + logger?: Logger; +}): ((state: GraphState) => Promise) => { + const refine = async (state: GraphState): Promise => { + logger?.debug(() => '---REFINE---'); + + const { + prompt, + combinedRefinements, + generationAttempts, + hallucinationFailures, + maxGenerationAttempts, + maxHallucinationFailures, + maxRepeatedGenerations, + refinements, + refinePrompt, + unrefinedResults, + } = state; + + let combinedResponse = ''; // mutable, because it must be accessed in the catch block + let partialResponse = ''; // mutable, because it must be accessed in the catch block + + try { + const query = getCombinedRefinePrompt({ + prompt, + combinedRefinements, + refinePrompt, + unrefinedResults, + }); + + const { chain, formatInstructions, llmType } = getChainWithFormatInstructions( + insightType, + llm + ); + + logger?.debug( + () => `refine node is invoking the chain (${llmType}), attempt ${generationAttempts}` + ); + + const rawResponse = (await chain.invoke({ + format_instructions: formatInstructions, + query, + })) as unknown as string; + + // LOCAL MUTATION: + partialResponse = extractJson(rawResponse); // remove the surrounding ```json``` + + // if the response is hallucinated, discard it: + if (responseIsHallucinated(partialResponse)) { + logger?.debug( + () => + `refine node detected a hallucination (${llmType}), on attempt ${generationAttempts}; discarding the accumulated refinements and starting over` + ); + + return discardPreviousRefinements({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected: true, + state, + }); + } + + // if the refinements are repeating, discard previous refinements and start over: + if ( + generationsAreRepeating({ + currentGeneration: partialResponse, + previousGenerations: refinements, + sampleLastNGenerations: maxRepeatedGenerations - 1, + }) + ) { + logger?.debug( + () => + `refine node detected (${llmType}), detected ${maxRepeatedGenerations} repeated generations on attempt ${generationAttempts}; discarding the accumulated results and starting over` + ); + + // discard the accumulated results and start over: + return discardPreviousRefinements({ + generationAttempts, + hallucinationFailures, + isHallucinationDetected: false, + state, + }); + } + + // LOCAL MUTATION: + combinedResponse = getCombined({ combinedGenerations: combinedRefinements, partialResponse }); // combine the new response with the previous ones + + const defendInsights = parseCombinedOrThrow({ + insightType, + combinedResponse, + generationAttempts, + llmType, + logger, + nodeName: 'refine', + }); + + return { + ...state, + insights: defendInsights, // the final, refined answer + generationAttempts: generationAttempts + 1, + combinedRefinements: combinedResponse, + refinements: [...refinements, partialResponse], + }; + } catch (error) { + const parsingError = `refine node is unable to parse (${llm._llmType()}) response from attempt ${generationAttempts}; (this may be an incomplete response from the model): ${error}`; + logger?.debug(() => parsingError); // logged at debug level because the error is expected when the model returns an incomplete response + + const maxRetriesReached = getMaxRetriesReached({ + generationAttempts: generationAttempts + 1, + maxGenerationAttempts, + }); + + const maxHallucinationFailuresReached = getMaxHallucinationFailuresReached({ + hallucinationFailures, + maxHallucinationFailures, + }); + + // we will use the unrefined results if we have reached the maximum number of retries or hallucination failures: + const useUnrefinedResults = getUseUnrefinedResults({ + maxHallucinationFailuresReached, + maxRetriesReached, + }); + + if (useUnrefinedResults) { + logger?.debug( + () => + `refine node is using unrefined results response (${llm._llmType()}) from attempt ${generationAttempts}, because all attempts have been used` + ); + } + + return { + ...state, + insights: useUnrefinedResults ? unrefinedResults : null, + combinedRefinements: combinedResponse, + errors: [...state.errors, parsingError], + generationAttempts: generationAttempts + 1, + refinements: [...refinements, partialResponse], + }; + } + }; + + return refine; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.test.ts new file mode 100644 index 0000000000000..d8cde706b3a36 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.test.ts @@ -0,0 +1,101 @@ +/* + * 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 type { ElasticsearchClient } from '@kbn/core/server'; +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { mockAnonymizationFields } from '../../../mock/mock_anonymization_fields'; +import { getAnonymizedEvents } from '../helpers/get_anonymized_events'; +import { mockAnonymizedEvents } from '../../../mock/mock_anonymized_events'; +import { AnonymizedEventsRetriever } from '.'; + +jest.mock('../helpers/get_anonymized_events', () => ({ + getAnonymizedEvents: jest.fn(), +})); + +describe('AnonymizedEventsRetriever', () => { + let esClient: ElasticsearchClient; + + beforeEach(() => { + jest.clearAllMocks(); + + esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + + (getAnonymizedEvents as jest.Mock).mockResolvedValue([...mockAnonymizedEvents]); + }); + + it('returns the expected pageContent and metadata', async () => { + const retriever = new AnonymizedEventsRetriever({ + insightType: 'incompatible_antivirus' as DefendInsightType, + endpointIds: ['endpoint-1'], + anonymizationFields: mockAnonymizationFields, + esClient, + size: 10, + }); + + const documents = await retriever._getRelevantDocuments('test-query'); + + expect(documents).toEqual([ + { + pageContent: mockAnonymizedEvents[0], + metadata: {}, + }, + { + pageContent: mockAnonymizedEvents[1], + metadata: {}, + }, + ]); + }); + + it('calls getAnonymizedEvents with the expected parameters', async () => { + const onNewReplacements = jest.fn(); + const mockReplacements = { + replacement1: 'SRVMAC08', + replacement2: 'SRVWIN01', + replacement3: 'SRVWIN02', + }; + + const retriever = new AnonymizedEventsRetriever({ + insightType: 'incompatible_antivirus' as DefendInsightType, + endpointIds: ['endpoint-1'], + anonymizationFields: mockAnonymizationFields, + esClient, + onNewReplacements, + replacements: mockReplacements, + size: 10, + }); + + await retriever._getRelevantDocuments('test-query'); + + expect(getAnonymizedEvents).toHaveBeenCalledWith({ + insightType: 'incompatible_antivirus', + endpointIds: ['endpoint-1'], + anonymizationFields: mockAnonymizationFields, + esClient, + onNewReplacements, + replacements: mockReplacements, + size: 10, + }); + }); + + it('handles empty anonymized events', async () => { + (getAnonymizedEvents as jest.Mock).mockResolvedValue([]); + + const retriever = new AnonymizedEventsRetriever({ + insightType: 'incompatible_antivirus' as DefendInsightType, + endpointIds: ['endpoint-1'], + anonymizationFields: mockAnonymizationFields, + esClient, + size: 10, + }); + + const documents = await retriever._getRelevantDocuments('test-query'); + + expect(documents).toHaveLength(0); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.ts new file mode 100644 index 0000000000000..fc8b5ab848f30 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/anonymized_events_retriever/index.ts @@ -0,0 +1,90 @@ +/* + * 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 type { CallbackManagerForRetrieverRun } from '@langchain/core/callbacks/manager'; +import type { Document } from '@langchain/core/documents'; +import type { DateMath } from '@elastic/elasticsearch/lib/api/types'; +import type { ElasticsearchClient } from '@kbn/core/server'; +import { BaseRetriever, type BaseRetrieverInput } from '@langchain/core/retrievers'; +import { DefendInsightType, Replacements } from '@kbn/elastic-assistant-common'; +import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +import { getAnonymizedEvents } from '../helpers/get_anonymized_events'; + +export type CustomRetrieverInput = BaseRetrieverInput; + +export class AnonymizedEventsRetriever extends BaseRetriever { + lc_namespace = ['langchain', 'retrievers']; + + private readonly insightType: DefendInsightType; + private readonly endpointIds: string[]; + private readonly anonymizationFields?: AnonymizationFieldResponse[]; + private readonly esClient: ElasticsearchClient; + private readonly onNewReplacements?: (newReplacements: Replacements) => void; + private readonly replacements?: Replacements; + private readonly size?: number; + private readonly start?: DateMath; + private readonly end?: DateMath; + + constructor({ + insightType, + endpointIds, + anonymizationFields, + fields, + esClient, + onNewReplacements, + replacements, + size, + start, + end, + }: { + insightType: DefendInsightType; + endpointIds: string[]; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + fields?: CustomRetrieverInput; + onNewReplacements?: (newReplacements: Replacements) => void; + replacements?: Replacements; + size?: number; + start?: DateMath; + end?: DateMath; + }) { + super(fields); + + this.insightType = insightType; + this.endpointIds = endpointIds; + this.anonymizationFields = anonymizationFields; + this.esClient = esClient; + this.onNewReplacements = onNewReplacements; + this.replacements = replacements; + this.size = size; + this.start = start; + this.end = end; + } + + async _getRelevantDocuments( + query: string, + runManager?: CallbackManagerForRetrieverRun + ): Promise { + const anonymizedEvents = await getAnonymizedEvents({ + insightType: this.insightType, + endpointIds: this.endpointIds, + anonymizationFields: this.anonymizationFields, + esClient: this.esClient, + onNewReplacements: this.onNewReplacements, + replacements: this.replacements, + size: this.size, + start: this.start, + end: this.end, + }); + + return anonymizedEvents.map((event) => ({ + pageContent: event, + metadata: {}, + })); + } +} diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/get_file_events_query.ts similarity index 72% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/get_file_events_query.ts index eef2e1ad28f16..ba624ca4729e9 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/get_file_events_query.ts @@ -5,13 +5,22 @@ * 2.0. */ -import type { SearchRequest } from '@elastic/elasticsearch/lib/api/types'; +import type { SearchRequest, DateMath } from '@elastic/elasticsearch/lib/api/types'; -import { FILE_EVENTS_INDEX_PATTERN } from '../../../../../common/endpoint/constants'; +const FILE_EVENTS_INDEX_PATTERN = 'logs-endpoint.events.file-*'; +const SIZE = 1500; -const SIZE = 200; - -export function getFileEventsQuery({ endpointIds }: { endpointIds: string[] }): SearchRequest { +export function getFileEventsQuery({ + endpointIds, + size, + gte, + lte, +}: { + endpointIds: string[]; + size?: number; + gte?: DateMath; + lte?: DateMath; +}): SearchRequest { return { allow_no_indices: true, query: { @@ -25,8 +34,9 @@ export function getFileEventsQuery({ endpointIds }: { endpointIds: string[] }): { range: { '@timestamp': { - gte: 'now-24h', - lte: 'now', + gte: gte ?? 'now-14d', + // gte: gte ?? 'now-24h', + lte: lte ?? 'now', }, }, }, @@ -38,7 +48,7 @@ export function getFileEventsQuery({ endpointIds }: { endpointIds: string[] }): unique_process_executable: { terms: { field: 'process.executable', - size: SIZE, + size: size ?? SIZE, }, aggs: { // Get the latest event for each process.executable diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/index.ts new file mode 100644 index 0000000000000..512eb5f891e6f --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/get_events/index.ts @@ -0,0 +1,29 @@ +/* + * 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 type { DateMath, SearchRequest } from '@elastic/elasticsearch/lib/api/types'; +import { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { getFileEventsQuery } from './get_file_events_query'; + +export function getQuery( + type: DefendInsightType, + options: { endpointIds: string[]; size?: number; gte?: DateMath; lte?: DateMath } +): SearchRequest { + if (type === DefendInsightType.Enum.incompatible_antivirus) { + const { endpointIds, size, gte, lte } = options; + + return getFileEventsQuery({ + endpointIds, + size, + gte, + lte, + }); + } + + throw new Error('Invalid defend insight type'); +} diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.test.ts new file mode 100644 index 0000000000000..11c8deafd2742 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.test.ts @@ -0,0 +1,271 @@ +/* + * 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 type { ElasticsearchClient } from '@kbn/core/server'; +import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; +import { + DefendInsightType, + getRawDataOrDefault, + transformRawData, +} from '@kbn/elastic-assistant-common'; + +import { mockAnonymizationFields } from '../../../../mock/mock_anonymization_fields'; +import { mockAnonymizedEventsReplacements } from '../../../../mock/mock_anonymized_events'; +import { getAnonymizedEvents } from '.'; + +jest.mock('@kbn/elastic-assistant-common', () => ({ + ...jest.requireActual('@kbn/elastic-assistant-common'), + getRawDataOrDefault: jest.fn(), + transformRawData: jest.fn(), +})); + +const createMockEsClient = () => { + return { + search: jest.fn(), + } as unknown as jest.Mocked; +}; + +const mockRawData: Record = { + field1: ['value1'], + field2: ['value2'], +}; + +describe('getAnonymizedEvents', () => { + const mockEsClient = createMockEsClient(); + const mockedGetRawDataOrDefault = getRawDataOrDefault as jest.MockedFunction< + typeof getRawDataOrDefault + >; + const mockedTransformRawData = transformRawData as jest.MockedFunction; + + beforeEach(() => { + jest.clearAllMocks(); + mockEsClient.search.mockResolvedValue({ + aggregations: { + unique_process_executable: { + buckets: [ + { + key: 'process1', + doc_count: 1, + latest_event: { + hits: { + hits: [ + { + _id: 'event1', + _source: { + agent: { id: 'agent1' }, + process: { executable: 'executable1' }, + }, + }, + ], + }, + }, + }, + ], + }, + }, + } as unknown as SearchResponse); + mockedGetRawDataOrDefault.mockReturnValue(mockRawData); + mockedTransformRawData.mockReturnValue('transformed data'); + }); + + it('should return an empty array when insightType is null', async () => { + const result = await getAnonymizedEvents({ + insightType: null as unknown as DefendInsightType, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(result).toEqual([]); + expect(mockEsClient.search).not.toHaveBeenCalled(); + }); + + it('should properly handle missing aggregations in response', async () => { + mockEsClient.search.mockResolvedValue({} as SearchResponse); + + const result = await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(result).toEqual([]); + }); + + it('should properly handle required parameters', async () => { + const result = await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(mockEsClient.search).toHaveBeenCalled(); + expect(result).toBeDefined(); + expect(Array.isArray(result)).toBe(true); + }); + + it('should call getRawDataOrDefault with correct fields', async () => { + await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(mockedGetRawDataOrDefault).toHaveBeenCalledWith({ + _id: ['event1'], + 'agent.id': ['agent1'], + 'process.executable': ['executable1'], + }); + }); + + it('should handle anonymizationFields when provided', async () => { + const onNewReplacements = jest.fn(); + + await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + anonymizationFields: mockAnonymizationFields, + onNewReplacements, + }); + + expect(mockedTransformRawData).toHaveBeenCalledWith( + expect.objectContaining({ + anonymizationFields: mockAnonymizationFields, + rawData: mockRawData, + }) + ); + }); + + it('should use existing replacements when provided', async () => { + const onNewReplacements = jest.fn(); + + await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + anonymizationFields: mockAnonymizationFields, + replacements: mockAnonymizedEventsReplacements, + onNewReplacements, + }); + + expect(mockedTransformRawData).toHaveBeenCalledWith( + expect.objectContaining({ + currentReplacements: expect.objectContaining(mockAnonymizedEventsReplacements), + }) + ); + }); + + it('should handle date range parameters', async () => { + const start = 'now-24h'; + const end = 'now'; + + await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + start, + end, + }); + + expect(mockEsClient.search).toHaveBeenCalledWith( + expect.objectContaining({ + query: { + bool: { + must: expect.arrayContaining([ + { + range: { + '@timestamp': { + gte: start, + lte: end, + }, + }, + }, + ]), + }, + }, + }) + ); + }); + + it('should handle size parameter', async () => { + const size = 5; + + await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + size, + }); + + expect(mockEsClient.search).toHaveBeenCalledWith( + expect.objectContaining({ + aggs: { + unique_process_executable: { + terms: { + field: 'process.executable', + size, + }, + aggs: { + latest_event: { + top_hits: { + _source: ['_id', 'agent.id', 'process.executable'], + size: 1, + sort: [{ '@timestamp': { order: 'desc' } }], + }, + }, + }, + }, + }, + }) + ); + }); + + it('should handle ES search errors', async () => { + const error = new Error('ES Search failed'); + mockEsClient.search.mockRejectedValue(error); + + await expect( + getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }) + ).rejects.toThrow(error); + }); + + it('should handle empty search results', async () => { + mockEsClient.search.mockResolvedValue({ + aggregations: { + unique_process_executable: { + buckets: [], + }, + }, + } as unknown as SearchResponse); + + const result = await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(result).toEqual([]); + }); + + it('should properly transform fields using anonymization rules', async () => { + const transformedValue = 'anonymized data'; + mockedTransformRawData.mockReturnValue(transformedValue); + + const result = await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + anonymizationFields: mockAnonymizationFields, + }); + + expect(result).toEqual(expect.arrayContaining([transformedValue])); + }); +}); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.ts similarity index 56% rename from x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.ts index 4f120e8e655f0..1902f86adbc7a 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/index.ts @@ -5,20 +5,18 @@ * 2.0. */ -import type { SearchRequest } from '@elastic/elasticsearch/lib/api/types'; +import type { DateMath, SearchResponse } from '@elastic/elasticsearch/lib/api/types'; import type { ElasticsearchClient } from '@kbn/core/server'; -import type { Replacements } from '@kbn/elastic-assistant-common'; -import type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; - import { - getAnonymizedValue, - transformRawData, DefendInsightType, + Replacements, + getAnonymizedValue, getRawDataOrDefault, + transformRawData, } from '@kbn/elastic-assistant-common'; -import { getFileEventsQuery } from './get_file_events_query'; -import { InvalidDefendInsightTypeError } from '../errors'; +import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; +import { getQuery } from './get_events'; interface AggregationResponse { unique_process_executable: { @@ -40,57 +38,34 @@ interface AggregationResponse { }; } -export async function getAnonymizedEvents({ +export const getAnonymizedEvents = async ({ + insightType, endpointIds, - type, anonymizationFields, esClient, onNewReplacements, replacements, + size, + start, + end, }: { + insightType: DefendInsightType; endpointIds: string[]; - type: DefendInsightType; anonymizationFields?: AnonymizationFieldResponse[]; esClient: ElasticsearchClient; onNewReplacements?: (replacements: Replacements) => void; replacements?: Replacements; -}): Promise { - const query = getQuery(type, { endpointIds }); - - return getAnonymized({ - query, - anonymizationFields, - esClient, - onNewReplacements, - replacements, - }); -} - -function getQuery(type: DefendInsightType, options: { endpointIds: string[] }): SearchRequest { - if (type === DefendInsightType.Enum.incompatible_antivirus) { - const { endpointIds } = options; - return getFileEventsQuery({ - endpointIds, - }); + size?: number; + start?: DateMath; + end?: DateMath; +}): Promise => { + if (insightType == null) { + return []; } - throw new InvalidDefendInsightTypeError(); -} - -const getAnonymized = async ({ - query, - anonymizationFields, - esClient, - onNewReplacements, - replacements, -}: { - query: SearchRequest; - anonymizationFields?: AnonymizationFieldResponse[]; - esClient: ElasticsearchClient; - onNewReplacements?: (replacements: Replacements) => void; - replacements?: Replacements; -}): Promise => { - const result = await esClient.search<{}, AggregationResponse>(query); + const query = getQuery(insightType, { endpointIds, size, gte: start, lte: end }); + // TODO add support for other insight types + const result = await esClient.search(query); const fileEvents = (result.aggregations?.unique_process_executable.buckets ?? []).map( (bucket) => { const latestEvent = bucket.latest_event.hits.hits[0]; @@ -102,8 +77,6 @@ const getAnonymized = async ({ } ); - // Accumulate replacements locally so we can, for example use the same - // replacement for a hostname when we see it in multiple alerts: let localReplacements = { ...(replacements ?? {}) }; const localOnNewReplacements = (newReplacements: Replacements) => { localReplacements = { ...localReplacements, ...newReplacements }; @@ -111,13 +84,13 @@ const getAnonymized = async ({ onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements }; - return fileEvents.map((fileEvent) => + return fileEvents.map((event) => transformRawData({ anonymizationFields, currentReplacements: localReplacements, // <-- the latest local replacements getAnonymizedValue, onNewReplacements: localOnNewReplacements, // <-- the local callback - rawData: getRawDataOrDefault(fileEvent), + rawData: getRawDataOrDefault(event), }) ); }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.test.ts new file mode 100644 index 0000000000000..7e9dfa389bd2b --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.test.ts @@ -0,0 +1,118 @@ +/* + * 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 { ElasticsearchClient, Logger } from '@kbn/core/server'; +import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks'; +import { DefendInsightType, Replacements } from '@kbn/elastic-assistant-common'; + +import type { GraphState } from '../../types'; +import { mockAnonymizedEvents } from '../../mock/mock_anonymized_events'; +import { getDefaultRefinePrompt } from '../refine/helpers/get_default_refine_prompt'; +import { getDefendInsightsPrompt } from '../helpers/prompts'; +import { getRetrieveAnonymizedEventsNode } from '.'; + +const insightType = DefendInsightType.Enum.incompatible_antivirus; +const initialGraphState: GraphState = { + insights: null, + prompt: getDefendInsightsPrompt({ type: insightType }), + anonymizedEvents: [], + combinedGenerations: '', + combinedRefinements: '', + errors: [], + generationAttempts: 0, + generations: [], + hallucinationFailures: 0, + maxGenerationAttempts: 10, + maxHallucinationFailures: 5, + maxRepeatedGenerations: 3, + refinements: [], + refinePrompt: getDefaultRefinePrompt(), + replacements: {}, + unrefinedResults: null, +}; + +jest.mock('./anonymized_events_retriever', () => ({ + AnonymizedEventsRetriever: jest + .fn() + .mockImplementation( + ({ + onNewReplacements, + replacements, + }: { + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; + }) => ({ + withConfig: jest.fn().mockReturnValue({ + invoke: jest.fn(async () => { + if (onNewReplacements != null && replacements != null) { + onNewReplacements(replacements); + } + + return mockAnonymizedEvents; + }), + }), + }) + ), +})); + +describe('getRetrieveAnonymizedEventsNode', () => { + const logger = { + debug: jest.fn(), + } as unknown as Logger; + + let esClient: ElasticsearchClient; + + beforeEach(() => { + esClient = elasticsearchClientMock.createScopedClusterClient().asCurrentUser; + }); + + it('returns a function', () => { + const result = getRetrieveAnonymizedEventsNode({ + insightType, + endpointIds: [], + esClient, + logger, + }); + expect(typeof result).toBe('function'); + }); + + it('updates state with anonymized events', async () => { + const state: GraphState = { ...initialGraphState }; + + const retrieveAnonymizedEvents = getRetrieveAnonymizedEventsNode({ + insightType, + endpointIds: [], + esClient, + logger, + }); + + const result = await retrieveAnonymizedEvents(state); + + expect(result).toHaveProperty('anonymizedEvents', mockAnonymizedEvents); + }); + + it('calls onNewReplacements with updated replacements', async () => { + const state: GraphState = { ...initialGraphState }; + const onNewReplacements = jest.fn(); + const replacements = { key: 'value' }; + + const retrieveAnonymizedEvents = getRetrieveAnonymizedEventsNode({ + insightType, + endpointIds: [], + esClient, + logger, + onNewReplacements, + replacements, + }); + + await retrieveAnonymizedEvents(state); + + expect(onNewReplacements).toHaveBeenCalledWith({ + ...replacements, + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.ts new file mode 100644 index 0000000000000..84a00d16812ca --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/index.ts @@ -0,0 +1,71 @@ +/* + * 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 type { ElasticsearchClient, Logger } from '@kbn/core/server'; +import type { DefendInsightType } from '@kbn/elastic-assistant-common'; +import { Replacements } from '@kbn/elastic-assistant-common'; +import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; + +import type { GraphState } from '../../types'; +import { AnonymizedEventsRetriever } from './anonymized_events_retriever'; + +export const getRetrieveAnonymizedEventsNode = ({ + insightType, + endpointIds, + anonymizationFields, + esClient, + logger, + onNewReplacements, + replacements, + size, +}: { + insightType: DefendInsightType; + endpointIds: string[]; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + logger?: Logger; + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; + size?: number; +}): ((state: GraphState) => Promise) => { + let localReplacements = { ...(replacements ?? {}) }; + const localOnNewReplacements = (newReplacements: Replacements) => { + localReplacements = { ...localReplacements, ...newReplacements }; + + onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements + }; + + const retrieveAnonymizedEvents = async (state: GraphState): Promise => { + logger?.debug(() => '---RETRIEVE ANONYMIZED EVENTS---'); + + const { start, end } = state; + + const retriever = new AnonymizedEventsRetriever({ + insightType, + endpointIds, + anonymizationFields, + esClient, + onNewReplacements: localOnNewReplacements, + replacements, + size, + start, + end, + }); + + const documents = await retriever + .withConfig({ runName: 'runAnonymizedEventsRetriever' }) + .invoke(''); + + return { + ...state, + anonymizedEvents: documents, + replacements: localReplacements, + }; + }; + + return retrieveAnonymizedEvents; +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.test.ts new file mode 100644 index 0000000000000..ca128e01f4633 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.test.ts @@ -0,0 +1,120 @@ +/* + * 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 { DefendInsightType } from '@kbn/elastic-assistant-common'; + +import { getDefaultRefinePrompt } from '../nodes/refine/helpers/get_default_refine_prompt'; +import { getDefendInsightsPrompt } from '../nodes/helpers/prompts'; +import { + DEFAULT_MAX_GENERATION_ATTEMPTS, + DEFAULT_MAX_HALLUCINATION_FAILURES, + DEFAULT_MAX_REPEATED_GENERATIONS, +} from '../constants'; +import { getDefaultGraphState } from '.'; + +const defaultInsightType = DefendInsightType.Enum.incompatible_antivirus; +const defaultDefendInsightsPrompt = getDefendInsightsPrompt({ + type: defaultInsightType, +}); +const defaultRefinePrompt = getDefaultRefinePrompt(); + +describe('getDefaultGraphState', () => { + it('returns the expected default defend insights', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.insights?.default?.()).toBeNull(); + }); + + it('returns the expected default prompt', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.prompt?.default?.()).toEqual(defaultDefendInsightsPrompt); + }); + + it('returns the expected default empty collection of anonymizedEvents', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.anonymizedEvents?.default?.()).toHaveLength(0); + }); + + it('returns the expected default combinedGenerations state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.combinedGenerations?.default?.()).toBe(''); + }); + + it('returns the expected default combinedRefinements state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.combinedRefinements?.default?.()).toBe(''); + }); + + it('returns the expected default errors state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.errors?.default?.()).toHaveLength(0); + }); + + it('return the expected default generationAttempts state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.generationAttempts?.default?.()).toBe(0); + }); + + it('returns the expected default generations state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.generations?.default?.()).toHaveLength(0); + }); + + it('returns the expected default hallucinationFailures state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.hallucinationFailures?.default?.()).toBe(0); + }); + + it('returns the expected default refinePrompt state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.refinePrompt?.default?.()).toEqual(defaultRefinePrompt); + }); + + it('returns the expected default maxGenerationAttempts state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.maxGenerationAttempts?.default?.()).toBe(DEFAULT_MAX_GENERATION_ATTEMPTS); + }); + + it('returns the expected default maxHallucinationFailures state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + expect(state.maxHallucinationFailures?.default?.()).toBe(DEFAULT_MAX_HALLUCINATION_FAILURES); + }); + + it('returns the expected default maxRepeatedGenerations state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.maxRepeatedGenerations?.default?.()).toBe(DEFAULT_MAX_REPEATED_GENERATIONS); + }); + + it('returns the expected default refinements state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.refinements?.default?.()).toHaveLength(0); + }); + + it('returns the expected default replacements state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.replacements?.default?.()).toEqual({}); + }); + + it('returns the expected default unrefinedResults state', () => { + const state = getDefaultGraphState({ insightType: defaultInsightType }); + + expect(state.unrefinedResults?.default?.()).toBeNull(); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.ts new file mode 100644 index 0000000000000..cf5bb590ad306 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/state/index.ts @@ -0,0 +1,105 @@ +/* + * 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 type { Document } from '@langchain/core/documents'; +import type { StateGraphArgs } from '@langchain/langgraph'; +import type { DateMath } from '@elastic/elasticsearch/lib/api/types'; +import type { DefendInsight, DefendInsightType, Replacements } from '@kbn/elastic-assistant-common'; + +import type { GraphState } from '../types'; +import { getDefendInsightsPrompt } from '../nodes/helpers/prompts'; +import { getDefaultRefinePrompt } from '../nodes/refine/helpers/get_default_refine_prompt'; +import { + DEFAULT_MAX_GENERATION_ATTEMPTS, + DEFAULT_MAX_HALLUCINATION_FAILURES, + DEFAULT_MAX_REPEATED_GENERATIONS, +} from '../constants'; + +export interface Options { + insightType: DefendInsightType; + start?: string; + end?: string; +} + +export const getDefaultGraphState = ({ + insightType, + start, + end, +}: Options): StateGraphArgs['channels'] => ({ + insights: { + value: (current: DefendInsight[] | null, next?: DefendInsight[] | null) => next ?? current, + default: () => null, + }, + prompt: { + value: (current: string, next?: string) => next ?? current, + default: () => getDefendInsightsPrompt({ type: insightType }), + }, + anonymizedEvents: { + value: (current: Document[], next?: Document[]) => next ?? current, + default: () => [], + }, + combinedGenerations: { + value: (current: string, next?: string) => next ?? current, + default: () => '', + }, + combinedRefinements: { + value: (current: string, next?: string) => next ?? current, + default: () => '', + }, + errors: { + value: (current: string[], next?: string[]) => next ?? current, + default: () => [], + }, + generationAttempts: { + value: (current: number, next?: number) => next ?? current, + default: () => 0, + }, + generations: { + value: (current: string[], next?: string[]) => next ?? current, + default: () => [], + }, + hallucinationFailures: { + value: (current: number, next?: number) => next ?? current, + default: () => 0, + }, + refinePrompt: { + value: (current: string, next?: string) => next ?? current, + default: () => getDefaultRefinePrompt(), + }, + maxGenerationAttempts: { + value: (current: number, next?: number) => next ?? current, + default: () => DEFAULT_MAX_GENERATION_ATTEMPTS, + }, + maxHallucinationFailures: { + value: (current: number, next?: number) => next ?? current, + default: () => DEFAULT_MAX_HALLUCINATION_FAILURES, + }, + maxRepeatedGenerations: { + value: (current: number, next?: number) => next ?? current, + default: () => DEFAULT_MAX_REPEATED_GENERATIONS, + }, + refinements: { + value: (current: string[], next?: string[]) => next ?? current, + default: () => [], + }, + replacements: { + value: (current: Replacements, next?: Replacements) => next ?? current, + default: () => ({}), + }, + unrefinedResults: { + value: (current: DefendInsight[] | null, next?: DefendInsight[] | null) => next ?? current, + default: () => null, + }, + start: { + value: (current?: DateMath, next?: DateMath) => next ?? current, + default: () => start, + }, + end: { + value: (current?: DateMath, next?: DateMath) => next ?? current, + default: () => end, + }, +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/types.ts new file mode 100644 index 0000000000000..aa4db037a6490 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/types.ts @@ -0,0 +1,31 @@ +/* + * 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 type { Document } from '@langchain/core/documents'; +import type { DefendInsight, Replacements } from '@kbn/elastic-assistant-common'; +import type { DateMath } from '@elastic/elasticsearch/lib/api/types'; + +export interface GraphState { + insights: DefendInsight[] | null; + prompt: string; + anonymizedEvents: Document[]; + combinedGenerations: string; + combinedRefinements: string; + errors: string[]; + generationAttempts: number; + generations: string[]; + hallucinationFailures: number; + maxGenerationAttempts: number; + maxHallucinationFailures: number; + maxRepeatedGenerations: number; + refinements: string[]; + refinePrompt: string; + replacements: Replacements; + unrefinedResults: DefendInsight[] | null; + start?: DateMath; + end?: DateMath; +} diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/field_maps_configuration.ts similarity index 100% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/field_maps_configuration.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/field_maps_configuration.ts diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts similarity index 95% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts index 415487534a1b6..981c598335bf1 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts @@ -6,11 +6,10 @@ */ import type { AuthenticatedUser } from '@kbn/core-security-common'; - import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; import { loggerMock } from '@kbn/logging-mocks'; -import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsightsSearchEsMock } from '../../../__mocks__/defend_insights_schema.mock'; import { getDefendInsight } from './get_defend_insight'; const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.ts similarity index 100% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/get_defend_insight.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.ts diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.test.ts similarity index 100% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.test.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.test.ts diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.ts similarity index 100% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/helpers.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.ts diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts similarity index 98% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts index 704ee9b962554..0d2192c82c2e9 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts @@ -14,14 +14,12 @@ import type { DefendInsightsGetRequestQuery, DefendInsightsResponse, } from '@kbn/elastic-assistant-common'; - import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; import { loggerMock } from '@kbn/logging-mocks'; import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; -import type { AIAssistantDataClientParams } from '..'; - -import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; +import type { AIAssistantDataClientParams } from '../../../ai_assistant_data_clients'; +import { getDefendInsightsSearchEsMock } from '../../../__mocks__/defend_insights_schema.mock'; import { getDefendInsight } from './get_defend_insight'; import { queryParamsToEsQuery, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts similarity index 97% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts index 064d918ec9a16..5518db11ced7c 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts @@ -5,8 +5,6 @@ * 2.0. */ -import { v4 as uuidv4 } from 'uuid'; - import type { DefendInsightCreateProps, DefendInsightUpdateProps, @@ -15,11 +13,11 @@ import type { DefendInsightsGetRequestQuery, } from '@kbn/elastic-assistant-common'; import type { AuthenticatedUser } from '@kbn/core-security-common'; +import { v4 as uuidv4 } from 'uuid'; -import type { AIAssistantDataClientParams } from '..'; +import type { AIAssistantDataClientParams } from '../../../ai_assistant_data_clients'; import type { EsDefendInsightSchema } from './types'; - -import { AIAssistantDataClient } from '..'; +import { AIAssistantDataClient } from '../../../ai_assistant_data_clients'; import { getDefendInsight } from './get_defend_insight'; import { queryParamsToEsQuery, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts similarity index 95% rename from x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts rename to x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts index f04c7ef505c2f..87b1c8edbd776 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts @@ -12,7 +12,7 @@ import type { UUID, } from '@kbn/elastic-assistant-common'; -import type { EsReplacementSchema } from '../conversations/types'; +import type { EsReplacementSchema } from '../../../ai_assistant_data_clients/conversations/types'; interface DefendInsightInsightEventSchema { id: string; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/index.ts index c1027b835765d..1d95ceaa27f8e 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/index.ts @@ -5,6 +5,8 @@ * 2.0. */ +import { DEFEND_INSIGHTS_ID } from '@kbn/elastic-assistant-common'; + import { getDefaultAssistantGraph, GetDefaultAssistantGraphParams, @@ -15,11 +17,19 @@ import { GetDefaultAttackDiscoveryGraphParams, getDefaultAttackDiscoveryGraph, } from '../../attack_discovery/graphs/default_attack_discovery_graph'; +import { + DefaultDefendInsightsGraph, + GetDefaultDefendInsightsGraphParams, + getDefaultDefendInsightsGraph, +} from '../../defend_insights/graphs/default_defend_insights_graph'; export type GetAssistantGraph = (params: GetDefaultAssistantGraphParams) => DefaultAssistantGraph; export type GetAttackDiscoveryGraph = ( params: GetDefaultAttackDiscoveryGraphParams ) => DefaultAttackDiscoveryGraph; +export type GetDefendInsightsGraph = ( + params: GetDefaultDefendInsightsGraphParams +) => DefaultDefendInsightsGraph; export interface AssistantGraphMetadata { getDefaultAssistantGraph: GetAssistantGraph; @@ -31,7 +41,15 @@ export interface AttackDiscoveryGraphMetadata { graphType: 'attack-discovery'; } -export type GraphMetadata = AssistantGraphMetadata | AttackDiscoveryGraphMetadata; +export interface DefendInsightsGraphMetadata { + getDefaultDefendInsightsGraph: GetDefendInsightsGraph; + graphType: typeof DEFEND_INSIGHTS_ID; +} + +export type GraphMetadata = + | AssistantGraphMetadata + | AttackDiscoveryGraphMetadata + | DefendInsightsGraphMetadata; /** * Map of the different Assistant Graphs. Useful for running evaluations. @@ -45,4 +63,8 @@ export const ASSISTANT_GRAPH_MAP: Record = { getDefaultAttackDiscoveryGraph, graphType: 'attack-discovery', }, + DefaultDefendInsightsGraph: { + getDefaultDefendInsightsGraph, + graphType: DEFEND_INSIGHTS_ID, + }, }; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts index e93e3786b123c..190cbd365a948 100755 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/plugin.ts @@ -24,7 +24,7 @@ import { AIAssistantService } from './ai_assistant_service'; import { RequestContextFactory } from './routes/request_context_factory'; import { PLUGIN_ID } from '../common/constants'; import { registerRoutes } from './routes/register_routes'; -import { appContextService } from './services/app_context'; +import { CallbackIds, appContextService } from './services/app_context'; import { createGetElserId, removeLegacyQuickPrompt } from './ai_assistant_service/helpers'; export class ElasticAssistantPlugin @@ -134,6 +134,9 @@ export class ElasticAssistantPlugin registerTools: (pluginName: string, tools: AssistantTool[]) => { return appContextService.registerTools(pluginName, tools); }, + registerCallback: (callbackId: CallbackIds, callback: Function) => { + return appContextService.registerCallback(callbackId, callback); + }, }; } diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts index 992f4d731eece..f1a600b1c314a 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insight.test.ts @@ -8,12 +8,10 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import type { AuthenticatedUser } from '@kbn/core-security-common'; - import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; -import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; - -import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import type { DefendInsightsDataClient } from '../../lib/defend_insights/persistence'; +import { transformESSearchToDefendInsights } from '../../lib/defend_insights/persistence/helpers'; import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; import { getDefendInsightRequest } from '../../__mocks__/request'; import { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts index b0e0edb251644..1ee33a836fd25 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/get_defend_insights.test.ts @@ -8,12 +8,10 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ import type { AuthenticatedUser } from '@kbn/core-security-common'; - import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; -import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; - -import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import type { DefendInsightsDataClient } from '../../lib/defend_insights/persistence'; +import { transformESSearchToDefendInsights } from '../../lib/defend_insights/persistence/helpers'; import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; import { getDefendInsightsRequest } from '../../__mocks__/request'; import { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts index aa422393a7ccc..7c5cdebc65ff2 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.test.ts @@ -7,11 +7,12 @@ /* eslint-disable @typescript-eslint/no-explicit-any */ +import type { Document } from '@langchain/core/documents'; +import type { DefendInsights } from '@kbn/elastic-assistant-common'; import moment from 'moment'; - import { ContentReferencesStore, - DEFEND_INSIGHTS_TOOL_ID, + DEFEND_INSIGHTS_ID, DefendInsightStatus, DefendInsightType, } from '@kbn/elastic-assistant-common'; @@ -36,9 +37,9 @@ describe('defend insights route helpers', () => { describe('getAssistantTool', () => { it('should return the defend-insights tool', () => { - const getRegisteredTools = jest.fn().mockReturnValue([{ id: DEFEND_INSIGHTS_TOOL_ID }]); + const getRegisteredTools = jest.fn().mockReturnValue([{ id: DEFEND_INSIGHTS_ID }]); const result = getAssistantTool(getRegisteredTools, 'pluginName'); - expect(result).toEqual({ id: DEFEND_INSIGHTS_TOOL_ID }); + expect(result).toEqual({ id: DEFEND_INSIGHTS_ID }); }); }); @@ -105,6 +106,7 @@ describe('defend insights route helpers', () => { describe('updateDefendInsights', () => { it('should update defend insights', async () => { const params = { + anonymizedEvents: [{}, {}, {}, {}, {}] as any as Document[], apiConfig: { connectorId: 'connector-id1', actionTypeId: 'action-type-id1', @@ -112,6 +114,7 @@ describe('defend insights route helpers', () => { provider: OpenAiProviderType.OpenAi, }, defendInsightId: 'insight-id1', + insights: ['insight1', 'insight2'] as any as DefendInsights, authenticatedUser: {} as any, dataClient: { getDefendInsight: jest.fn().mockResolvedValueOnce({ @@ -157,40 +160,6 @@ describe('defend insights route helpers', () => { expect.any(Object) ); }); - - it('should handle error if rawDefendInsights is null', async () => { - const params = { - apiConfig: { - connectorId: 'connector-id1', - actionTypeId: 'action-type-id1', - model: 'model', - provider: OpenAiProviderType.OpenAi, - }, - defendInsightId: 'id', - authenticatedUser: {} as any, - dataClient: { - getDefendInsight: jest.fn().mockResolvedValueOnce({ - status: DefendInsightStatus.Enum.running, - backingIndex: 'index', - generationIntervals: [], - }), - updateDefendInsight: jest.fn(), - } as any, - latestReplacements: {}, - logger: { error: jest.fn() } as any, - rawDefendInsights: null, - startTime: moment(), - telemetry: { reportEvent: jest.fn() } as any, - }; - await updateDefendInsights(params); - - expect(params.logger.error).toHaveBeenCalledTimes(1); - expect(params.telemetry.reportEvent).toHaveBeenCalledTimes(1); - expect(params.telemetry.reportEvent).toHaveBeenCalledWith( - DEFEND_INSIGHT_ERROR_EVENT.eventType, - expect.any(Object) - ); - }); }); describe('updateDefendInsightLastViewedAt', () => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts index 88ba28c06b8c7..66363db3ae7fc 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts @@ -5,8 +5,7 @@ * 2.0. */ -import moment, { Moment } from 'moment'; - +import type { Document } from '@langchain/core/documents'; import type { AnalyticsServiceSetup, AuthenticatedUser, @@ -17,44 +16,39 @@ import type { ElasticsearchClient } from '@kbn/core-elasticsearch-server'; import type { ApiConfig, ContentReferencesStore, - DefendInsight, DefendInsightGenerationInterval, + DefendInsights, DefendInsightsPostRequestBody, DefendInsightsResponse, Replacements, } from '@kbn/elastic-assistant-common'; import type { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; import type { ActionsClient } from '@kbn/actions-plugin/server'; - +import moment, { Moment } from 'moment'; import { ActionsClientLlm } from '@kbn/langchain/server'; import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith'; import { PublicMethodsOf } from '@kbn/utility-types'; import { transformError } from '@kbn/securitysolution-es-utils'; import { - DEFEND_INSIGHTS_TOOL_ID, + DEFEND_INSIGHTS_ID, DefendInsightStatus, DefendInsightType, DefendInsightsGetRequestQuery, } from '@kbn/elastic-assistant-common'; +import type { GraphState } from '../../lib/defend_insights/graphs/default_defend_insights_graph/types'; import type { GetRegisteredTools } from '../../services/app_context'; import type { AssistantTool, ElasticAssistantApiRequestHandlerContext } from '../../types'; - -import { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; +import { DefendInsightsDataClient } from '../../lib/defend_insights/persistence'; import { DEFEND_INSIGHT_ERROR_EVENT, DEFEND_INSIGHT_SUCCESS_EVENT, } from '../../lib/telemetry/event_based_telemetry'; -import { getLlmType } from '../utils'; +import { getDefaultDefendInsightsGraph } from '../../lib/defend_insights/graphs/default_defend_insights_graph'; +import { DEFEND_INSIGHTS_GRAPH_RUN_NAME } from '../../lib/defend_insights/graphs/default_defend_insights_graph/constants'; import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; - -function getDataFromJSON(defendInsightStringified: string): { - eventsContextCount: number; - insights: DefendInsight[]; -} { - const { eventsContextCount, insights } = JSON.parse(defendInsightStringified); - return { eventsContextCount, insights }; -} +import { getLlmType } from '../utils'; +import { MAX_GENERATION_ATTEMPTS, MAX_HALLUCINATION_FAILURES } from './translations'; function addGenerationInterval( generationIntervals: DefendInsightGenerationInterval[], @@ -93,7 +87,7 @@ export function getAssistantTool( pluginName: string ): AssistantTool | undefined { const assistantTools = getRegisteredTools(pluginName); - return assistantTools.find((tool) => tool.id === DEFEND_INSIGHTS_TOOL_ID); + return assistantTools.find((tool) => tool.id === DEFEND_INSIGHTS_ID); } export function getAssistantToolParams({ @@ -271,30 +265,29 @@ export async function createDefendInsight( } export async function updateDefendInsights({ + anonymizedEvents, apiConfig, defendInsightId, + insights, authenticatedUser, dataClient, latestReplacements, logger, - rawDefendInsights, startTime, telemetry, }: { + anonymizedEvents: Document[]; apiConfig: ApiConfig; defendInsightId: string; + insights: DefendInsights | null; authenticatedUser: AuthenticatedUser; dataClient: DefendInsightsDataClient; latestReplacements: Replacements; logger: Logger; - rawDefendInsights: string | null; startTime: Moment; telemetry: AnalyticsServiceSetup; }) { try { - if (rawDefendInsights == null) { - throw new Error('tool returned no Defend insights'); - } const currentInsight = await dataClient.getDefendInsight({ id: defendInsightId, authenticatedUser, @@ -304,12 +297,12 @@ export async function updateDefendInsights({ } const endTime = moment(); const durationMs = endTime.diff(startTime); - const { eventsContextCount, insights } = getDataFromJSON(rawDefendInsights); + const eventsContextCount = anonymizedEvents.length; const updateProps = { eventsContextCount, - insights, + insights: insights ?? undefined, status: DefendInsightStatus.Enum.succeeded, - ...(!eventsContextCount || !insights.length + ...(!eventsContextCount || !insights ? {} : { generationIntervals: addGenerationInterval(currentInsight.generationIntervals, { @@ -329,7 +322,7 @@ export async function updateDefendInsights({ telemetry.reportEvent(DEFEND_INSIGHT_SUCCESS_EVENT.eventType, { actionTypeId: apiConfig.actionTypeId, eventsContextCount: updateProps.eventsContextCount, - insightsGenerated: updateProps.insights.length, + insightsGenerated: updateProps.insights?.length ?? 0, durationMs, model: apiConfig.model, provider: apiConfig.provider, @@ -390,3 +383,207 @@ export async function updateDefendInsightLastViewedAt({ await updateDefendInsightsLastViewedAt({ params: { ids: [id] }, authenticatedUser, dataClient }) )[0]; } + +export const invokeDefendInsightsGraph = async ({ + insightType, + endpointIds, + actionsClient, + anonymizationFields, + apiConfig, + connectorTimeout, + esClient, + langSmithProject, + langSmithApiKey, + latestReplacements, + logger, + onNewReplacements, + size, + start, + end, +}: { + insightType: DefendInsightType; + endpointIds: string[]; + actionsClient: PublicMethodsOf; + anonymizationFields: AnonymizationFieldResponse[]; + apiConfig: ApiConfig; + connectorTimeout: number; + esClient: ElasticsearchClient; + langSmithProject?: string; + langSmithApiKey?: string; + latestReplacements: Replacements; + logger: Logger; + onNewReplacements: (newReplacements: Replacements) => void; + size?: number; + start?: string; + end?: string; +}): Promise<{ + anonymizedEvents: Document[]; + insights: DefendInsights | null; +}> => { + const llmType = getLlmType(apiConfig.actionTypeId); + const model = apiConfig.model; + const tags = [DEFEND_INSIGHTS_ID, llmType, model].flatMap((tag) => tag ?? []); + + const traceOptions = { + projectName: langSmithProject, + tracers: [ + ...getLangSmithTracer({ + apiKey: langSmithApiKey, + projectName: langSmithProject, + logger, + }), + ], + }; + + const llm = new ActionsClientLlm({ + actionsClient, + connectorId: apiConfig.connectorId, + llmType, + logger, + temperature: 0, + timeout: connectorTimeout, + traceOptions, + }); + + if (llm == null) { + throw new Error('LLM is required for Defend insights'); + } + + const graph = getDefaultDefendInsightsGraph({ + insightType, + endpointIds, + anonymizationFields, + esClient, + llm, + logger, + onNewReplacements, + replacements: latestReplacements, + size, + start, + end, + }); + + logger?.debug(() => 'invokeDefendInsightsGraph: invoking the Defend insights graph'); + + const result: GraphState = (await graph.invoke( + {}, + { + callbacks: [...(traceOptions?.tracers ?? [])], + runName: DEFEND_INSIGHTS_GRAPH_RUN_NAME, + tags, + } + )) as GraphState; + const { + insights, + anonymizedEvents, + errors, + generationAttempts, + hallucinationFailures, + maxGenerationAttempts, + maxHallucinationFailures, + } = result; + + throwIfErrorCountsExceeded({ + errors, + generationAttempts, + hallucinationFailures, + logger, + maxGenerationAttempts, + maxHallucinationFailures, + }); + + return { anonymizedEvents, insights }; +}; + +export const handleGraphError = async ({ + apiConfig, + defendInsightId, + authenticatedUser, + dataClient, + err, + latestReplacements, + logger, + telemetry, +}: { + apiConfig: ApiConfig; + defendInsightId: string; + authenticatedUser: AuthenticatedUser; + dataClient: DefendInsightsDataClient; + err: Error; + latestReplacements: Replacements; + logger: Logger; + telemetry: AnalyticsServiceSetup; +}) => { + try { + logger.error(err); + const error = transformError(err); + const currentInsight = await dataClient.getDefendInsight({ + id: defendInsightId, + authenticatedUser, + }); + + if (currentInsight === null || currentInsight?.status === 'canceled') { + return; + } + + await dataClient.updateDefendInsight({ + defendInsightUpdateProps: { + insights: [], + status: DefendInsightStatus.Enum.failed, + id: defendInsightId, + replacements: latestReplacements, + backingIndex: currentInsight.backingIndex, + failureReason: error.message, + }, + authenticatedUser, + }); + telemetry.reportEvent(DEFEND_INSIGHT_ERROR_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + errorMessage: error.message, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } catch (updateErr) { + const updateError = transformError(updateErr); + telemetry.reportEvent(DEFEND_INSIGHT_ERROR_EVENT.eventType, { + actionTypeId: apiConfig.actionTypeId, + errorMessage: updateError.message, + model: apiConfig.model, + provider: apiConfig.provider, + }); + } +}; + +export const throwIfErrorCountsExceeded = ({ + errors, + generationAttempts, + hallucinationFailures, + logger, + maxGenerationAttempts, + maxHallucinationFailures, +}: { + errors: string[]; + generationAttempts: number; + hallucinationFailures: number; + logger?: Logger; + maxGenerationAttempts: number; + maxHallucinationFailures: number; +}): void => { + if (hallucinationFailures >= maxHallucinationFailures) { + const hallucinationFailuresError = `${MAX_HALLUCINATION_FAILURES( + hallucinationFailures + )}\n${errors.join(',\n')}`; + + logger?.error(hallucinationFailuresError); + throw new Error(hallucinationFailuresError); + } + + if (generationAttempts >= maxGenerationAttempts) { + const generationAttemptsError = `${MAX_GENERATION_ATTEMPTS(generationAttempts)}\n${errors.join( + ',\n' + )}`; + + logger?.error(generationAttemptsError); + throw new Error(generationAttemptsError); + } +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts index d602aa4d962bf..0b8bd879ef514 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts @@ -9,31 +9,27 @@ import type { AuthenticatedUser } from '@kbn/core-security-common'; import type { DefendInsightsPostRequestBody } from '@kbn/elastic-assistant-common'; - -import { getPrompt } from '@kbn/security-ai-prompts'; import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks'; import { actionsMock } from '@kbn/actions-plugin/server/mocks'; import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants'; import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common'; +import { licensingMock } from '@kbn/licensing-plugin/public/mocks'; -import type { DefendInsightsDataClient } from '../../ai_assistant_data_clients/defend_insights'; - +import type { DefendInsightsDataClient } from '../../lib/defend_insights/persistence'; import { serverMock } from '../../__mocks__/server'; import { ElasticAssistantRequestHandlerContextMock, requestContextMock, } from '../../__mocks__/request_context'; -import { transformESSearchToDefendInsights } from '../../ai_assistant_data_clients/defend_insights/helpers'; +import { transformESSearchToDefendInsights } from '../../lib/defend_insights/persistence/helpers'; import { getDefendInsightsSearchEsMock } from '../../__mocks__/defend_insights_schema.mock'; import { postDefendInsightsRequest } from '../../__mocks__/request'; -import { getAssistantTool, createDefendInsight, isDefendInsightsEnabled } from './helpers'; +import { createDefendInsight, isDefendInsightsEnabled, invokeDefendInsightsGraph } from './helpers'; import { postDefendInsightsRoute } from './post_defend_insights'; -import { licensingMock } from '@kbn/licensing-plugin/public/mocks'; jest.mock('@kbn/security-ai-prompts'); jest.mock('./helpers'); -const getPromptMock = getPrompt as jest.Mock; describe('postDefendInsightsRoute', () => { let server: ReturnType; let context: ElasticAssistantRequestHandlerContextMock; @@ -83,7 +79,6 @@ describe('postDefendInsightsRoute', () => { langSmithApiKey: 'langSmithApiKey', }; } - const getTool = jest.fn(); beforeEach(() => { const tools = requestContextMock.createTools(); @@ -98,13 +93,15 @@ describe('postDefendInsightsRoute', () => { mockDataClient = getDefaultDataClient(); mockApiConfig = getDefaultApiConfig(); mockRequestBody = getDefaultRequestBody(); - getPromptMock.mockResolvedValue('prompt'); - (getAssistantTool as jest.Mock).mockReturnValue({ getTool, name: 'test-tool' }); (createDefendInsight as jest.Mock).mockResolvedValue({ currentInsight: mockCurrentInsight, defendInsightId: mockCurrentInsight.id, }); (isDefendInsightsEnabled as jest.Mock).mockResolvedValue(true); + (invokeDefendInsightsGraph as jest.Mock).mockResolvedValue({ + anonymizedEvents: [], + insights: [mockCurrentInsight], + }); context.elasticAssistant.getCurrentUser.mockResolvedValue(mockUser); context.elasticAssistant.getDefendInsightsDataClient.mockResolvedValue(mockDataClient); @@ -168,15 +165,6 @@ describe('postDefendInsightsRoute', () => { }); }); - it('should handle assistantTool null response', async () => { - (getAssistantTool as jest.Mock).mockReturnValueOnce(null); - const response = await server.inject( - postDefendInsightsRequest(mockRequestBody), - requestContextMock.convertContext(context) - ); - expect(response.status).toEqual(404); - }); - it('should 404 if feature flag disabled', async () => { (isDefendInsightsEnabled as jest.Mock).mockReturnValueOnce(false); const response = await server.inject( @@ -201,23 +189,4 @@ describe('postDefendInsightsRoute', () => { status_code: 500, }); }); - - it('should call getPrompt for tool description', async () => { - await server.inject( - postDefendInsightsRequest(mockRequestBody), - requestContextMock.convertContext(context) - ); - expect(getPromptMock).toHaveBeenCalledWith( - expect.objectContaining({ - connectorId: 'connector-id', - promptId: 'test-tool', - promptGroupId: 'security-tools', - }) - ); - expect(getTool).toHaveBeenCalledWith( - expect.objectContaining({ - description: 'prompt', - }) - ); - }); }); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts index 46fb5a4301a28..90bf6c267c141 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts @@ -5,10 +5,8 @@ * 2.0. */ -import moment from 'moment/moment'; - import type { IKibanaResponse } from '@kbn/core/server'; - +import moment from 'moment/moment'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; import { DEFEND_INSIGHTS, @@ -20,19 +18,16 @@ import { import { transformError } from '@kbn/securitysolution-es-utils'; import { IRouter, Logger } from '@kbn/core/server'; -import { getPrompt } from '@kbn/security-ai-prompts'; -import { localToolPrompts, promptGroupId } from '../../lib/prompt/tool_prompts'; import { buildResponse } from '../../lib/build_response'; import { ElasticAssistantRequestHandlerContext } from '../../types'; -import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers'; import { - getAssistantTool, - getAssistantToolParams, - handleToolError, createDefendInsight, updateDefendInsights, isDefendInsightsEnabled, + invokeDefendInsightsGraph, + handleGraphError, } from './helpers'; +import { CallbackIds, appContextService } from '../../services/app_context'; const ROUTE_HANDLER_TIMEOUT = 10 * 60 * 1000; // 10 * 60 seconds = 10 minutes const LANG_CHAIN_TIMEOUT = ROUTE_HANDLER_TIMEOUT - 10_000; // 9 minutes 50 seconds @@ -73,8 +68,6 @@ export const postDefendInsightsRoute = (router: IRouter + }) + .then(({ anonymizedEvents, insights }) => updateDefendInsights({ + anonymizedEvents, apiConfig, defendInsightId, + insights, authenticatedUser, dataClient, latestReplacements, logger, - rawDefendInsights, startTime, telemetry, - }) + }).then(() => insights) + ) + .then((insights) => + appContextService + .getRegisteredCallbacks(CallbackIds.DefendInsightsPostCreate) + .map((cb) => cb(insights, request)) ) .catch((err) => - handleToolError({ + handleGraphError({ apiConfig, defendInsightId, authenticatedUser, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/translations.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/translations.ts new file mode 100644 index 0000000000000..467b966365a73 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/translations.ts @@ -0,0 +1,28 @@ +/* + * 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 { i18n } from '@kbn/i18n'; + +export const MAX_HALLUCINATION_FAILURES = (hallucinationFailures: number) => + i18n.translate( + 'xpack.elasticAssistantPlugin.defendInsights.defaultDefendInsightsGraph.nodes.retriever.helpers.throwIfErrorCountsExceeded.maxHallucinationFailuresErrorMessage', + { + defaultMessage: + 'Maximum hallucination failures ({hallucinationFailures}) reached. Try sending fewer events to this model.', + values: { hallucinationFailures }, + } + ); + +export const MAX_GENERATION_ATTEMPTS = (generationAttempts: number) => + i18n.translate( + 'xpack.elasticAssistantPlugin.defendInsights.defaultDefendInsightsGraph.nodes.retriever.helpers.throwIfErrorCountsExceeded.maxGenerationAttemptsErrorMessage', + { + defaultMessage: + 'Maximum generation attempts ({generationAttempts}) reached. Try sending fewer events to this model.', + values: { generationAttempts }, + } + ); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_graphs_from_names/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_graphs_from_names/index.ts index c0320c9ff6adf..21b9951c4623b 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_graphs_from_names/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/get_graphs_from_names/index.ts @@ -21,9 +21,12 @@ export const getGraphsFromNames = (graphNames: string[]): GetGraphsFromNamesResu (acc, graphName) => { const graph = ASSISTANT_GRAPH_MAP[graphName]; if (graph != null) { - return graph.graphType === 'assistant' - ? { ...acc, assistantGraphs: [...acc.assistantGraphs, graph] } - : { ...acc, attackDiscoveryGraphs: [...acc.attackDiscoveryGraphs, graph] }; + switch (graph.graphType) { + case 'assistant': + return { ...acc, assistantGraphs: [...acc.assistantGraphs, graph] }; + case 'attack-discovery': + return { ...acc, attackDiscoveryGraphs: [...acc.attackDiscoveryGraphs, graph] }; + } } return acc; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts index 17b7f0729deca..2c9d4b2dacfe5 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts @@ -21,7 +21,7 @@ import { Message, Replacements, replaceAnonymizedValuesWithOriginalValues, - DEFEND_INSIGHTS_TOOL_ID, + DEFEND_INSIGHTS_ID, ContentReferencesStore, ContentReferences, MessageMetadata, @@ -290,7 +290,7 @@ export const langChainExecute = async ({ }); const assistantContext = context.elasticAssistant; // We don't (yet) support invoking these tools interactively - const unsupportedTools = new Set(['attack-discovery', DEFEND_INSIGHTS_TOOL_ID]); + const unsupportedTools = new Set(['attack-discovery', DEFEND_INSIGHTS_ID]); const assistantTools = assistantContext .getRegisteredTools(pluginName) .filter((tool) => !unsupportedTools.has(tool.id)); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts index 9708602db7f29..3a50fbdd15e7d 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/services/app_context.ts @@ -17,6 +17,10 @@ export type GetRegisteredFeatures = (pluginName: string) => AssistantFeatures; export interface ElasticAssistantAppContext { logger: Logger; } +export enum CallbackIds { + DefendInsightsPostCreate = 'defend-insights:post-create', +} +export type RegisteredCallbacks = Map; /** * Service for managing context specific to the Elastic Assistant @@ -27,6 +31,7 @@ class AppContextService { private logger: Logger | undefined; private registeredTools: RegisteredToolsStorage = new Map>(); private registeredFeatures: RegisteredFeaturesStorage = new Map(); + private registeredCallbacks: RegisteredCallbacks = new Map(); public start(appContext: ElasticAssistantAppContext) { this.logger = appContext.logger; @@ -116,6 +121,24 @@ class AppContextService { return features; } + + /** + * Register a callback to a callbackId + * @param callbackId + * @param callback + */ + public registerCallback(callbackId: CallbackIds, callback: Function) { + const callbacks = this.registeredCallbacks.get(callbackId) ?? []; + this.registeredCallbacks.set(callbackId, [...callbacks, callback]); + } + + /** + * Get all registered callbacks for a callbackId + * @param callbackId + */ + public getRegisteredCallbacks(callbackId: CallbackIds): Function[] { + return this.registeredCallbacks.get(callbackId) ?? []; + } } export const appContextService = new AppContextService(); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts index 1dfe451a915f8..507164a43e394 100755 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts @@ -57,11 +57,13 @@ import { GetAIAssistantConversationsDataClientParams, } from './ai_assistant_data_clients/conversations'; import type { GetRegisteredFeatures, GetRegisteredTools } from './services/app_context'; +import { CallbackIds } from './services/app_context'; import { AIAssistantDataClient } from './ai_assistant_data_clients'; import { AIAssistantKnowledgeBaseDataClient } from './ai_assistant_data_clients/knowledge_base'; -import type { DefendInsightsDataClient } from './ai_assistant_data_clients/defend_insights'; +import type { DefendInsightsDataClient } from './lib/defend_insights/persistence'; export const PLUGIN_ID = 'elasticAssistant' as const; +export { CallbackIds }; /** The plugin setup interface */ export interface ElasticAssistantPluginSetup { @@ -108,6 +110,12 @@ export interface ElasticAssistantPluginStart { * @param pluginName Name of the plugin to get the tools for */ getRegisteredTools: GetRegisteredTools; + /** + * Register a callback to be used by the elastic assistant. + * @param callbackId + * @param callback + */ + registerCallback: (callbackId: CallbackIds, callback: Function) => void; } export interface ElasticAssistantPluginSetupDependencies { diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts deleted file mode 100644 index e35a3fee54866..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.test.ts +++ /dev/null @@ -1,183 +0,0 @@ -/* - * 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 type { ElasticsearchClient } from '@kbn/core/server'; - -import { DefendInsightType, transformRawData } from '@kbn/elastic-assistant-common'; - -import { InvalidDefendInsightTypeError } from '../errors'; -import { getFileEventsQuery } from './get_file_events_query'; -import { getAnonymizedEvents } from '.'; -import type { SearchResponse } from '@elastic/elasticsearch/lib/api/types'; - -jest.mock('@kbn/elastic-assistant-common', () => { - const originalModule = jest.requireActual('@kbn/elastic-assistant-common'); - return { - ...originalModule, - transformRawData: jest.fn(), - }; -}); - -jest.mock('./get_file_events_query', () => ({ - getFileEventsQuery: jest.fn(), -})); - -describe('getAnonymizedEvents', () => { - let mockEsClient: jest.Mocked; - - const mockAggregations = { - unique_process_executable: { - buckets: [ - { - key: 'process1', - doc_count: 10, - latest_event: { - hits: { - hits: [ - { - _id: 'event1', - _source: { - agent: { id: 'agent1' }, - process: { executable: 'process1' }, - }, - }, - ], - }, - }, - }, - { - key: 'process2', - doc_count: 5, - latest_event: { - hits: { - hits: [ - { - _id: 'event2', - _source: { - agent: { id: 'agent2' }, - process: { executable: 'process2' }, - }, - }, - ], - }, - }, - }, - ], - }, - }; - - beforeEach(() => { - (getFileEventsQuery as jest.Mock).mockReturnValue({ index: 'test-index', body: {} }); - (transformRawData as jest.Mock).mockImplementation( - ({ rawData }) => `anonymized_${Object.values(rawData)[0]}` - ); - mockEsClient = { - search: jest.fn().mockResolvedValue({ - took: 1, - timed_out: false, - _shards: { - total: 1, - successful: 1, - skipped: 0, - failed: 0, - }, - aggregations: mockAggregations, - }), - } as unknown as jest.Mocked; - }); - - afterEach(() => { - jest.clearAllMocks(); - }); - - it('should return anonymized events successfully from aggregations', async () => { - const result = await getAnonymizedEvents({ - endpointIds: ['endpoint1'], - type: DefendInsightType.Enum.incompatible_antivirus, - esClient: mockEsClient, - }); - - expect(result).toEqual(['anonymized_event1', 'anonymized_event2']); - expect(getFileEventsQuery).toHaveBeenCalledWith({ endpointIds: ['endpoint1'] }); - expect(mockEsClient.search).toHaveBeenCalledWith({ index: 'test-index', body: {} }); - expect(transformRawData).toHaveBeenCalledTimes(2); - expect(transformRawData).toHaveBeenCalledWith( - expect.objectContaining({ - rawData: expect.objectContaining({ - _id: ['event1'], - }), - }) - ); - }); - - it('should map aggregation response correctly into fileEvents structure', async () => { - await getAnonymizedEvents({ - endpointIds: ['endpoint1'], - type: DefendInsightType.Enum.incompatible_antivirus, - esClient: mockEsClient, - }); - - expect(mockEsClient.search).toHaveBeenCalledWith({ index: 'test-index', body: {} }); - - expect(transformRawData).toHaveBeenCalledWith( - expect.objectContaining({ - rawData: { - _id: ['event1'], - 'agent.id': ['agent1'], - 'process.executable': ['process1'], - }, - }) - ); - - expect(transformRawData).toHaveBeenCalledWith( - expect.objectContaining({ - rawData: { - _id: ['event2'], - 'agent.id': ['agent2'], - 'process.executable': ['process2'], - }, - }) - ); - }); - - it('should throw InvalidDefendInsightTypeError for invalid type', async () => { - await expect( - getAnonymizedEvents({ - endpointIds: ['endpoint1'], - type: 'invalid_type' as DefendInsightType, - esClient: mockEsClient, - }) - ).rejects.toThrow(InvalidDefendInsightTypeError); - }); - - it('should handle empty aggregation response gracefully', async () => { - mockEsClient.search.mockResolvedValueOnce({ - took: 1, - timed_out: false, - _shards: { - total: 1, - successful: 1, - skipped: 0, - failed: 0, - }, - aggregations: { - unique_process_executable: { - buckets: [], - }, - }, - } as unknown as SearchResponse); - - const result = await getAnonymizedEvents({ - endpointIds: ['endpoint1'], - type: DefendInsightType.Enum.incompatible_antivirus, - esClient: mockEsClient, - }); - - expect(result).toEqual([]); - expect(transformRawData).not.toHaveBeenCalled(); - }); -}); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts deleted file mode 100644 index 5ef5aaeedf364..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.test.ts +++ /dev/null @@ -1,69 +0,0 @@ -/* - * 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 { DynamicTool } from '@langchain/core/tools'; - -import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; -import { DEFEND_INSIGHTS_TOOL_ID, DefendInsightType } from '@kbn/elastic-assistant-common'; -import { elasticsearchServiceMock } from '@kbn/core/server/mocks'; - -import type { DefendInsightsToolParams } from '.'; - -import { APP_UI_ID } from '../../../../common'; -import { DEFEND_INSIGHTS_TOOL, DEFEND_INSIGHTS_TOOL_DESCRIPTION } from '.'; - -jest.mock('@kbn/elastic-assistant-plugin/server/lib/langchain/helpers', () => ({ - requestHasRequiredAnonymizationParams: jest.fn(), -})); - -describe('DEFEND_INSIGHTS_TOOL', () => { - const mockLLM = {}; - const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); - const mockRequest = {}; - const mockParams: DefendInsightsToolParams = { - endpointIds: ['endpoint1'], - insightType: DefendInsightType.Enum.incompatible_antivirus, - anonymizationFields: [], - esClient: mockEsClient, - langChainTimeout: 1000, - llm: mockLLM, - onNewReplacements: jest.fn(), - replacements: {}, - request: mockRequest, - } as unknown as DefendInsightsToolParams; - - afterEach(() => { - jest.clearAllMocks(); - }); - - it('should have correct properties', () => { - expect(DEFEND_INSIGHTS_TOOL.id).toBe(DEFEND_INSIGHTS_TOOL_ID); - expect(DEFEND_INSIGHTS_TOOL.name).toBe('defendInsightsTool'); - expect(DEFEND_INSIGHTS_TOOL.description).toBe(DEFEND_INSIGHTS_TOOL_DESCRIPTION); - expect(DEFEND_INSIGHTS_TOOL.sourceRegister).toBe(APP_UI_ID); - }); - - it('should return tool if supported', () => { - (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(true); - const tool = DEFEND_INSIGHTS_TOOL.getTool(mockParams); - expect(tool).toBeInstanceOf(DynamicTool); - }); - - it('should return null if not request missing anonymization params', () => { - (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(false); - const tool = DEFEND_INSIGHTS_TOOL.getTool(mockParams); - expect(tool).toBeNull(); - }); - - it('should return null if LLM is not provided', () => { - (requestHasRequiredAnonymizationParams as jest.Mock).mockReturnValue(true); - const paramsWithoutLLM = { ...mockParams, llm: undefined }; - const tool = DEFEND_INSIGHTS_TOOL.getTool(paramsWithoutLLM) as DynamicTool; - - expect(tool).toBeNull(); - }); -}); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.ts deleted file mode 100644 index 0851642388550..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.ts +++ /dev/null @@ -1,127 +0,0 @@ -/* - * 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 { PromptTemplate } from '@langchain/core/prompts'; -import { DynamicTool } from '@langchain/core/tools'; -import { LLMChain } from 'langchain/chains'; -import { OutputFixingParser } from 'langchain/output_parsers'; - -import type { AssistantTool, AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; -import type { - DefendInsight, - DefendInsightType, - DefendInsightsPostRequestBody, -} from '@kbn/elastic-assistant-common'; -import type { KibanaRequest } from '@kbn/core/server'; - -import { requestHasRequiredAnonymizationParams } from '@kbn/elastic-assistant-plugin/server/lib/langchain/helpers'; -import { DEFEND_INSIGHTS_TOOL_ID } from '@kbn/elastic-assistant-common'; - -import { APP_UI_ID } from '../../../../common'; -import { securityWorkflowInsightsService } from '../../../endpoint/services'; -import { getAnonymizedEvents } from './get_events'; -import { getDefendInsightsOutputParser } from './output_parsers'; -import { getDefendInsightsPrompt } from './prompts'; - -export const DEFEND_INSIGHTS_TOOL_DESCRIPTION = 'Call this for Elastic Defend insights.'; - -export interface DefendInsightsToolParams extends AssistantToolParams { - endpointIds: string[]; - insightType: DefendInsightType; - request: KibanaRequest; -} - -/** - * Returns a tool for generating Elastic Defend configuration insights - */ -export const DEFEND_INSIGHTS_TOOL: AssistantTool = Object.freeze({ - id: DEFEND_INSIGHTS_TOOL_ID, - name: 'defendInsightsTool', - // note: this description is overwritten when `getTool` is called - // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts - // local definitions can be overwritten by security-ai-prompt integration definitions - description: DEFEND_INSIGHTS_TOOL_DESCRIPTION, - sourceRegister: APP_UI_ID, - - isSupported: (params: AssistantToolParams): boolean => { - const { llm, request } = params; - - return requestHasRequiredAnonymizationParams(request) && llm != null; - }, - - getTool(params: AssistantToolParams): DynamicTool | null { - if (!this.isSupported(params)) return null; - - const { - endpointIds, - insightType, - anonymizationFields, - esClient, - langChainTimeout, - llm, - onNewReplacements, - replacements, - request, - } = params as DefendInsightsToolParams; - - return new DynamicTool({ - name: 'DefendInsightsTool', - description: params.description || DEFEND_INSIGHTS_TOOL_DESCRIPTION, - func: async () => { - if (llm == null) { - throw new Error('LLM is required for Defend Insights'); - } - - const anonymizedEvents = await getAnonymizedEvents({ - endpointIds, - type: insightType, - anonymizationFields, - esClient, - onNewReplacements, - replacements, - }); - - const eventsContextCount = anonymizedEvents.length; - if (eventsContextCount === 0) { - return JSON.stringify({ eventsContextCount, insights: [] }, null, 2); - } - - const outputParser = getDefendInsightsOutputParser({ type: insightType }); - const outputFixingParser = OutputFixingParser.fromLLM(llm, outputParser); - - const prompt = new PromptTemplate({ - template: `Answer the user's question as best you can:\n{format_instructions}\n{query}`, - inputVariables: ['query'], - partialVariables: { - format_instructions: outputFixingParser.getFormatInstructions(), - }, - }); - - const answerFormattingChain = new LLMChain({ - llm, - prompt, - outputKey: 'records', - outputParser: outputFixingParser, - }); - - const result = await answerFormattingChain.call({ - query: getDefendInsightsPrompt({ - type: insightType, - events: anonymizedEvents, - }), - timeout: langChainTimeout, - }); - const insights: DefendInsight[] = result.records; - - await securityWorkflowInsightsService.createFromDefendInsights(insights, request); - - return JSON.stringify({ eventsContextCount, insights }, null, 2); - }, - tags: [DEFEND_INSIGHTS_TOOL_ID], - }); - }, -}); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts deleted file mode 100644 index 516de86a30975..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/prompts/incompatible_antivirus.ts +++ /dev/null @@ -1,16 +0,0 @@ -/* - * 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. - */ - -export function getIncompatibleAntivirusPrompt({ events }: { events: string[] }): string { - return `You are an Elastic Security user tasked with analyzing file events from Elastic Security to identify antivirus processes. Only focus on detecting antivirus processes. Ignore processes that belong to Elastic Agent or Elastic Defend, that are not antivirus processes, or are typical processes built into the operating system. Accuracy is of the utmost importance, try to minimize false positives. Group the processes by the antivirus program, keeping track of the agent.id and _id associated to each of the individual events as endpointId and eventId respectively. If there are no events, ignore the group field. Escape backslashes to respect JSON validation. New lines must always be escaped with double backslashes, i.e. \\\\n to ensure valid JSON. Only return JSON output, as described above. Do not add any additional text to describe your output. - - Use context from the following process events to provide insights: - """ - ${events.join('\n\n')} - """ - `; -} diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts index dd2aa8e54ebdf..a86b65bedf989 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/index.ts @@ -9,7 +9,6 @@ import { PRODUCT_DOCUMENTATION_TOOL } from './product_docs/product_documentation import { NL_TO_ESQL_TOOL } from './esql/nl_to_esql_tool'; import { ALERT_COUNTS_TOOL } from './alert_counts/alert_counts_tool'; import { OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL } from './open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool'; -import { DEFEND_INSIGHTS_TOOL } from './defend_insights'; import { KNOWLEDGE_BASE_RETRIEVAL_TOOL } from './knowledge_base/knowledge_base_retrieval_tool'; import { KNOWLEDGE_BASE_WRITE_TOOL } from './knowledge_base/knowledge_base_write_tool'; import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs/security_labs_tool'; @@ -18,7 +17,6 @@ import { SECURITY_LABS_KNOWLEDGE_BASE_TOOL } from './security_labs/security_labs // x-pack/solutions/security/plugins/elastic_assistant/server/lib/telemetry/event_based_telemetry.ts export const assistantTools = [ ALERT_COUNTS_TOOL, - DEFEND_INSIGHTS_TOOL, KNOWLEDGE_BASE_RETRIEVAL_TOOL, KNOWLEDGE_BASE_WRITE_TOOL, NL_TO_ESQL_TOOL, diff --git a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/builders/index.ts b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/builders/index.ts index a19e800629598..7c3d5bbbde02a 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/builders/index.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/builders/index.ts @@ -6,14 +6,12 @@ */ import type { ElasticsearchClient, KibanaRequest } from '@kbn/core/server'; - import type { DefendInsight, DefendInsightsPostRequestBody } from '@kbn/elastic-assistant-common'; - import { DefendInsightType } from '@kbn/elastic-assistant-common'; +import { InvalidDefendInsightTypeError } from '@kbn/elastic-assistant-plugin/server/lib/defend_insights/errors'; import type { SecurityWorkflowInsight } from '../../../../../common/endpoint/types/workflow_insights'; -import { InvalidDefendInsightTypeError } from '../../../../assistant/tools/defend_insights/errors'; import type { EndpointMetadataService } from '../../metadata'; import { buildIncompatibleAntivirusWorkflowInsights } from './incompatible_antivirus'; diff --git a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.test.ts b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.test.ts index f4ecaaf875061..927da3aa5ab3f 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.test.ts @@ -193,7 +193,10 @@ describe('SecurityWorkflowInsightsService', () => { expect(createDatastreamMock).toHaveBeenCalledTimes(1); expect(createDatastreamMock).toHaveBeenCalledWith(kibanaPackageJson.version); - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); expect(createPipelineMock).toHaveBeenCalledTimes(1); expect(createPipelineMock).toHaveBeenCalledWith(esClient); @@ -220,7 +223,10 @@ describe('SecurityWorkflowInsightsService', () => { throw new Error('test error'); }); - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); expect(logger.warn).toHaveBeenCalledTimes(2); expect(logger.warn).toHaveBeenNthCalledWith(1, expect.stringContaining('test error')); @@ -262,7 +268,10 @@ describe('SecurityWorkflowInsightsService', () => { _version: 1, }; jest.spyOn(esClient, 'index').mockResolvedValue(esClientIndexResp); - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const result = await securityWorkflowInsightsService.createFromDefendInsights( defendInsights, request @@ -282,7 +291,10 @@ describe('SecurityWorkflowInsightsService', () => { describe('create', () => { it('should index the doc correctly', async () => { - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const insight = getDefaultInsight(); await securityWorkflowInsightsService.create(insight); @@ -300,7 +312,10 @@ describe('SecurityWorkflowInsightsService', () => { }); it('should not index the doc if remediation exists', async () => { - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const insight = getDefaultInsight(); const remediationExistsMock = checkIfRemediationExists as jest.Mock; @@ -325,7 +340,10 @@ describe('SecurityWorkflowInsightsService', () => { const updateSpy = jest .spyOn(securityWorkflowInsightsService, 'update') .mockResolvedValueOnce({} as UpdateResponse); - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const insight = getDefaultInsight(); await securityWorkflowInsightsService.create(insight); @@ -349,7 +367,10 @@ describe('SecurityWorkflowInsightsService', () => { describe('update', () => { it('should update the doc correctly', async () => { - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const insightId = 'some-insight-id'; const insight = getDefaultInsight(); const indexName = 'backing-index-name'; @@ -370,7 +391,10 @@ describe('SecurityWorkflowInsightsService', () => { describe('fetch', () => { it('should fetch the docs with the correct params', async () => { - await securityWorkflowInsightsService.start({ esClient }); + await securityWorkflowInsightsService.start({ + esClient, + registerDefendInsightsCallback: jest.fn(), + }); const searchParams: SearchParams = { size: 50, from: 50, diff --git a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.ts b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.ts index e3559df9c7e0f..ae45b249e2821 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.ts @@ -5,8 +5,6 @@ * 2.0. */ -import { ReplaySubject, firstValueFrom, combineLatest } from 'rxjs'; - import type { SearchHit, UpdateResponse, @@ -15,13 +13,14 @@ import type { import type { ElasticsearchClient, KibanaRequest, Logger } from '@kbn/core/server'; import type { DataStreamSpacesAdapter } from '@kbn/data-stream-adapter'; import type { DefendInsight, DefendInsightsPostRequestBody } from '@kbn/elastic-assistant-common'; +import { ReplaySubject, firstValueFrom, combineLatest } from 'rxjs'; +import { CallbackIds } from '@kbn/elastic-assistant-plugin/server/types'; import type { SearchParams, SecurityWorkflowInsight, } from '../../../../common/endpoint/types/workflow_insights'; import type { EndpointAppContextService } from '../../endpoint_app_context_services'; - import { SecurityWorkflowInsightsFailedInitialized } from './errors'; import { buildEsQueryParams, @@ -45,6 +44,7 @@ interface SetupInterface { interface StartInterface { esClient: ElasticsearchClient; + registerDefendInsightsCallback: (callbackId: CallbackIds, callback: Function) => void; } class SecurityWorkflowInsightsService { @@ -83,7 +83,7 @@ class SecurityWorkflowInsightsService { this.setup$.next(); } - public async start({ esClient }: StartInterface) { + public async start({ esClient, registerDefendInsightsCallback }: StartInterface) { if (!this.isFeatureEnabled) { return; } @@ -92,6 +92,7 @@ class SecurityWorkflowInsightsService { await firstValueFrom(this.setup$); try { + this.registerDefendInsightsCallbacks(registerDefendInsightsCallback); await createPipeline(esClient); await this.ds?.install({ logger: this.logger, @@ -129,6 +130,10 @@ class SecurityWorkflowInsightsService { defendInsights: DefendInsight[], request: KibanaRequest ): Promise>> { + if (!defendInsights || !defendInsights.length) { + return []; + } + await this.isInitialized; const workflowInsights = await buildWorkflowInsights({ @@ -245,6 +250,15 @@ class SecurityWorkflowInsightsService { return this._endpointContext; } + + private registerDefendInsightsCallbacks( + registerCallback: (callbackId: CallbackIds, callback: Function) => void + ) { + registerCallback( + CallbackIds.DefendInsightsPostCreate, + this.createFromDefendInsights.bind(this) + ); + } } export const securityWorkflowInsightsService = new SecurityWorkflowInsightsService(); diff --git a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts index c99b30401c45b..86945dca3e351 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts @@ -729,6 +729,7 @@ export class Plugin implements ISecuritySolutionPlugin { securityWorkflowInsightsService .start({ esClient: core.elasticsearch.client.asInternalUser, + registerDefendInsightsCallback: plugins.elasticAssistant.registerCallback, }) .catch(() => {}); From d7445380b366ed94e769199e878fec74b3660d23 Mon Sep 17 00:00:00 2001 From: Kevin Delemme Date: Tue, 4 Mar 2025 14:35:28 -0500 Subject: [PATCH 7/9] chore(slo): replace error log level with debug (#212975) ## Summary Resolves https://github.com/elastic/kibana/issues/212972 This PR replaces the info and error log levels with debug since most of these errors are for developers or users. Not operations. --- .../plugins/slo/server/services/create_slo.ts | 4 +-- .../plugins/slo/server/services/reset_slo.ts | 2 +- .../slo/server/services/resource_installer.ts | 34 +++++-------------- .../slo/server/services/slo_repository.ts | 2 +- .../summary_search_client.ts | 2 +- .../services/summay_transform_manager.ts | 12 +++---- .../tasks/orphan_summary_cleanup_task.ts | 4 +-- .../tasks/temp_summary_cleanup_task.ts | 6 ++-- .../services/transform_generators/common.ts | 2 +- .../slo/server/services/transform_manager.ts | 18 +++++----- .../remote_summary_doc_to_slo.ts | 8 ++--- .../plugins/slo/server/services/update_slo.ts | 8 ++--- 12 files changed, 43 insertions(+), 59 deletions(-) diff --git a/x-pack/solutions/observability/plugins/slo/server/services/create_slo.ts b/x-pack/solutions/observability/plugins/slo/server/services/create_slo.ts index e135e35119a2f..25a6c91eaf0a8 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/create_slo.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/create_slo.ts @@ -100,7 +100,7 @@ export class CreateSLO { this.summaryTransformManager.start(summaryTransformId), ]); } catch (err) { - this.logger.error( + this.logger.debug( `Cannot create the SLO [id: ${slo.id}, revision: ${slo.revision}]. Rolling back. ${err}` ); @@ -108,7 +108,7 @@ export class CreateSLO { try { await operation(); } catch (rollbackErr) { - this.logger.error(`Rollback operation failed. ${rollbackErr}`); + this.logger.debug(`Rollback operation failed. ${rollbackErr}`); } }); diff --git a/x-pack/solutions/observability/plugins/slo/server/services/reset_slo.ts b/x-pack/solutions/observability/plugins/slo/server/services/reset_slo.ts index 24aa6a705741c..c4d4091060050 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/reset_slo.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/reset_slo.ts @@ -86,7 +86,7 @@ export class ResetSLO { { logger: this.logger } ); } catch (err) { - this.logger.error( + this.logger.debug( `Cannot reset the SLO [id: ${slo.id}, revision: ${slo.revision}]. Rolling back. ${err}` ); diff --git a/x-pack/solutions/observability/plugins/slo/server/services/resource_installer.ts b/x-pack/solutions/observability/plugins/slo/server/services/resource_installer.ts index e69d37262b69f..64b9f90968440 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/resource_installer.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/resource_installer.ts @@ -42,7 +42,7 @@ export class DefaultResourceInstaller implements ResourceInstaller { try { installTimeout = setTimeout(() => (this.isInstalling = false), 60000); - this.logger.info('Installing SLO shared resources'); + this.logger.debug('Installing SLO shared resources'); await Promise.all([ this.createOrUpdateComponentTemplate(SLI_MAPPINGS_TEMPLATE), this.createOrUpdateComponentTemplate(SLI_SETTINGS_TEMPLATE), @@ -71,9 +71,9 @@ export class DefaultResourceInstaller implements ResourceInstaller { this.esClient ); if (template._meta?.version && currentVersion === template._meta.version) { - this.logger.info(`SLO component template found with version [${template._meta.version}]`); + this.logger.debug(`SLO component template found with version [${template._meta.version}]`); } else { - this.logger.info(`Installing SLO component template [${template.name}]`); + this.logger.debug(`Installing SLO component template [${template.name}]`); return this.execute(() => this.esClient.cluster.putComponentTemplate(template)); } } @@ -86,9 +86,9 @@ export class DefaultResourceInstaller implements ResourceInstaller { ); if (template._meta?.version && currentVersion === template._meta.version) { - this.logger.info(`SLO index template found with version [${template._meta.version}]`); + this.logger.debug(`SLO index template found with version [${template._meta.version}]`); } else { - this.logger.info(`Installing SLO index template [${template.name}]`); + this.logger.debug(`Installing SLO index template [${template.name}]`); return this.execute(() => this.esClient.indices.putIndexTemplate(template)); } } @@ -114,19 +114,11 @@ async function fetchComponentTemplateVersion( esClient: ElasticsearchClient ) { const getTemplateRes = await retryTransientEsErrors( - () => - esClient.cluster.getComponentTemplate( - { - name, - }, - { - ignore: [404], - } - ), + () => esClient.cluster.getComponentTemplate({ name }, { ignore: [404] }), { logger } ); - return getTemplateRes?.component_templates?.[0]?.component_template?._meta?.version || null; + return getTemplateRes?.component_templates?.[0]?.component_template?._meta?.version ?? null; } async function fetchIndexTemplateVersion( @@ -135,17 +127,9 @@ async function fetchIndexTemplateVersion( esClient: ElasticsearchClient ) { const getTemplateRes = await retryTransientEsErrors( - () => - esClient.indices.getIndexTemplate( - { - name, - }, - { - ignore: [404], - } - ), + () => esClient.indices.getIndexTemplate({ name }, { ignore: [404] }), { logger } ); - return getTemplateRes?.index_templates?.[0]?.index_template?._meta?.version || null; + return getTemplateRes?.index_templates?.[0]?.index_template?._meta?.version ?? null; } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/slo_repository.ts b/x-pack/solutions/observability/plugins/slo/server/services/slo_repository.ts index 01580bc92708b..d3bdc04f830e6 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/slo_repository.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/slo_repository.ts @@ -166,7 +166,7 @@ export class KibanaSavedObjectsSLORepository implements SLORepository { }); if (isLeft(result)) { - this.logger.error(`Invalid stored SLO with id [${storedSLO.id}]`); + this.logger.debug(`Invalid stored SLO with id [${storedSLO.id}]`); return undefined; } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/summary_search_client/summary_search_client.ts b/x-pack/solutions/observability/plugins/slo/server/services/summary_search_client/summary_search_client.ts index 3cdae8788d06c..ff4f69e73412e 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/summary_search_client/summary_search_client.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/summary_search_client/summary_search_client.ts @@ -160,7 +160,7 @@ export class DefaultSummarySearchClient implements SummarySearchClient { }), }; } catch (err) { - this.logger.error(`Error while searching SLO summary documents. ${err}`); + this.logger.debug(`Error while searching SLO summary documents. ${err}`); return { total: 0, ...pagination, results: [] }; } } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/summay_transform_manager.ts b/x-pack/solutions/observability/plugins/slo/server/services/summay_transform_manager.ts index 139827b0425a3..078a56b29ca9f 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/summay_transform_manager.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/summay_transform_manager.ts @@ -32,7 +32,7 @@ export class DefaultSummaryTransformManager implements TransformManager { } ); } catch (err) { - this.logger.error(`Cannot create summary transform for SLO [${slo.id}]. ${err}`); + this.logger.debug(`Cannot create summary transform for SLO [${slo.id}]. ${err}`); if (err.meta?.body?.error?.type === 'security_exception') { throw new SecurityException(err.meta.body.error.reason); } @@ -57,7 +57,7 @@ export class DefaultSummaryTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot preview SLO summary transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot preview SLO summary transform [${transformId}]. ${err}`); throw err; } } @@ -75,7 +75,7 @@ export class DefaultSummaryTransformManager implements TransformManager { } ); } catch (err) { - this.logger.error(`Cannot start SLO summary transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot start SLO summary transform [${transformId}]. ${err}`); throw err; } } @@ -91,7 +91,7 @@ export class DefaultSummaryTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot stop SLO summary transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot stop SLO summary transform [${transformId}]. ${err}`); throw err; } } @@ -107,7 +107,7 @@ export class DefaultSummaryTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot delete SLO summary transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot delete SLO summary transform [${transformId}]. ${err}`); throw err; } } @@ -124,7 +124,7 @@ export class DefaultSummaryTransformManager implements TransformManager { ); return response?.transforms[0]?._meta?.version; } catch (err) { - this.logger.error(`Cannot retrieve SLO transform version [${transformId}]. ${err}`); + this.logger.debug(`Cannot retrieve SLO transform version [${transformId}]. ${err}`); throw err; } } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/tasks/orphan_summary_cleanup_task.ts b/x-pack/solutions/observability/plugins/slo/server/services/tasks/orphan_summary_cleanup_task.ts index 55316671ef05f..644b1ef8c562a 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/tasks/orphan_summary_cleanup_task.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/tasks/orphan_summary_cleanup_task.ts @@ -100,7 +100,7 @@ export class SloOrphanSummaryCleanupTask { ); if (sloSummaryIdsToDelete.length > 0) { - this.logger.info( + this.logger.debug( `[SLO] Deleting ${sloSummaryIdsToDelete.length} SLO Summary documents from the summary index` ); @@ -227,7 +227,7 @@ export class SloOrphanSummaryCleanupTask { this.esClient = esClient; if (!taskManager) { - this.logger.info( + this.logger.debug( 'Missing required service during startup, skipping orphan-slo-summary-cleanup task.' ); return; diff --git a/x-pack/solutions/observability/plugins/slo/server/services/tasks/temp_summary_cleanup_task.ts b/x-pack/solutions/observability/plugins/slo/server/services/tasks/temp_summary_cleanup_task.ts index 5d61ae15d65b5..4d3e20b380ac6 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/tasks/temp_summary_cleanup_task.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/tasks/temp_summary_cleanup_task.ts @@ -63,7 +63,7 @@ export class TempSummaryCleanupTask { } if (!plugins.taskManager) { - this.logger.error('Missing required service during start'); + this.logger.debug('Missing required service during start'); return; } @@ -108,7 +108,7 @@ export class TempSummaryCleanupTask { return getDeleteTaskRunResult(); } - this.logger.debug(`runTask() started`); + this.logger.debug(`runTask started`); const [coreStart] = await core.getStartServices(); const esClient = coreStart.elasticsearch.client.asInternalUser; @@ -126,7 +126,7 @@ export class TempSummaryCleanupTask { return; } - this.logger.error(`Error: ${err}`); + this.logger.debug(`Error: ${err}`); } } } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/transform_generators/common.ts b/x-pack/solutions/observability/plugins/slo/server/services/transform_generators/common.ts index 9e8456ccd1b5a..42879c244eddb 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/transform_generators/common.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/transform_generators/common.ts @@ -39,7 +39,7 @@ export function parseStringFilters(filters: string, logger: Logger) { try { return JSON.parse(filters); } catch (e) { - logger.info(`Failed to parse filters: ${e}`); + logger.debug(`Failed to parse filters: ${e}`); return {}; } } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/transform_manager.ts b/x-pack/solutions/observability/plugins/slo/server/services/transform_manager.ts index 0c0666246b66e..3a303bffb8c2e 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/transform_manager.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/transform_manager.ts @@ -34,7 +34,7 @@ export class DefaultTransformManager implements TransformManager { async install(slo: SLODefinition): Promise { const generator = this.generators[slo.indicator.type]; if (!generator) { - this.logger.error(`No transform generator found for indicator type [${slo.indicator.type}]`); + this.logger.debug(`No transform generator found for indicator type [${slo.indicator.type}]`); throw new Error(`Unsupported indicator type [${slo.indicator.type}]`); } @@ -47,7 +47,7 @@ export class DefaultTransformManager implements TransformManager { } ); } catch (err) { - this.logger.error( + this.logger.debug( `Cannot create SLO transform for indicator type [${slo.indicator.type}]. ${err}` ); if (err.meta?.body?.error?.type === 'security_exception') { @@ -63,7 +63,7 @@ export class DefaultTransformManager implements TransformManager { async inspect(slo: SLODefinition): Promise { const generator = this.generators[slo.indicator.type]; if (!generator) { - this.logger.error(`No transform generator found for indicator type [${slo.indicator.type}]`); + this.logger.debug(`No transform generator found for indicator type [${slo.indicator.type}]`); throw new Error(`Unsupported indicator type [${slo.indicator.type}]`); } @@ -80,7 +80,7 @@ export class DefaultTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot preview SLO transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot preview SLO transform [${transformId}]. ${err}`); throw err; } } @@ -97,7 +97,7 @@ export class DefaultTransformManager implements TransformManager { ); await this.scheduleNowTransform(transformId); } catch (err) { - this.logger.error(`Cannot start SLO transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot start SLO transform [${transformId}]. ${err}`); throw err; } } @@ -113,7 +113,7 @@ export class DefaultTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot stop SLO transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot stop SLO transform [${transformId}]. ${err}`); throw err; } } @@ -129,7 +129,7 @@ export class DefaultTransformManager implements TransformManager { { logger: this.logger } ); } catch (err) { - this.logger.error(`Cannot delete SLO transform [${transformId}]. ${err}`); + this.logger.debug(`Cannot delete SLO transform [${transformId}]. ${err}`); throw err; } } @@ -146,7 +146,7 @@ export class DefaultTransformManager implements TransformManager { ); return response?.transforms[0]?._meta?.version; } catch (err) { - this.logger.error(`Cannot retrieve SLO transform version [${transformId}]. ${err}`); + this.logger.debug(`Cannot retrieve SLO transform version [${transformId}]. ${err}`); throw err; } } @@ -158,7 +158,7 @@ export class DefaultTransformManager implements TransformManager { this.logger.debug(`SLO transform [${transformId}] scheduled now successfully`); }) .catch((e) => { - this.logger.error(`Cannot schedule now SLO transform [${transformId}]. ${e}`); + this.logger.debug(`Cannot schedule now SLO transform [${transformId}]. ${e}`); }); } } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/unsafe_federated/remote_summary_doc_to_slo.ts b/x-pack/solutions/observability/plugins/slo/server/services/unsafe_federated/remote_summary_doc_to_slo.ts index 1efe9f58c3bdd..2ed764468197a 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/unsafe_federated/remote_summary_doc_to_slo.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/unsafe_federated/remote_summary_doc_to_slo.ts @@ -48,8 +48,8 @@ export function fromRemoteSummaryDocumentToSloDefinition( if (isLeft(res)) { const errors = formatErrors(res.left); - logger.error(`Invalid remote stored summary SLO with id [${summaryDoc.slo.id}]`); - logger.error(errors.join('|')); + logger.debug(`Invalid remote stored summary SLO with id [${summaryDoc.slo.id}]`); + logger.debug(errors.join('|')); return undefined; } @@ -65,10 +65,10 @@ function getIndicator(summaryDoc: EsSummaryDocument, logger: Logger): Indicator if (isLeft(res)) { const errors = formatErrors(res.left); - logger.info( + logger.debug( `Invalid indicator from remote summary SLO id [${summaryDoc.slo.id}] - Fallback on dummy indicator` ); - logger.info(errors.join('|')); + logger.debug(errors.join('|')); return getDummyIndicator(summaryDoc); } diff --git a/x-pack/solutions/observability/plugins/slo/server/services/update_slo.ts b/x-pack/solutions/observability/plugins/slo/server/services/update_slo.ts index b526d4ff6fba9..d8ff5a508b6d3 100644 --- a/x-pack/solutions/observability/plugins/slo/server/services/update_slo.ts +++ b/x-pack/solutions/observability/plugins/slo/server/services/update_slo.ts @@ -82,7 +82,7 @@ export class UpdateSLO { { logger: this.logger } ); } catch (err) { - this.logger.error( + this.logger.debug( `Cannot update the SLO summary pipeline [id: ${updatedSlo.id}, revision: ${updatedSlo.revision}]. ${err}` ); @@ -90,7 +90,7 @@ export class UpdateSLO { try { await operation(); } catch (rollbackErr) { - this.logger.error(`Rollback operation failed. ${rollbackErr}`); + this.logger.debug(`Rollback operation failed. ${rollbackErr}`); } }); @@ -161,7 +161,7 @@ export class UpdateSLO { { logger: this.logger } ); } catch (err) { - this.logger.error( + this.logger.debug( `Cannot update the SLO [id: ${updatedSlo.id}, revision: ${updatedSlo.revision}]. Rolling back. ${err}` ); @@ -169,7 +169,7 @@ export class UpdateSLO { try { await operation(); } catch (rollbackErr) { - this.logger.error(`Rollback operation failed. ${rollbackErr}`); + this.logger.debug(`Rollback operation failed. ${rollbackErr}`); } }); From ef0c364f11b485c29e4c9d065a93c2a4d12fb5eb Mon Sep 17 00:00:00 2001 From: Joe Reuter Date: Tue, 4 Mar 2025 21:11:24 +0100 Subject: [PATCH 8/9] =?UTF-8?q?=F0=9F=8C=8A=20Streams:=20Improve=20integra?= =?UTF-8?q?tion=20tests=20(#213115)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit I noticed that a couple integration tests were not actually validating whether the documents got routed the right way (the number and string tests). This PR fixes this by introducing a helper that can easily do the same check we had a couple of times in there. --- .../apis/observability/streams/full_flow.ts | 50 ++++--------------- .../observability/streams/helpers/requests.ts | 11 ++++ 2 files changed, 22 insertions(+), 39 deletions(-) diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/full_flow.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/full_flow.ts index b1b547f67e365..e9db9f9e57706 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/full_flow.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/full_flow.ts @@ -16,6 +16,7 @@ import { enableStreams, fetchDocument, forkStream, + indexAndAssertTargetStream, indexDocument, } from './helpers/requests'; @@ -153,10 +154,7 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { stream: 'somethingelse', // a field named stream should work as well }), }; - const response = await indexDocument(esClient, 'logs', doc); - expect(response.result).to.eql('created'); - const result = await fetchDocument(esClient, 'logs', response._id); - expect(result._index).to.match(/^\.ds\-logs-.*/); + const result = await indexAndAssertTargetStream(esClient, 'logs', doc); expect(result._source).to.eql({ '@timestamp': '2024-01-01T00:00:00.000Z', message: 'test', @@ -191,11 +189,7 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { message: 'test', }), }; - const response = await indexDocument(esClient, 'logs', doc); - expect(response.result).to.eql('created'); - - const result = await fetchDocument(esClient, 'logs.nginx', response._id); - expect(result._index).to.match(/^\.ds\-logs.nginx-.*/); + const result = await indexAndAssertTargetStream(esClient, 'logs.nginx', doc); expect(result._source).to.eql({ '@timestamp': '2024-01-01T00:00:10.000Z', message: 'test', @@ -225,11 +219,7 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { message: 'test', }), }; - const response = await indexDocument(esClient, 'logs', doc); - expect(response.result).to.eql('created'); - - const result = await fetchDocument(esClient, 'logs.nginx.access', response._id); - expect(result._index).to.match(/^\.ds\-logs.nginx.access-.*/); + const result = await indexAndAssertTargetStream(esClient, 'logs.nginx.access', doc); expect(result._source).to.eql({ '@timestamp': '2024-01-01T00:00:20.000Z', message: 'test', @@ -259,11 +249,7 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { message: 'test', }), }; - const response = await indexDocument(esClient, 'logs', doc); - expect(response.result).to.eql('created'); - - const result = await fetchDocument(esClient, 'logs.nginx', response._id); - expect(result._index).to.match(/^\.ds\-logs.nginx-.*/); + const result = await indexAndAssertTargetStream(esClient, 'logs.nginx', doc); expect(result._source).to.eql({ '@timestamp': '2024-01-01T00:00:20.000Z', message: 'test', @@ -299,10 +285,8 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { message: 'test', }), }; - const response1 = await indexDocument(esClient, 'logs', doc1); - expect(response1.result).to.eql('created'); - const response2 = await indexDocument(esClient, 'logs', doc2); - expect(response2.result).to.eql('created'); + await indexAndAssertTargetStream(esClient, 'logs.number-test', doc1); + await indexAndAssertTargetStream(esClient, 'logs.number-test', doc2); }); it('Fork logs to logs.string-test', async () => { @@ -334,11 +318,8 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { message: 'status_code: 400', }), }; - const response1 = await indexDocument(esClient, 'logs', doc1); - expect(response1.result).to.eql('created'); - - const response2 = await indexDocument(esClient, 'logs', doc2); - expect(response2.result).to.eql('created'); + await indexAndAssertTargetStream(esClient, 'logs.string-test', doc1); + await indexAndAssertTargetStream(esClient, 'logs.string-test', doc2); }); it('Fork logs to logs.weird-characters', async () => { @@ -369,17 +350,8 @@ export default function ({ getService }: DeploymentAgnosticFtrProviderContext) { 'weird fieldname': 'Keep where it is', }, }; - const response1 = await indexDocument(esClient, 'logs', doc1); - expect(response1.result).to.eql('created'); - - const result1 = await fetchDocument(esClient, 'logs.weird-characters', response1._id); - expect(result1._index).to.match(/^\.ds\-logs.weird-characters-.*/); - - const response2 = await indexDocument(esClient, 'logs', doc2); - expect(response2.result).to.eql('created'); - - const result2 = await fetchDocument(esClient, 'logs', response2._id); - expect(result2._index).to.match(/^\.ds\-logs-.*/); + await indexAndAssertTargetStream(esClient, 'logs.weird-characters', doc1); + await indexAndAssertTargetStream(esClient, 'logs', doc2); }); }); }); diff --git a/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/helpers/requests.ts b/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/helpers/requests.ts index eed7a6e1d7e55..f7ec43d19d72b 100644 --- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/helpers/requests.ts +++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/streams/helpers/requests.ts @@ -26,6 +26,17 @@ export async function indexDocument(esClient: Client, index: string, document: J return response; } +export async function indexAndAssertTargetStream( + esClient: Client, + target: string, + document: JsonObject +) { + const response = await esClient.index({ index: 'logs', document, refresh: 'wait_for' }); + const result = await fetchDocument(esClient, target, response._id); + expect(result._index).to.match(new RegExp(`^\.ds\-${target}-.*`)); + return result; +} + export async function fetchDocument(esClient: Client, index: string, id: string) { const query = { ids: { values: [id] }, From b32f0fe1e863a599c9d61a38e99c974deec6519a Mon Sep 17 00:00:00 2001 From: Hannah Mudge Date: Tue, 4 Mar 2025 13:34:39 -0700 Subject: [PATCH 9/9] [kbn-grid-layout] Store rows in object instead of array (#212965) Closes https://github.com/elastic/kibana/issues/211930 ## Summary This PR makes it so that `kbn-grid-layout` stores its rows as an object / dictionary (`{ [key: string]: GridRowData }`) rather than an array (`Array`). This is a prerequisite for https://github.com/elastic/kibana/issues/190381 , since it allows us to re-order rows without re-rendering their contents. It also means that deleting a row will no longer cause the rows below it to re-render, since re-rendering is now dependant on the row's **ID** rather than the row's order. **Before** https://github.com/user-attachments/assets/83651b24-a32c-4953-8ad5-c0eced163eb5 **After** https://github.com/user-attachments/assets/9cef6dbc-3d62-46aa-bc40-ab24fc4e5556 ### Checklist - [x] [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 - [x] The PR description includes the appropriate Release Notes section, and the correct `release_note:*` label is applied per the [guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com> --- examples/grid_example/public/app.tsx | 26 +++++---- .../public/logs_dashboard_panels.json | 58 +++++++++---------- .../public/serialized_grid_layout.ts | 12 ++-- examples/grid_example/public/types.ts | 6 +- .../public/use_mock_dashboard_api.tsx | 2 +- examples/grid_example/public/utils.ts | 25 ++++---- .../kbn-grid-layout/grid/drag_preview.tsx | 6 +- .../kbn-grid-layout/grid/grid_layout.test.tsx | 14 ++--- .../kbn-grid-layout/grid/grid_layout.tsx | 42 +++++--------- .../drag_handle/use_drag_handle_api.tsx | 6 +- .../grid/grid_panel/grid_panel.test.tsx | 2 +- .../grid/grid_panel/grid_panel.tsx | 26 ++++----- .../grid/grid_panel/resize_handle.tsx | 38 ++++++------ .../grid/grid_row/delete_grid_row_modal.tsx | 22 ++++--- .../grid/grid_row/grid_row.test.tsx | 10 ++-- .../grid/grid_row/grid_row.tsx | 48 +++++++-------- .../grid/grid_row/grid_row_header.test.tsx | 56 +++++++++--------- .../grid/grid_row/grid_row_header.tsx | 26 ++++----- .../grid/grid_row/grid_row_title.tsx | 24 ++++---- .../grid/test_utils/sample_layout.ts | 16 +++-- .../private/kbn-grid-layout/grid/types.ts | 14 +++-- .../grid/use_grid_layout_events/index.ts | 9 ++- .../state_manager_actions.ts | 42 +++++++------- .../grid/use_grid_layout_state.ts | 12 ++-- .../grid/utils/equality_checks.ts | 12 ++-- .../grid/utils/resolve_grid_row.test.ts | 22 +++++++ .../grid/utils/resolve_grid_row.ts | 17 ++++-- .../grid/utils/row_management.ts | 6 +- .../component/grid/dashboard_grid.tsx | 12 ++-- 29 files changed, 325 insertions(+), 286 deletions(-) diff --git a/examples/grid_example/public/app.tsx b/examples/grid_example/public/app.tsx index 4ea2e4ed7bf42..0b51915012acd 100644 --- a/examples/grid_example/public/app.tsx +++ b/examples/grid_example/public/app.tsx @@ -8,9 +8,11 @@ */ import deepEqual from 'fast-deep-equal'; +import { cloneDeep } from 'lodash'; import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import ReactDOM from 'react-dom'; import { Subject, combineLatest, debounceTime, map, skip, take } from 'rxjs'; +import { v4 as uuidv4 } from 'uuid'; import { EuiBadge, @@ -88,8 +90,8 @@ export const GridExample = ({ const currentPanel = panels[panelId]; const savedPanel = savedState.current.panels[panelId]; panelsAreEqual = deepEqual( - { row: 0, ...currentPanel.gridData }, - { row: 0, ...savedPanel.gridData } + { row: 'first', ...currentPanel.gridData }, + { row: 'first', ...savedPanel.gridData } ); } const hasChanges = !(panelsAreEqual && deepEqual(rows, savedState.current.rows)); @@ -147,15 +149,17 @@ export const GridExample = ({ ); const addNewSection = useCallback(() => { - mockDashboardApi.rows$.next([ - ...mockDashboardApi.rows$.getValue(), - { - title: i18n.translate('examples.gridExample.defaultSectionTitle', { - defaultMessage: 'New collapsible section', - }), - collapsed: false, - }, - ]); + const rows = cloneDeep(mockDashboardApi.rows$.getValue()); + const id = uuidv4(); + rows[id] = { + id, + order: Object.keys(rows).length, + title: i18n.translate('examples.gridExample.defaultSectionTitle', { + defaultMessage: 'New collapsible section', + }), + collapsed: false, + }; + mockDashboardApi.rows$.next(rows); // scroll to bottom after row is added layoutUpdated$.pipe(skip(1), take(1)).subscribe(() => { diff --git a/examples/grid_example/public/logs_dashboard_panels.json b/examples/grid_example/public/logs_dashboard_panels.json index 29d63baa2d33b..3b5522975db7c 100644 --- a/examples/grid_example/public/logs_dashboard_panels.json +++ b/examples/grid_example/public/logs_dashboard_panels.json @@ -1004,7 +1004,7 @@ "w": 48, "h": 17, "i": "4", - "row": 1 + "row": "second" }, "explicitInput": { "id": "4", @@ -1035,7 +1035,7 @@ "w": 18, "h": 8, "i": "05da0d2b-0145-4068-b21c-00be3184d465", - "row": 1 + "row": "second" }, "explicitInput": { "id": "05da0d2b-0145-4068-b21c-00be3184d465", @@ -1073,7 +1073,7 @@ "w": 18, "h": 16, "i": "b7da9075-4742-47e3-b4f8-fc9ba82de74c", - "row": 1 + "row": "second" }, "explicitInput": { "id": "b7da9075-4742-47e3-b4f8-fc9ba82de74c", @@ -1111,7 +1111,7 @@ "w": 12, "h": 16, "i": "5c409557-644d-4c05-a284-ffe54bb28db0", - "row": 1 + "row": "second" }, "explicitInput": { "id": "5c409557-644d-4c05-a284-ffe54bb28db0", @@ -1234,7 +1234,7 @@ "w": 6, "h": 8, "i": "af4b5c07-506e-44c2-b2bb-2113d0c5b274", - "row": 1 + "row": "second" }, "explicitInput": { "id": "af4b5c07-506e-44c2-b2bb-2113d0c5b274", @@ -1400,7 +1400,7 @@ "w": 6, "h": 8, "i": "d42c4870-c028-4d8a-abd0-0effbc190ce3", - "row": 1 + "row": "second" }, "explicitInput": { "id": "d42c4870-c028-4d8a-abd0-0effbc190ce3", @@ -1520,7 +1520,7 @@ "w": 6, "h": 8, "i": "4092d42c-f93b-4c71-a6db-8f12abf12791", - "row": 1 + "row": "second" }, "explicitInput": { "id": "4092d42c-f93b-4c71-a6db-8f12abf12791", @@ -1641,7 +1641,7 @@ "w": 30, "h": 15, "i": "15", - "row": 2 + "row": "third" }, "explicitInput": { "id": "15", @@ -1887,7 +1887,7 @@ "w": 18, "h": 8, "i": "4e64d6d7-4f92-4d5e-abbb-13796604db30", - "row": 2 + "row": "third" }, "explicitInput": { "id": "4e64d6d7-4f92-4d5e-abbb-13796604db30v", @@ -1925,7 +1925,7 @@ "w": 6, "h": 7, "i": "ddce4ad8-6a82-44f0-9995-57f46f153f50", - "row": 2 + "row": "third" }, "explicitInput": { "id": "ddce4ad8-6a82-44f0-9995-57f46f153f50", @@ -2120,7 +2120,7 @@ "w": 6, "h": 7, "i": "a2884704-db3b-4b92-a19a-cdfe668dec39", - "row": 2 + "row": "third" }, "explicitInput": { "id": "a2884704-db3b-4b92-a19a-cdfe668dec39", @@ -2315,7 +2315,7 @@ "w": 6, "h": 7, "i": "529eec49-10e2-4a40-9c77-5c81f4eb3943", - "row": 2 + "row": "third" }, "explicitInput": { "id": "529eec49-10e2-4a40-9c77-5c81f4eb3943", @@ -2510,7 +2510,7 @@ "w": 48, "h": 12, "i": "1d5f0b3f-d9d2-4b26-997b-83bc5ca3090b", - "row": 2 + "row": "third" }, "explicitInput": { "id": "1d5f0b3f-d9d2-4b26-997b-83bc5ca3090b", @@ -2905,7 +2905,7 @@ "w": 48, "h": 15, "i": "9f79ecca-123f-4098-a658-6b0e998da003", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "9f79ecca-123f-4098-a658-6b0e998da003", @@ -2922,7 +2922,7 @@ "w": 24, "h": 9, "i": "7", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "7", @@ -3161,7 +3161,7 @@ "w": 24, "h": 11, "i": "10", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "10", @@ -3346,7 +3346,7 @@ "w": 24, "h": 22, "i": "23", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "23", @@ -3371,7 +3371,7 @@ "w": 24, "h": 22, "i": "31", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "31", @@ -3388,7 +3388,7 @@ "w": 24, "h": 8, "i": "6afc61f7-e2d5-45a3-9e7a-281160ad3eb9", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "6afc61f7-e2d5-45a3-9e7a-281160ad3eb9", @@ -3420,7 +3420,7 @@ "w": 8, "h": 8, "i": "392b4936-f753-47bc-a98d-a4e41a0a4cd4", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "392b4936-f753-47bc-a98d-a4e41a0a4cd4", @@ -3485,7 +3485,7 @@ "w": 8, "h": 4, "i": "9271deff-5a61-4665-83fc-f9fdc6bf0c0b", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "9271deff-5a61-4665-83fc-f9fdc6bf0c0b", @@ -3613,7 +3613,7 @@ "w": 8, "h": 4, "i": "aa591c29-1a31-4ee1-a71d-b829c06fd162", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "aa591c29-1a31-4ee1-a71d-b829c06fd162", @@ -3777,7 +3777,7 @@ "w": 8, "h": 4, "i": "b766e3b8-4544-46ed-99e6-9ecc4847e2a2", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "b766e3b8-4544-46ed-99e6-9ecc4847e2a2", @@ -3905,7 +3905,7 @@ "w": 8, "h": 4, "i": "2e33ade5-96e5-40b4-b460-493e5d4fa834", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "2e33ade5-96e5-40b4-b460-493e5d4fa834", @@ -4069,7 +4069,7 @@ "w": 24, "h": 8, "i": "086ac2e9-dd16-4b45-92b8-1e43ff7e3f65", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "086ac2e9-dd16-4b45-92b8-1e43ff7e3f65", @@ -4190,7 +4190,7 @@ "w": 24, "h": 28, "i": "fb86b32f-fb7a-45cf-9511-f366fef51bbd", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "fb86b32f-fb7a-45cf-9511-f366fef51bbd", @@ -4500,7 +4500,7 @@ "w": 24, "h": 11, "i": "0cc42484-16f7-42ec-b38c-9bf8be69cde7", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "0cc42484-16f7-42ec-b38c-9bf8be69cde7", @@ -4643,7 +4643,7 @@ "w": 12, "h": 11, "i": "5d53db36-2d5a-4adc-af7b-cec4c1a294e0", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "5d53db36-2d5a-4adc-af7b-cec4c1a294e0", @@ -4773,7 +4773,7 @@ "w": 12, "h": 11, "i": "ecd89a7c-9124-4472-bdc6-9bdbd70d45d5", - "row": 3 + "row": "fourth" }, "explicitInput": { "id": "ecd89a7c-9124-4472-bdc6-9bdbd70d45d5", diff --git a/examples/grid_example/public/serialized_grid_layout.ts b/examples/grid_example/public/serialized_grid_layout.ts index e5664bd1f22e5..38000e20d74ef 100644 --- a/examples/grid_example/public/serialized_grid_layout.ts +++ b/examples/grid_example/public/serialized_grid_layout.ts @@ -28,10 +28,10 @@ export function setSerializedGridLayout(state: MockSerializedDashboardState) { const initialState: MockSerializedDashboardState = { panels: logsPanels, - rows: [ - { title: 'Request Sizes', collapsed: false }, - { title: 'Visitors', collapsed: false }, - { title: 'Response Codes', collapsed: false }, - { title: 'Entire Flights Dashboard', collapsed: true }, - ], + rows: { + first: { id: 'first', order: 0, title: 'Request Sizes', collapsed: false }, + second: { id: 'second', order: 1, title: 'Visitors', collapsed: false }, + third: { id: 'third', order: 2, title: 'Response Codes', collapsed: false }, + fourth: { id: 'fourth', order: 3, title: 'Entire Flights Dashboard', collapsed: true }, + }, }; diff --git a/examples/grid_example/public/types.ts b/examples/grid_example/public/types.ts index 705b652e3d6bf..70e6b840f61f9 100644 --- a/examples/grid_example/public/types.ts +++ b/examples/grid_example/public/types.ts @@ -26,7 +26,7 @@ export interface DashboardGridData { interface DashboardPanelState { type: string; - gridData: DashboardGridData & { row?: number }; + gridData: DashboardGridData & { row?: string }; explicitInput: Partial & { id: string }; version?: string; } @@ -35,7 +35,9 @@ export interface MockedDashboardPanelMap { [key: string]: DashboardPanelState; } -export type MockedDashboardRowMap = Array<{ title: string; collapsed: boolean }>; +export interface MockedDashboardRowMap { + [id: string]: { id: string; order: number; title: string; collapsed: boolean }; +} export interface MockSerializedDashboardState { panels: MockedDashboardPanelMap; diff --git a/examples/grid_example/public/use_mock_dashboard_api.tsx b/examples/grid_example/public/use_mock_dashboard_api.tsx index ec91fc762cabd..e013ba07a6111 100644 --- a/examples/grid_example/public/use_mock_dashboard_api.tsx +++ b/examples/grid_example/public/use_mock_dashboard_api.tsx @@ -99,7 +99,7 @@ export const useMockDashboardApi = ({ [newId]: { type: panelPackage.panelType, gridData: { - row: 0, + row: 'first', x: 0, y: 0, w: DEFAULT_PANEL_WIDTH, diff --git a/examples/grid_example/public/utils.ts b/examples/grid_example/public/utils.ts index 8c9db472cffa7..b40f0f846f08c 100644 --- a/examples/grid_example/public/utils.ts +++ b/examples/grid_example/public/utils.ts @@ -15,10 +15,11 @@ export const gridLayoutToDashboardPanelMap = ( layout: GridLayoutData ): { panels: MockedDashboardPanelMap; rows: MockedDashboardRowMap } => { const panels: MockedDashboardPanelMap = {}; - const rows: MockedDashboardRowMap = []; - layout.forEach((row, rowIndex) => { - rows.push({ title: row.title, collapsed: row.isCollapsed }); - Object.values(row.panels).forEach((panelGridData) => { + const rows: MockedDashboardRowMap = {}; + Object.entries(layout).forEach(([rowId, row]) => { + const { panels: rowPanels, isCollapsed, ...rest } = row; // drop panels + rows[rowId] = { ...rest, collapsed: isCollapsed }; + Object.values(rowPanels).forEach((panelGridData) => { panels[panelGridData.id] = { ...panelState[panelGridData.id], gridData: { @@ -27,7 +28,7 @@ export const gridLayoutToDashboardPanelMap = ( x: panelGridData.column, w: panelGridData.width, h: panelGridData.height, - row: rowIndex, + row: rowId, }, }; }); @@ -42,15 +43,19 @@ export const dashboardInputToGridLayout = ({ panels: MockedDashboardPanelMap; rows: MockedDashboardRowMap; }): GridLayoutData => { - const layout: GridLayoutData = []; - - rows.forEach((row) => { - layout.push({ title: row.title, isCollapsed: row.collapsed, panels: {} }); + const layout: GridLayoutData = {}; + Object.values(rows).forEach((row) => { + const { collapsed, ...rest } = row; + layout[row.id] = { + ...rest, + panels: {}, + isCollapsed: collapsed, + }; }); Object.keys(panels).forEach((panelId) => { const gridData = panels[panelId].gridData; - layout[gridData.row ?? 0].panels[panelId] = { + layout[gridData.row ?? 'first'].panels[panelId] = { id: panelId, row: gridData.y, column: gridData.x, diff --git a/src/platform/packages/private/kbn-grid-layout/grid/drag_preview.tsx b/src/platform/packages/private/kbn-grid-layout/grid/drag_preview.tsx index 0d34c33b674fe..dd489bdc915be 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/drag_preview.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/drag_preview.tsx @@ -13,7 +13,7 @@ import { combineLatest, skip } from 'rxjs'; import { css } from '@emotion/react'; import { useGridLayoutContext } from './use_grid_layout_context'; -export const DragPreview = React.memo(({ rowIndex }: { rowIndex: number }) => { +export const DragPreview = React.memo(({ rowId }: { rowId: string }) => { const { gridLayoutStateManager } = useGridLayoutContext(); const dragPreviewRef = useRef(null); @@ -29,10 +29,10 @@ export const DragPreview = React.memo(({ rowIndex }: { rowIndex: number }) => { .subscribe(([activePanel, proposedGridLayout]) => { if (!dragPreviewRef.current) return; - if (!activePanel || !proposedGridLayout?.[rowIndex].panels[activePanel.id]) { + if (!activePanel || !proposedGridLayout?.[rowId].panels[activePanel.id]) { dragPreviewRef.current.style.display = 'none'; } else { - const panel = proposedGridLayout[rowIndex].panels[activePanel.id]; + const panel = proposedGridLayout[rowId].panels[activePanel.id]; dragPreviewRef.current.style.display = 'block'; dragPreviewRef.current.style.gridColumnStart = `${panel.column + 1}`; dragPreviewRef.current.style.gridColumnEnd = `${panel.column + 1 + panel.width}`; diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.test.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.test.tsx index fa54d83019129..4d4b577f4a8f8 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.test.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.test.tsx @@ -96,10 +96,10 @@ describe('GridLayout', () => { // if layout **has** changed, call `onLayoutChange` const newLayout = cloneDeep(layout); - newLayout[0] = { - ...newLayout[0], + newLayout.first = { + ...newLayout.first, panels: { - ...newLayout[0].panels, + ...newLayout.first.panels, panel1: { id: 'panel1', row: 100, @@ -217,7 +217,7 @@ describe('GridLayout', () => { it('after removing a panel', async () => { const { rerender } = renderGridLayout(); const sampleLayoutWithoutPanel1 = cloneDeep(getSampleLayout()); - delete sampleLayoutWithoutPanel1[0].panels.panel1; + delete sampleLayoutWithoutPanel1.first.panels.panel1; rerender({ layout: sampleLayoutWithoutPanel1 }); expect(getAllThePanelIds()).toEqual([ @@ -236,9 +236,9 @@ describe('GridLayout', () => { it('after replacing a panel id', async () => { const { rerender } = renderGridLayout(); const modifiedLayout = cloneDeep(getSampleLayout()); - const newPanel = { ...modifiedLayout[0].panels.panel1, id: 'panel11' }; - delete modifiedLayout[0].panels.panel1; - modifiedLayout[0].panels.panel11 = newPanel; + const newPanel = { ...modifiedLayout.first.panels.panel1, id: 'panel11' }; + delete modifiedLayout.first.panels.panel1; + modifiedLayout.first.panels.panel11 = newPanel; rerender({ layout: modifiedLayout }); diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.tsx index 9dd0a2d44d7a3..d2d20ca7dec26 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_layout.tsx @@ -8,9 +8,10 @@ */ import classNames from 'classnames'; +import deepEqual from 'fast-deep-equal'; import { cloneDeep } from 'lodash'; import React, { useEffect, useMemo, useRef, useState } from 'react'; -import { combineLatest, distinctUntilChanged, map, pairwise, skip } from 'rxjs'; +import { combineLatest, pairwise } from 'rxjs'; import { css } from '@emotion/react'; @@ -20,7 +21,7 @@ import { GridAccessMode, GridLayoutData, GridSettings, UseCustomDragHandle } fro import { GridLayoutContext, GridLayoutContextType } from './use_grid_layout_context'; import { useGridLayoutState } from './use_grid_layout_state'; import { isLayoutEqual } from './utils/equality_checks'; -import { resolveGridRow } from './utils/resolve_grid_row'; +import { getRowKeysInOrder, resolveGridRow } from './utils/resolve_grid_row'; export type GridLayoutProps = { layout: GridLayoutData; @@ -50,10 +51,7 @@ export const GridLayout = ({ accessMode, }); - const [rowCount, setRowCount] = useState( - gridLayoutStateManager.gridLayout$.getValue().length - ); - + const [rowIdsInOrder, setRowIdsInOrder] = useState(getRowKeysInOrder(layout)); /** * Update the `gridLayout$` behaviour subject in response to the `layout` prop changing */ @@ -64,8 +62,8 @@ export const GridLayout = ({ * the layout sent in as a prop is not guaranteed to be valid (i.e it may have floating panels) - * so, we need to loop through each row and ensure it is compacted */ - newLayout.forEach((row, rowIndex) => { - newLayout[rowIndex] = resolveGridRow(row); + Object.entries(newLayout).forEach(([rowId, row]) => { + newLayout[rowId] = resolveGridRow(row); }); gridLayoutStateManager.gridLayout$.next(newLayout); } @@ -77,27 +75,18 @@ export const GridLayout = ({ */ useEffect(() => { /** - * The only thing that should cause the entire layout to re-render is adding a new row; - * this subscription ensures this by updating the `rowCount` state when it changes. - */ - const rowCountSubscription = gridLayoutStateManager.gridLayout$ - .pipe( - skip(1), // we initialized `rowCount` above, so skip the initial emit - map((newLayout) => newLayout.length), - distinctUntilChanged() - ) - .subscribe((newRowCount) => { - setRowCount(newRowCount); - }); - - /** - * This subscription calls the passed `onLayoutChange` callback when the layout changes + * This subscription calls the passed `onLayoutChange` callback when the layout changes; + * if the row IDs have changed, it also sets `rowIdsInOrder` to trigger a re-render */ const onLayoutChangeSubscription = gridLayoutStateManager.gridLayout$ .pipe(pairwise()) .subscribe(([layoutBefore, layoutAfter]) => { if (!isLayoutEqual(layoutBefore, layoutAfter)) { onLayoutChange(layoutAfter); + + if (!deepEqual(Object.keys(layoutBefore), Object.keys(layoutAfter))) { + setRowIdsInOrder(getRowKeysInOrder(layoutAfter)); + } } }); @@ -125,7 +114,6 @@ export const GridLayout = ({ }); return () => { - rowCountSubscription.unsubscribe(); onLayoutChangeSubscription.unsubscribe(); gridLayoutClassSubscription.unsubscribe(); }; @@ -158,9 +146,9 @@ export const GridLayout = ({ styles.hasExpandedPanel, ]} > - {Array.from({ length: rowCount }, (_, rowIndex) => { - return ; - })} + {rowIdsInOrder.map((rowId) => ( + + ))} diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/drag_handle/use_drag_handle_api.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/drag_handle/use_drag_handle_api.tsx index a2e3c632d1364..b3470018cd5b3 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/drag_handle/use_drag_handle_api.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/drag_handle/use_drag_handle_api.tsx @@ -20,17 +20,17 @@ export interface DragHandleApi { export const useDragHandleApi = ({ panelId, - rowIndex, + rowId, }: { panelId: string; - rowIndex: number; + rowId: string; }): DragHandleApi => { const { useCustomDragHandle } = useGridLayoutContext(); const startInteraction = useGridLayoutEvents({ interactionType: 'drag', panelId, - rowIndex, + rowId, }); const removeEventListenersRef = useRef<(() => void) | null>(null); diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.test.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.test.tsx index 5df0b7d831e5e..24b8934621011 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.test.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.test.tsx @@ -25,7 +25,7 @@ describe('GridPanel', () => { } as GridLayoutContextType; const panelProps = { panelId: 'panel1', - rowIndex: 0, + rowId: 'first', ...(overrides?.propsOverrides ?? {}), }; const { rerender, ...rtlRest } = render( diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.tsx index c29a9ccae1bd4..f78d3976e610c 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/grid_panel.tsx @@ -20,20 +20,20 @@ import { ResizeHandle } from './resize_handle'; export interface GridPanelProps { panelId: string; - rowIndex: number; + rowId: string; } -export const GridPanel = React.memo(({ panelId, rowIndex }: GridPanelProps) => { +export const GridPanel = React.memo(({ panelId, rowId }: GridPanelProps) => { const { gridLayoutStateManager, useCustomDragHandle, renderPanelContents } = useGridLayoutContext(); const { euiTheme } = useEuiTheme(); - const dragHandleApi = useDragHandleApi({ panelId, rowIndex }); + const dragHandleApi = useDragHandleApi({ panelId, rowId }); /** Set initial styles based on state at mount to prevent styles from "blipping" */ const initialStyles = useMemo(() => { const initialPanel = (gridLayoutStateManager.proposedGridLayout$.getValue() ?? - gridLayoutStateManager.gridLayout$.getValue())[rowIndex].panels[panelId]; + gridLayoutStateManager.gridLayout$.getValue())[rowId].panels[panelId]; return css` position: relative; height: calc( @@ -48,7 +48,7 @@ export const GridPanel = React.memo(({ panelId, rowIndex }: GridPanelProps) => { grid-row-start: ${initialPanel.row + 1}; grid-row-end: ${initialPanel.row + 1 + initialPanel.height}; `; - }, [gridLayoutStateManager, rowIndex, panelId]); + }, [gridLayoutStateManager, rowId, panelId]); useEffect( () => { @@ -60,8 +60,8 @@ export const GridPanel = React.memo(({ panelId, rowIndex }: GridPanelProps) => { ]) .pipe(skip(1)) // skip the first emit because the `initialStyles` will take care of it .subscribe(([activePanel, gridLayout, proposedGridLayout]) => { - const ref = gridLayoutStateManager.panelRefs.current[rowIndex][panelId]; - const panel = (proposedGridLayout ?? gridLayout)[rowIndex].panels[panelId]; + const ref = gridLayoutStateManager.panelRefs.current[rowId][panelId]; + const panel = (proposedGridLayout ?? gridLayout)[rowId]?.panels[panelId]; if (!ref || !panel) return; const currentInteractionEvent = gridLayoutStateManager.interactionEvent$.getValue(); @@ -128,9 +128,9 @@ export const GridPanel = React.memo(({ panelId, rowIndex }: GridPanelProps) => { */ const expandedPanelSubscription = gridLayoutStateManager.expandedPanelId$.subscribe( (expandedPanelId) => { - const ref = gridLayoutStateManager.panelRefs.current[rowIndex][panelId]; + const ref = gridLayoutStateManager.panelRefs.current[rowId][panelId]; const gridLayout = gridLayoutStateManager.gridLayout$.getValue(); - const panel = gridLayout[rowIndex].panels[panelId]; + const panel = gridLayout[rowId].panels[panelId]; if (!ref || !panel) return; if (expandedPanelId && expandedPanelId === panelId) { @@ -160,17 +160,17 @@ export const GridPanel = React.memo(({ panelId, rowIndex }: GridPanelProps) => { return (
{ - if (!gridLayoutStateManager.panelRefs.current[rowIndex]) { - gridLayoutStateManager.panelRefs.current[rowIndex] = {}; + if (!gridLayoutStateManager.panelRefs.current[rowId]) { + gridLayoutStateManager.panelRefs.current[rowId] = {}; } - gridLayoutStateManager.panelRefs.current[rowIndex][panelId] = element; + gridLayoutStateManager.panelRefs.current[rowId][panelId] = element; }} css={initialStyles} className="kbnGridPanel" > {!useCustomDragHandle && } {panelContents} - +
); }); diff --git a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/resize_handle.tsx b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/resize_handle.tsx index ccb264d2294e1..8fb0254bccba4 100644 --- a/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/resize_handle.tsx +++ b/src/platform/packages/private/kbn-grid-layout/grid/grid_panel/resize_handle.tsx @@ -15,27 +15,25 @@ import { i18n } from '@kbn/i18n'; import { useGridLayoutEvents } from '../use_grid_layout_events'; -export const ResizeHandle = React.memo( - ({ rowIndex, panelId }: { rowIndex: number; panelId: string }) => { - const startInteraction = useGridLayoutEvents({ - interactionType: 'resize', - panelId, - rowIndex, - }); +export const ResizeHandle = React.memo(({ rowId, panelId }: { rowId: string; panelId: string }) => { + const startInteraction = useGridLayoutEvents({ + interactionType: 'resize', + panelId, + rowId, + }); - return ( -