[Defend Insights] langgraph upgrade

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));

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/errors.ts → 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');
   }

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;

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');
+  });
+});

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';

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');
+  });
+});

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;
+};

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');
+  });
+});

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';
+  }
+};