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..f19915770c58a --- /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 '../ai_assistant_data_clients/defend_insights/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/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..42c63b18459ed --- /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, + hasZeroAlerts: 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, + hasZeroAlerts: 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, + hasZeroAlerts: 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..e262db05ec8ce --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/index.ts @@ -0,0 +1,131 @@ +/* + * 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. + * + * Refer to the following diagram for this graph: + * x-pack/solutions/security/plugins/elastic_assistant/docs/img/default_defend_insights_graph.png + */ +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..c69a87f9b5462 --- /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,88 @@ +/* + * 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\n' + 'agent.id,agent-id1\n' + 'process.executable,some/file/path.exe\n', + }, + { + metadata: {}, + pageContent: + '_id,event-id2\n' + 'agent.id,agent-id2\n' + 'process.executable,another/file/path.exe\n', + }, + ], + 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..be0c1ce670c8d --- /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 insightType: DefendInsightType; + private endpointIds: string[]; + private anonymizationFields?: AnonymizationFieldResponse[]; + private esClient: ElasticsearchClient; + private onNewReplacements?: (newReplacements: Replacements) => void; + private replacements?: Replacements; + private size?: number; + private start?: DateMath; + private 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/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_events/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 new file mode 100644 index 0000000000000..c32dfc69b96ea --- /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/get_file_events_query.ts @@ -0,0 +1,58 @@ +/* + * 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 { SearchRequest, DateMath } from '@elastic/elasticsearch/lib/api/types'; + +const FILE_EVENTS_INDEX_PATTERN = 'logs-endpoint.events.file-*'; +const SIZE = 1500; + +export function getFileEventsQuery({ + endpointIds, + size, + gte, + lte, +}: { + endpointIds: string[]; + size?: number; + gte?: DateMath; + lte?: DateMath; +}): SearchRequest { + return { + allow_no_indices: true, + fields: ['_id', 'agent.id', 'process.executable'], + query: { + bool: { + must: [ + { + terms: { + 'agent.id': endpointIds, + }, + }, + { + range: { + '@timestamp': { + gte: gte ?? 'now-24h', + lte: lte ?? 'now', + }, + }, + }, + ], + }, + }, + size: size ?? SIZE, + sort: [ + { + '@timestamp': { + order: 'desc', + }, + }, + ], + _source: false, + ignore_unavailable: true, + index: [FILE_EVENTS_INDEX_PATTERN], + }; +} 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..44e2ba925faf0 --- /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,232 @@ +/* + * 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 { mockFileEventsQueryResults } from '../../../../mock/mock_file_events_query_results'; +import { mockAnonymizedEventsReplacements } from '../../../../mock/mock_anonymized_events'; +import { getAnonymizedEvents } from './index'; + +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({ + ...mockFileEventsQueryResults, + hits: { + ...mockFileEventsQueryResults.hits, + hits: mockFileEventsQueryResults.hits.hits.map((hit) => ({ + ...hit, + fields: { ...mockRawData }, + })), + }, + } 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 undefined hits in response', async () => { + mockEsClient.search.mockResolvedValue({} as SearchResponse); + + const result = await getAnonymizedEvents({ + insightType: DefendInsightType.Enum.incompatible_antivirus, + endpointIds: ['test-endpoint'], + esClient: mockEsClient, + }); + + expect(result).toBeUndefined(); + }); + + 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, + }); + + // Verify getRawDataOrDefault is called with fields + expect(mockedGetRawDataOrDefault).toHaveBeenCalledWith(expect.objectContaining(mockRawData)); + }); + + 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({ + size, + }) + ); + }); + + 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({ hits: { hits: [] } } 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/elastic_assistant/server/lib/defend_insights/graphs/default_defend_insights_graph/nodes/retriever/helpers/get_anonymized_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 new file mode 100644 index 0000000000000..9c5ab1175efac --- /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.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 { DateMath, SearchResponse } from '@elastic/elasticsearch/lib/api/types'; +import type { ElasticsearchClient } from '@kbn/core/server'; +import { + DefendInsightType, + Replacements, + getAnonymizedValue, + getRawDataOrDefault, + transformRawData, +} from '@kbn/elastic-assistant-common'; + +import { AnonymizationFieldResponse } from '@kbn/elastic-assistant-common/impl/schemas/anonymization_fields/bulk_crud_anonymization_fields_route.gen'; +import { getQuery } from './get_events'; + +export const getAnonymizedEvents = async ({ + insightType, + endpointIds, + anonymizationFields, + esClient, + onNewReplacements, + replacements, + size, + start, + end, +}: { + insightType: DefendInsightType; + endpointIds: string[]; + anonymizationFields?: AnonymizationFieldResponse[]; + esClient: ElasticsearchClient; + onNewReplacements?: (replacements: Replacements) => void; + replacements?: Replacements; + size?: number; + start?: DateMath; + end?: DateMath; +}): Promise => { + if (insightType == null) { + return []; + } + + const query = getQuery(insightType, { endpointIds, size, gte: start, lte: end }); + const result = await esClient.search(query); + + let localReplacements = { ...(replacements ?? {}) }; + const localOnNewReplacements = (newReplacements: Replacements) => { + localReplacements = { ...localReplacements, ...newReplacements }; + + onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements + }; + + return result.hits?.hits?.map((x) => + transformRawData({ + anonymizationFields, + currentReplacements: localReplacements, // <-- the latest local replacements + getAnonymizedValue, + onNewReplacements: localOnNewReplacements, // <-- the local callback + rawData: getRawDataOrDefault(x.fields), + }) + ); +}; 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/lib/defend_insights/persistence/field_maps_configuration.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/field_maps_configuration.ts new file mode 100644 index 0000000000000..5769ab4557102 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/field_maps_configuration.ts @@ -0,0 +1,174 @@ +/* + * 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 { FieldMap } from '@kbn/data-stream-adapter'; + +export const defendInsightsFieldMap: FieldMap = { + '@timestamp': { + type: 'date', + array: false, + required: false, + }, + users: { + type: 'nested', + array: true, + required: false, + }, + 'users.id': { + type: 'keyword', + array: false, + required: true, + }, + 'users.name': { + type: 'keyword', + array: false, + required: false, + }, + id: { + type: 'keyword', + array: false, + required: true, + }, + last_viewed_at: { + type: 'date', + array: false, + required: true, + }, + updated_at: { + type: 'date', + array: false, + required: true, + }, + created_at: { + type: 'date', + array: false, + required: true, + }, + endpoint_ids: { + type: 'keyword', + array: true, + required: false, + }, + insight_type: { + type: 'keyword', + required: true, + }, + insights: { + type: 'nested', + array: true, + required: false, + }, + 'insights.group': { + type: 'keyword', + array: true, + required: true, + }, + 'insights.events': { + type: 'nested', + array: true, + required: false, + }, + 'insights.events.endpoint_id': { + type: 'keyword', + array: false, + required: true, + }, + 'insights.events.id': { + type: 'keyword', + array: false, + required: true, + }, + 'insights.events.value': { + type: 'text', + array: false, + required: true, + }, + replacements: { + type: 'object', + array: false, + required: false, + }, + 'replacements.value': { + type: 'keyword', + array: false, + required: false, + }, + 'replacements.uuid': { + type: 'keyword', + array: false, + required: false, + }, + api_config: { + type: 'object', + array: false, + required: true, + }, + 'api_config.connector_id': { + type: 'keyword', + array: false, + required: true, + }, + 'api_config.action_type_id': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.default_system_prompt_id': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.provider': { + type: 'keyword', + array: false, + required: false, + }, + 'api_config.model': { + type: 'keyword', + array: false, + required: false, + }, + events_context_count: { + type: 'integer', + array: false, + required: false, + }, + status: { + type: 'keyword', + array: false, + required: true, + }, + namespace: { + type: 'keyword', + array: false, + required: true, + }, + average_interval_ms: { + type: 'integer', + array: false, + required: false, + }, + failure_reason: { + type: 'keyword', + array: false, + required: false, + }, + generation_intervals: { + type: 'nested', + array: true, + required: false, + }, + 'generation_intervals.date': { + type: 'date', + array: false, + required: true, + }, + 'generation_intervals.duration_ms': { + type: 'integer', + array: false, + required: true, + }, +} as const; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts new file mode 100644 index 0000000000000..981c598335bf1 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.test.ts @@ -0,0 +1,68 @@ +/* + * 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 { 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 { getDefendInsight } from './get_defend_insight'; + +const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); +const mockLogger = loggerMock.create(); + +const mockResponse = getDefendInsightsSearchEsMock(); + +const user = { + username: 'test_user', + profile_uid: '1234', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, +} as AuthenticatedUser; +const mockRequest = { + esClient: mockEsClient, + index: 'defend-insights-index', + id: 'insight-id', + user, + logger: mockLogger, +}; +describe('getDefendInsight', () => { + afterEach(() => { + jest.clearAllMocks(); + }); + + it('should get defend insight by id successfully', async () => { + mockEsClient.search.mockResolvedValueOnce(mockResponse); + + const response = await getDefendInsight(mockRequest); + + expect(response).not.toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).not.toHaveBeenCalled(); + }); + + it('should return null if no defend insights found', async () => { + mockEsClient.search.mockResolvedValueOnce({ ...mockResponse, hits: { hits: [] } }); + + const response = await getDefendInsight(mockRequest); + + expect(response).toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).not.toHaveBeenCalled(); + }); + + it('should throw error on elasticsearch search failure', async () => { + mockEsClient.search.mockRejectedValueOnce(new Error('Elasticsearch error')); + + await expect(getDefendInsight(mockRequest)).rejects.toThrowError('Elasticsearch error'); + + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).toHaveBeenCalledTimes(1); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.ts new file mode 100644 index 0000000000000..4eeef2afd8738 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/get_defend_insight.ts @@ -0,0 +1,78 @@ +/* + * 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 { AuthenticatedUser, ElasticsearchClient, Logger } from '@kbn/core/server'; +import { DefendInsightsResponse } from '@kbn/elastic-assistant-common'; + +import { EsDefendInsightSchema } from './types'; +import { transformESSearchToDefendInsights } from './helpers'; + +export interface GetDefendInsightParams { + esClient: ElasticsearchClient; + logger: Logger; + index: string; + id: string; + user: AuthenticatedUser; +} + +export const getDefendInsight = async ({ + esClient, + logger, + index, + id, + user, +}: GetDefendInsightParams): Promise => { + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + try { + const response = await esClient.search({ + query: { + bool: { + must: [ + { + bool: { + should: [ + { + term: { + _id: id, + }, + }, + ], + }, + }, + ...filterByUser, + ], + }, + }, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + }); + const insights = transformESSearchToDefendInsights(response); + return insights[0] ?? null; + } catch (err) { + logger.error(`Error fetching Defend insight: ${err} with id: ${id}`); + throw err; + } +}; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.test.ts new file mode 100644 index 0000000000000..8e0793218154a --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.test.ts @@ -0,0 +1,64 @@ +/* + * 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 { DefendInsightsGetRequestQuery } from '@kbn/elastic-assistant-common'; + +import { DefendInsightType, DefendInsightStatus } from '@kbn/elastic-assistant-common'; + +import { queryParamsToEsQuery } from './helpers'; + +describe('defend insights data client helpers', () => { + describe('queryParamsToEsQuery', () => { + let queryParams: DefendInsightsGetRequestQuery; + let expectedQuery: object[]; + + function getDefaultQueryParams(): DefendInsightsGetRequestQuery { + return { + ids: ['insight-id1', 'insight-id2'], + endpoint_ids: ['endpoint-id1', 'endpoint-id2'], + connector_id: 'connector-id1', + type: DefendInsightType.Enum.incompatible_antivirus, + status: DefendInsightStatus.Enum.succeeded, + }; + } + + function getDefaultExpectedQuery(): object[] { + return [ + { terms: { _id: queryParams.ids } }, + { terms: { endpoint_ids: queryParams.endpoint_ids } }, + { term: { 'api_config.connector_id': queryParams.connector_id } }, + { term: { insight_type: queryParams.type } }, + { term: { status: queryParams.status } }, + ]; + } + + beforeEach(() => { + queryParams = getDefaultQueryParams(); + expectedQuery = getDefaultExpectedQuery(); + }); + + it('should correctly convert valid query parameters to Elasticsearch query format', () => { + const result = queryParamsToEsQuery(queryParams); + expect(result).toEqual(expectedQuery); + }); + + it('should ignore invalid query parameters', () => { + const badParams = { + ...queryParams, + invalid_param: 'invalid value', + }; + + const result = queryParamsToEsQuery(badParams); + expect(result).toEqual(expectedQuery); + }); + + it('should handle empty query parameters', () => { + const result = queryParamsToEsQuery({}); + expect(result).toEqual([]); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.ts new file mode 100644 index 0000000000000..b8164f53d9815 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/helpers.ts @@ -0,0 +1,221 @@ +/* + * 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 { get as _get, isArray } from 'lodash'; + +import type { estypes } from '@elastic/elasticsearch'; +import type { QueryDslQueryContainer } from '@elastic/elasticsearch/lib/api/types'; +import type { AuthenticatedUser } from '@kbn/core/server'; +import type { + DefendInsightCreateProps, + DefendInsightUpdateProps, + DefendInsightsResponse, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; + +import type { + CreateDefendInsightSchema, + EsDefendInsightSchema, + UpdateDefendInsightSchema, +} from './types'; + +export const transformESSearchToDefendInsights = ( + response: estypes.SearchResponse +): DefendInsightsResponse[] => { + return response.hits.hits + .filter((hit) => hit._source !== undefined) + .map((hit) => { + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + const insightSchema = hit._source!; + const defendInsight: DefendInsightsResponse = { + timestamp: insightSchema['@timestamp'], + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + id: hit._id!, + backingIndex: hit._index, + createdAt: insightSchema.created_at, + updatedAt: insightSchema.updated_at, + lastViewedAt: insightSchema.last_viewed_at, + users: + insightSchema.users?.map((user) => ({ + id: user.id, + name: user.name, + })) ?? [], + namespace: insightSchema.namespace, + status: insightSchema.status, + eventsContextCount: insightSchema.events_context_count, + apiConfig: { + connectorId: insightSchema.api_config.connector_id, + actionTypeId: insightSchema.api_config.action_type_id, + defaultSystemPromptId: insightSchema.api_config.default_system_prompt_id, + model: insightSchema.api_config.model, + provider: insightSchema.api_config.provider, + }, + endpointIds: insightSchema.endpoint_ids, + insightType: insightSchema.insight_type, + insights: insightSchema.insights.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpointId: event.endpoint_id, + value: event.value, + })), + })), + replacements: insightSchema.replacements?.reduce((acc: Record, r) => { + acc[r.uuid] = r.value; + return acc; + }, {}), + generationIntervals: + insightSchema.generation_intervals?.map((interval) => ({ + date: interval.date, + durationMs: interval.duration_ms, + })) ?? [], + averageIntervalMs: insightSchema.average_interval_ms ?? 0, + failureReason: insightSchema.failure_reason, + }; + + return defendInsight; + }); +}; + +export const transformToCreateScheme = ( + createdAt: string, + spaceId: string, + user: AuthenticatedUser, + { + endpointIds, + insightType, + insights, + apiConfig, + eventsContextCount, + replacements, + status, + }: DefendInsightCreateProps +): CreateDefendInsightSchema => { + return { + '@timestamp': createdAt, + created_at: createdAt, + users: [ + { + id: user.profile_uid, + name: user.username, + }, + ], + status, + api_config: { + action_type_id: apiConfig.actionTypeId, + connector_id: apiConfig.connectorId, + default_system_prompt_id: apiConfig.defaultSystemPromptId, + model: apiConfig.model, + provider: apiConfig.provider, + }, + events_context_count: eventsContextCount, + endpoint_ids: endpointIds, + insight_type: insightType, + insights: insights?.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpoint_id: event.endpointId, + value: event.value, + })), + })), + updated_at: createdAt, + last_viewed_at: createdAt, + replacements: replacements + ? Object.keys(replacements).map((key) => ({ + uuid: key, + value: replacements[key], + })) + : undefined, + namespace: spaceId, + }; +}; + +export const transformToUpdateScheme = ( + updatedAt: string, + { + eventsContextCount, + apiConfig, + insights, + failureReason, + generationIntervals, + id, + replacements, + lastViewedAt, + status, + }: DefendInsightUpdateProps +): UpdateDefendInsightSchema => { + const averageIntervalMsObj = + generationIntervals && generationIntervals.length > 0 + ? { + average_interval_ms: Math.trunc( + generationIntervals.reduce((acc, interval) => acc + interval.durationMs, 0) / + generationIntervals.length + ), + generation_intervals: generationIntervals.map((interval) => ({ + date: interval.date, + duration_ms: interval.durationMs, + })), + } + : {}; + return { + events_context_count: eventsContextCount, + ...(apiConfig + ? { + api_config: { + action_type_id: apiConfig.actionTypeId, + connector_id: apiConfig.connectorId, + default_system_prompt_id: apiConfig.defaultSystemPromptId, + model: apiConfig.model, + provider: apiConfig.provider, + }, + } + : {}), + ...(insights + ? { + insights: insights.map((insight) => ({ + group: insight.group, + events: insight.events?.map((event) => ({ + id: event.id, + endpoint_id: event.endpointId, + value: event.value, + })), + })), + } + : {}), + failure_reason: failureReason, + id, + replacements: replacements + ? Object.keys(replacements).map((key) => ({ + uuid: key, + value: replacements[key], + })) + : undefined, + ...(status ? { status } : {}), + // only update updated_at time if this is not an update to last_viewed_at + ...(lastViewedAt ? { last_viewed_at: lastViewedAt } : { updated_at: updatedAt }), + ...averageIntervalMsObj, + }; +}; + +const validParams = new Set(['ids', 'endpoint_ids', 'connector_id', 'type', 'status']); +const paramKeyMap = { ids: '_id', connector_id: 'api_config.connector_id', type: 'insight_type' }; +export function queryParamsToEsQuery( + queryParams: DefendInsightsGetRequestQuery +): QueryDslQueryContainer[] { + return Object.entries(queryParams).reduce((acc: object[], [k, v]) => { + if (!validParams.has(k)) { + return acc; + } + + const filterKey = isArray(v) ? 'terms' : 'term'; + const paramKey = _get(paramKeyMap, k, k); + const next = { [filterKey]: { [paramKey]: v } }; + + return [...acc, next]; + }, []); +} diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts new file mode 100644 index 0000000000000..0d2192c82c2e9 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.test.ts @@ -0,0 +1,408 @@ +/* + * 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. + */ + +/* eslint-disable @typescript-eslint/no-explicit-any */ + +import type { AuthenticatedUser } from '@kbn/core-security-common'; +import type { + DefendInsightCreateProps, + DefendInsightsUpdateProps, + 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 '../../../ai_assistant_data_clients'; +import { getDefendInsightsSearchEsMock } from '../../../__mocks__/defend_insights_schema.mock'; +import { getDefendInsight } from './get_defend_insight'; +import { + queryParamsToEsQuery, + transformESSearchToDefendInsights, + transformToUpdateScheme, +} from './helpers'; +import { DefendInsightsDataClient } from '.'; + +jest.mock('./get_defend_insight'); +jest.mock('./helpers', () => { + const original = jest.requireActual('./helpers'); + return { + ...original, + queryParamsToEsQuery: jest.fn(), + }; +}); + +describe('DefendInsightsDataClient', () => { + const mockEsClient = elasticsearchServiceMock.createElasticsearchClient(); + const mockLogger = loggerMock.create(); + const mockGetDefendInsight = jest.mocked(getDefendInsight); + let user: AuthenticatedUser; + let dataClientParams: AIAssistantDataClientParams; + let dataClient: DefendInsightsDataClient; + + function getDefaultUser(): AuthenticatedUser { + return { + username: 'test_user', + profile_uid: '1234', + authentication_realm: { + type: 'my_realm_type', + name: 'my_realm_name', + }, + } as AuthenticatedUser; + } + + function getDefaultDataClientParams(): AIAssistantDataClientParams { + return { + logger: mockLogger, + currentUser: user, + elasticsearchClientPromise: new Promise((resolve) => resolve(mockEsClient)), + indexPatternsResourceName: 'defend-insights-index', + kibanaVersion: '9.0.0', + spaceId: 'space-1', + } as AIAssistantDataClientParams; + } + + beforeEach(() => { + user = getDefaultUser(); + dataClientParams = getDefaultDataClientParams(); + dataClient = new DefendInsightsDataClient(dataClientParams); + }); + + afterEach(() => { + jest.clearAllMocks(); + }); + + describe('getDefendInsight', () => { + it('should correctly get defend insight', async () => { + const id = 'some-id'; + mockGetDefendInsight.mockResolvedValueOnce({ id } as DefendInsightsResponse); + const response = await dataClient.getDefendInsight({ id, authenticatedUser: user }); + + expect(mockGetDefendInsight).toHaveBeenCalledTimes(1); + expect(response).not.toBeNull(); + expect(response!.id).toEqual(id); + }); + }); + + describe('createDefendInsight', () => { + const defendInsightCreate: DefendInsightCreateProps = { + endpointIds: [], + insightType: DefendInsightType.Enum.incompatible_antivirus, + insights: [], + apiConfig: { + actionTypeId: 'action-type-id', + connectorId: 'connector-id', + defaultSystemPromptId: 'default-prompt-id', + model: 'model-name', + provider: 'OpenAI', + }, + eventsContextCount: 10, + replacements: { key1: 'value1', key2: 'value2' }, + status: DefendInsightStatus.Enum.running, + }; + + it('should create defend insight successfully', async () => { + const id = 'created-id'; + // @ts-expect-error not full response interface + mockEsClient.create.mockResolvedValueOnce({ _id: id }); + mockGetDefendInsight.mockResolvedValueOnce({ id } as DefendInsightsResponse); + + const response = await dataClient.createDefendInsight({ + defendInsightCreate, + authenticatedUser: user, + }); + expect(mockEsClient.create).toHaveBeenCalledTimes(1); + expect(mockGetDefendInsight).toHaveBeenCalledTimes(1); + expect(response).not.toBeNull(); + expect(response!.id).toEqual(id); + }); + + it('should throw error on elasticsearch create failure', async () => { + mockEsClient.create.mockRejectedValueOnce(new Error('Elasticsearch error')); + const responsePromise = dataClient.createDefendInsight({ + defendInsightCreate, + authenticatedUser: user, + }); + await expect(responsePromise).rejects.toThrowError('Elasticsearch error'); + expect(mockEsClient.create).toHaveBeenCalledTimes(1); + expect(mockGetDefendInsight).not.toHaveBeenCalled(); + }); + }); + + describe('findDefendInsightsByParams', () => { + let mockQueryParamsToEsQuery: Function; + let queryParams: DefendInsightsGetRequestQuery; + let expectedTermFilters: object[]; + + function getDefaultQueryParams() { + return { + ids: ['insight-id1', 'insight-id2'], + endpoint_ids: ['endpoint-id1', 'endpoint-id2'], + connector_id: 'connector-id1', + type: DefendInsightType.Enum.incompatible_antivirus, + status: DefendInsightStatus.Enum.succeeded, + }; + } + + function getDefaultExpectedTermFilters() { + return [ + { terms: { _id: queryParams.ids } }, + { terms: { endpoint_ids: queryParams.endpoint_ids } }, + { term: { 'api_config.connector_id': queryParams.connector_id } }, + { term: { insight_type: queryParams.type } }, + { term: { status: queryParams.status } }, + ]; + } + + beforeEach(() => { + queryParams = getDefaultQueryParams(); + expectedTermFilters = getDefaultExpectedTermFilters(); + mockQueryParamsToEsQuery = jest + .mocked(queryParamsToEsQuery) + .mockReturnValueOnce(expectedTermFilters); + }); + + it('should return defend insights successfully', async () => { + const mockResponse = getDefendInsightsSearchEsMock(); + mockEsClient.search.mockResolvedValueOnce(mockResponse); + + const result = await dataClient.findDefendInsightsByParams({ + params: queryParams, + authenticatedUser: user, + }); + const expectedResult = transformESSearchToDefendInsights(mockResponse); + + expect(mockQueryParamsToEsQuery).toHaveBeenCalledTimes(1); + expect(mockQueryParamsToEsQuery).toHaveBeenCalledWith(queryParams); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockEsClient.search).toHaveBeenCalledWith( + expect.objectContaining({ + index: `${dataClientParams.indexPatternsResourceName}-${dataClientParams.spaceId}`, + size: 10, + sort: [ + { + '@timestamp': 'desc', + }, + ], + query: { + bool: { + must: [ + ...expectedTermFilters, + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: { 'users.id': user.profile_uid }, + }, + ], + }, + }, + }, + }, + ], + }, + }, + }) + ); + expect(result).toEqual(expectedResult); + }); + + it('should log and throw an error if search fails', async () => { + const mockError = new Error('Search failed'); + mockEsClient.search.mockRejectedValue(mockError); + + await expect( + dataClient.findDefendInsightsByParams({ + params: queryParams, + authenticatedUser: user, + }) + ).rejects.toThrow(mockError); + expect(mockLogger.error).toHaveBeenCalledWith( + `error fetching Defend insights: ${mockError} with params: ${JSON.stringify(queryParams)}` + ); + }); + }); + + describe('findAllDefendInsights', () => { + it('should correctly query ES', async () => { + const mockResponse = getDefendInsightsSearchEsMock(); + mockEsClient.search.mockResolvedValueOnce(mockResponse); + const searchParams = { + query: { + bool: { + must: [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: { 'users.id': user.profile_uid }, + }, + ], + }, + }, + }, + }, + ], + }, + }, + size: 10000, + _source: true, + ignore_unavailable: true, + index: `${dataClientParams.indexPatternsResourceName}-${dataClientParams.spaceId}`, + seq_no_primary_term: true, + }; + + const response = await dataClient.findAllDefendInsights({ + authenticatedUser: user, + }); + expect(response).not.toBeNull(); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockEsClient.search).toHaveBeenCalledWith(searchParams); + }); + + it('should throw error on elasticsearch search failure', async () => { + mockEsClient.search.mockRejectedValueOnce(new Error('Elasticsearch error')); + await expect( + dataClient.findAllDefendInsights({ + authenticatedUser: user, + }) + ).rejects.toThrowError('Elasticsearch error'); + expect(mockEsClient.search).toHaveBeenCalledTimes(1); + expect(mockLogger.error).toHaveBeenCalledTimes(1); + }); + }); + + describe('updateDefendInsights', () => { + let defendInsightsUpdateProps: DefendInsightsUpdateProps; + + function getDefaultProps() { + return [ + { + id: 'insight-id1', + backingIndex: 'defend-insights-index', + status: DefendInsightStatus.Enum.succeeded, + insights: [ + { + group: 'windows_defender', + events: [ + { + id: 'event-id-1', + endpointId: 'endpoint-id-1', + value: '/windows/defender/scan.exe', + }, + ], + }, + ], + }, + ]; + } + + beforeEach(async () => { + defendInsightsUpdateProps = getDefaultProps(); + }); + + it('should update defend insights successfully', async () => { + // ensure startTime is before updatedAt timestamp + const startTime = new Date().getTime() - 1; + const mockResponse: DefendInsightsResponse[] = [ + { id: defendInsightsUpdateProps[0].id } as DefendInsightsResponse, + ]; + + const findDefendInsightsByParamsSpy = jest.spyOn(dataClient, 'findDefendInsightsByParams'); + findDefendInsightsByParamsSpy.mockResolvedValueOnce(mockResponse); + + const result = await dataClient.updateDefendInsights({ + defendInsightsUpdateProps, + authenticatedUser: user, + }); + const expectedDoc = transformToUpdateScheme('', defendInsightsUpdateProps[0]); + delete expectedDoc.updated_at; + + expect(mockEsClient.bulk).toHaveBeenCalledTimes(1); + expect(mockEsClient.bulk).toHaveBeenCalledWith({ + body: [ + { + update: { + _index: defendInsightsUpdateProps[0].backingIndex, + _id: defendInsightsUpdateProps[0].id, + }, + }, + { + doc: expect.objectContaining({ ...expectedDoc }), + }, + ], + refresh: 'wait_for', + }); + const updatedAt = (mockEsClient.bulk.mock.calls[0][0] as { body: any[] }).body[1].doc + .updated_at; + expect(new Date(updatedAt).getTime()).toBeGreaterThan(startTime); + expect(dataClient.findDefendInsightsByParams).toHaveBeenCalledTimes(1); + expect(dataClient.findDefendInsightsByParams).toHaveBeenCalledWith({ + params: { ids: [defendInsightsUpdateProps[0].id] }, + authenticatedUser: user, + }); + expect(result).toEqual(mockResponse); + }); + + it('should log a warning and throw an error if update fails', async () => { + const mockError = new Error('Update failed'); + mockEsClient.bulk.mockRejectedValue(mockError); + + await expect( + dataClient.updateDefendInsights({ + defendInsightsUpdateProps, + authenticatedUser: user, + }) + ).rejects.toThrow(mockError); + + expect(mockLogger.warn).toHaveBeenCalledWith( + `error updating Defend insights: ${mockError} for IDs: ${defendInsightsUpdateProps[0].id}` + ); + }); + }); + + describe('updateDefendInsight', () => { + it('correctly calls updateDefendInsights', async () => { + const defendInsightUpdateProps = { + id: 'insight-id1', + backingIndex: 'defend-insights-index', + status: DefendInsightStatus.Enum.succeeded, + insights: [ + { + group: 'windows_defender', + events: [ + { + id: 'event-id-1', + endpointId: 'endpoint-id-1', + value: '/windows/defender/scan.exe', + }, + ], + }, + ], + }; + const updateDefendInsightsSpy = jest.spyOn(dataClient, 'updateDefendInsights'); + updateDefendInsightsSpy.mockResolvedValueOnce([]); + await dataClient.updateDefendInsight({ + defendInsightUpdateProps, + authenticatedUser: user, + }); + + expect(updateDefendInsightsSpy).toHaveBeenCalledTimes(1); + expect(updateDefendInsightsSpy).toHaveBeenCalledWith({ + defendInsightsUpdateProps: [defendInsightUpdateProps], + authenticatedUser: user, + }); + }); + }); +}); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts new file mode 100644 index 0000000000000..b4ef64b53f18a --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/index.ts @@ -0,0 +1,284 @@ +/* + * 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 { + DefendInsightCreateProps, + DefendInsightUpdateProps, + DefendInsightsUpdateProps, + DefendInsightsResponse, + DefendInsightsGetRequestQuery, +} from '@kbn/elastic-assistant-common'; +import type { AuthenticatedUser } from '@kbn/core-security-common'; +import { v4 as uuidv4 } from 'uuid'; + +import type { AIAssistantDataClientParams } from '../../../ai_assistant_data_clients'; +import type { EsDefendInsightSchema } from './types'; +import { AIAssistantDataClient } from '../../../ai_assistant_data_clients'; +import { getDefendInsight } from './get_defend_insight'; +import { + queryParamsToEsQuery, + transformESSearchToDefendInsights, + transformToCreateScheme, + transformToUpdateScheme, +} from './helpers'; + +const DEFAULT_PAGE_SIZE = 10; + +export class DefendInsightsDataClient extends AIAssistantDataClient { + constructor(public readonly options: AIAssistantDataClientParams) { + super(options); + } + + /** + * Fetches a Defend insight + * @param options + * @param options.id The existing Defend insight id. + * @param options.authenticatedUser Current authenticated user. + * @returns The Defend insight response + */ + public getDefendInsight = async ({ + id, + authenticatedUser, + }: { + id: string; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + return getDefendInsight({ + esClient, + logger: this.options.logger, + index: this.indexTemplateAndPattern.alias, + id, + user: authenticatedUser, + }); + }; + + /** + * Creates a Defend insight, if given at least the "apiConfig" + * @param options + * @param options.defendInsightCreate + * @param options.authenticatedUser + * @returns The Defend insight created + */ + public createDefendInsight = async ({ + defendInsightCreate, + authenticatedUser, + }: { + defendInsightCreate: DefendInsightCreateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const id = defendInsightCreate?.id || uuidv4(); + const createdAt = new Date().toISOString(); + + const body = transformToCreateScheme(createdAt, this.spaceId, user, defendInsightCreate); + try { + const response = await esClient.create({ + body, + id, + index, + refresh: 'wait_for', + }); + + const createdDefendInsight = await getDefendInsight({ + esClient, + index, + id: response._id, + logger, + user, + }); + return createdDefendInsight; + } catch (err) { + logger.error(`error creating Defend insight: ${err} with id: ${id}`); + throw err; + } + }; + + /** + * Find Defend insights by params + * @param options + * @param options.params + * @param options.authenticatedUser + * @returns The Defend insights found + */ + public findDefendInsightsByParams = async ({ + params, + authenticatedUser, + }: { + params: DefendInsightsGetRequestQuery; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const termFilters = queryParamsToEsQuery(params); + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + + try { + const query = { + bool: { + must: [...termFilters, ...filterByUser], + }, + }; + const response = await esClient.search({ + query, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + sort: [{ '@timestamp': 'desc' }], + size: params.size || DEFAULT_PAGE_SIZE, + }); + return transformESSearchToDefendInsights(response); + } catch (err) { + logger.error(`error fetching Defend insights: ${err} with params: ${JSON.stringify(params)}`); + throw err; + } + }; + + /** + * Finds all Defend insight for authenticated user + * @param options + * @param options.authenticatedUser + * @returns The Defend insight + */ + public findAllDefendInsights = async ({ + authenticatedUser, + }: { + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const index = this.indexTemplateAndPattern.alias; + const user = authenticatedUser; + const MAX_ITEMS = 10000; + const filterByUser = [ + { + nested: { + path: 'users', + query: { + bool: { + must: [ + { + match: user.profile_uid + ? { 'users.id': user.profile_uid } + : { 'users.name': user.username }, + }, + ], + }, + }, + }, + }, + ]; + + try { + const response = await esClient.search({ + query: { + bool: { + must: [...filterByUser], + }, + }, + size: MAX_ITEMS, + _source: true, + ignore_unavailable: true, + index, + seq_no_primary_term: true, + }); + const insights = transformESSearchToDefendInsights(response); + return insights ?? []; + } catch (err) { + logger.error(`error fetching Defend insights: ${err}`); + throw err; + } + }; + + /** + * Updates Defend insights + * @param options + * @param options.defendInsightsUpdateProps + * @param options.authenticatedUser + */ + public updateDefendInsights = async ({ + defendInsightsUpdateProps, + authenticatedUser, + }: { + defendInsightsUpdateProps: DefendInsightsUpdateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + const esClient = await this.options.elasticsearchClientPromise; + const logger = this.options.logger; + const updatedAt = new Date().toISOString(); + + let ids: string[] = []; + const bulkParams = defendInsightsUpdateProps.flatMap((updateProp) => { + const index = updateProp.backingIndex; + const params = transformToUpdateScheme(updatedAt, updateProp); + ids = [...ids, params.id]; + return [ + { + update: { + _index: index, + _id: params.id, + }, + }, + { + doc: params, + }, + ]; + }); + + try { + await esClient.bulk({ body: bulkParams, refresh: 'wait_for' }); + return this.findDefendInsightsByParams({ params: { ids }, authenticatedUser }); + } catch (err) { + logger.warn(`error updating Defend insights: ${err} for IDs: ${ids}`); + throw err; + } + }; + + /** + * Updates a Defend insight + * @param options + * @param options.defendInsightUpdateProps + * @param options.authenticatedUser + */ + public updateDefendInsight = async ({ + defendInsightUpdateProps, + authenticatedUser, + }: { + defendInsightUpdateProps: DefendInsightUpdateProps; + authenticatedUser: AuthenticatedUser; + }): Promise => { + return ( + await this.updateDefendInsights({ + defendInsightsUpdateProps: [defendInsightUpdateProps], + authenticatedUser, + }) + )[0]; + }; +} diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts new file mode 100644 index 0000000000000..87b1c8edbd776 --- /dev/null +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights/persistence/types.ts @@ -0,0 +1,88 @@ +/* + * 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 { + DefendInsightStatus, + DefendInsightType, + Provider, + UUID, +} from '@kbn/elastic-assistant-common'; + +import type { EsReplacementSchema } from '../../../ai_assistant_data_clients/conversations/types'; + +interface DefendInsightInsightEventSchema { + id: string; + endpoint_id: string; + value: string; +} + +interface DefendInsightInsightSchema { + group: string; + events?: DefendInsightInsightEventSchema[]; +} + +interface BaseDefendInsightSchema { + '@timestamp': string; + created_at: string; + updated_at: string; + last_viewed_at: string; + status: DefendInsightStatus; + events_context_count?: number; + endpoint_ids: string[]; + insight_type: DefendInsightType; + insights: DefendInsightInsightSchema[]; + api_config: { + connector_id: string; + action_type_id: string; + default_system_prompt_id?: string; + provider?: Provider; + model?: string; + }; + replacements?: EsReplacementSchema[]; +} + +export interface EsDefendInsightSchema extends BaseDefendInsightSchema { + id: string; + namespace: string; + failure_reason?: string; + users?: Array<{ + id?: string; + name?: string; + }>; + average_interval_ms?: number; + generation_intervals?: Array<{ date: string; duration_ms: number }>; +} + +export interface CreateDefendInsightSchema extends BaseDefendInsightSchema { + id?: string | undefined; + users: Array<{ + id?: string; + name?: string; + }>; + namespace: string; +} + +export interface UpdateDefendInsightSchema { + id: UUID; + '@timestamp'?: string; + updated_at?: string; + last_viewed_at?: string; + status?: DefendInsightStatus; + events_context_count?: number; + insights?: DefendInsightInsightSchema[]; + api_config?: { + action_type_id?: string; + connector_id?: string; + default_system_prompt_id?: string; + provider?: Provider; + model?: string; + }; + replacements?: EsReplacementSchema[]; + average_interval_ms?: number; + generation_intervals?: Array<{ date: string; duration_ms: number }>; + failure_reason?: 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..810f44658f919 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_TOOL_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_TOOL_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_TOOL_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/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts index 88ba28c06b8c7..78b3d2d703295 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,15 +16,15 @@ 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'; @@ -37,24 +36,19 @@ import { 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 { 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[], @@ -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_TOOL_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, + } + ); + 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.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts index c8ba9e9819b19..3792dc88b475b 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/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 a979e1b20aba9..f3413c686804e 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'; 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/scripts/endpoint/endpoint_agent_runner/elastic_endpoint.ts b/x-pack/solutions/security/plugins/security_solution/scripts/endpoint/endpoint_agent_runner/elastic_endpoint.ts index e18f91457eca7..197d9e088d8db 100644 --- a/x-pack/solutions/security/plugins/security_solution/scripts/endpoint/endpoint_agent_runner/elastic_endpoint.ts +++ b/x-pack/solutions/security/plugins/security_solution/scripts/endpoint/endpoint_agent_runner/elastic_endpoint.ts @@ -52,6 +52,7 @@ export const enrollEndpointHost = async (): Promise => { version, useClosestVersionMatch: false, disk: '8G', + memory: '4G', }); log.info(hostVm.info()); 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/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts deleted file mode 100644 index eef2e1ad28f16..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/get_file_events_query.ts +++ /dev/null @@ -1,64 +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 { SearchRequest } from '@elastic/elasticsearch/lib/api/types'; - -import { FILE_EVENTS_INDEX_PATTERN } from '../../../../../common/endpoint/constants'; - -const SIZE = 200; - -export function getFileEventsQuery({ endpointIds }: { endpointIds: string[] }): SearchRequest { - return { - allow_no_indices: true, - query: { - bool: { - must: [ - { - terms: { - 'agent.id': endpointIds, - }, - }, - { - range: { - '@timestamp': { - gte: 'now-24h', - lte: 'now', - }, - }, - }, - ], - }, - }, - size: 0, // Aggregations only - aggs: { - unique_process_executable: { - terms: { - field: 'process.executable', - size: SIZE, - }, - aggs: { - // Get the latest event for each process.executable - latest_event: { - top_hits: { - size: 1, - sort: [ - { - '@timestamp': { - order: 'desc', - }, - }, - ], - _source: ['_id', 'agent.id', 'process.executable'], // Include only necessary fields - }, - }, - }, - }, - }, - ignore_unavailable: true, - index: [FILE_EVENTS_INDEX_PATTERN], - }; -} 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/get_events/index.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts deleted file mode 100644 index 4f120e8e655f0..0000000000000 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/get_events/index.ts +++ /dev/null @@ -1,123 +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 { SearchRequest } 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, - getRawDataOrDefault, -} from '@kbn/elastic-assistant-common'; - -import { getFileEventsQuery } from './get_file_events_query'; -import { InvalidDefendInsightTypeError } from '../errors'; - -interface AggregationResponse { - unique_process_executable: { - buckets: Array<{ - key: string; - doc_count: number; - latest_event: { - hits: { - hits: Array<{ - _id: string; - _source: { - agent: { id: string }; - process: { executable: string }; - }; - }>; - }; - }; - }>; - }; -} - -export async function getAnonymizedEvents({ - endpointIds, - type, - anonymizationFields, - esClient, - onNewReplacements, - replacements, -}: { - 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, - }); - } - - 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 fileEvents = (result.aggregations?.unique_process_executable.buckets ?? []).map( - (bucket) => { - const latestEvent = bucket.latest_event.hits.hits[0]; - return { - _id: [latestEvent._id], - 'agent.id': [latestEvent._source.agent.id], - 'process.executable': [latestEvent._source.process.executable], - }; - } - ); - - // 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 }; - - onNewReplacements?.(localReplacements); // invoke the callback with the latest replacements - }; - - return fileEvents.map((fileEvent) => - transformRawData({ - anonymizationFields, - currentReplacements: localReplacements, // <-- the latest local replacements - getAnonymizedValue, - onNewReplacements: localOnNewReplacements, // <-- the local callback - rawData: getRawDataOrDefault(fileEvent), - }) - ); -}; 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.ts b/x-pack/solutions/security/plugins/security_solution/server/endpoint/services/workflow_insights/index.ts index 29e060298214b..5a892776047f4 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({ @@ -247,6 +252,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 dea527b73fd76..257afe40506d8 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/plugin.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/plugin.ts @@ -731,6 +731,7 @@ export class Plugin implements ISecuritySolutionPlugin { securityWorkflowInsightsService .start({ esClient: core.elasticsearch.client.asInternalUser, + registerDefendInsightsCallback: plugins.elasticAssistant.registerCallback, }) .catch(() => {});