[Security Solution] defend insights langgraph upgrade (elastic#211038)

## Summary

This is intended to be a "minimal" migration for Defend Insights to
langgraph + output chunking. Other than the increased events due to the
context increase from output chunking, the functionality is unchanged.

* migrates defend insights to langgraph
* adds output chunking / refinement


### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

Author: joeypoon

.github/CODEOWNERS

@@ -2507,6 +2507,8 @@ x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/defend
 x-pack/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts @elastic/security-defend-workflows
 x-pack/plugins/elastic_assistant/server/ai_assistant_data_clients/defend_insights @elastic/security-defend-workflows
 x-pack/plugins/elastic_assistant/server/routes/defend_insights @elastic/security-defend-workflows
+x-pack/solutions/security/plugins/elastic_assistant/server/lib/defend_insights @elastic/security-defend-workflows
+x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights @elastic/security-defend-workflows
 /x-pack/solutions/security/plugins/security_solution/public/common/components/response_actions @elastic/security-defend-workflows
 /x-pack/solutions/security/plugins/security_solution_serverless/public/upselling/pages/osquery_automated_response_actions.tsx @elastic/security-defend-workflows
 

x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts

@@ -57,6 +57,6 @@ export const ELASTIC_AI_ASSISTANT_EVALUATE_URL =
   `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/evaluate` as const;
 
 // Defend insights
-export const DEFEND_INSIGHTS_TOOL_ID = 'defend-insights';
+export const DEFEND_INSIGHTS_ID = 'defend-insights';
 export const DEFEND_INSIGHTS = `${ELASTIC_AI_ASSISTANT_INTERNAL_URL}/defend_insights`;
 export const DEFEND_INSIGHTS_BY_ID = `${DEFEND_INSIGHTS}/{id}`;

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/alerts/helpers/get_raw_data_or_default/index.ts

@@ -5,8 +5,8 @@
  * 2.0.
  */
 
-import { isRawDataValid } from '../is_raw_data_valid';
 import type { MaybeRawData } from '../types';
+import { isRawDataValid } from '../is_raw_data_valid';
 
 /** Returns the raw data if it valid, or a default if it's not */
 export const getRawDataOrDefault = (rawData: MaybeRawData): Record<string, unknown[]> =>

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/defend_insights_schema.mock.ts

@@ -8,7 +8,7 @@
 import type { estypes } from '@elastic/elasticsearch';
 import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common';
 
-import type { EsDefendInsightSchema } from '../ai_assistant_data_clients/defend_insights/types';
+import type { EsDefendInsightSchema } from '../lib/defend_insights/persistence/types';
 
 export const getDefendInsightsSearchEsMock = () => {
   const searchResponse: estypes.SearchResponse<EsDefendInsightSchema> = {

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/raw_defend_insights.ts

@@ -0,0 +1,108 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { DefendInsightStatus, DefendInsightType } from '@kbn/elastic-assistant-common';
+
+import type { EsDefendInsightSchema } from '../lib/defend_insights/persistence/types';
+
+export const getParsedDefendInsightsMock = (timestamp: string): EsDefendInsightSchema[] => [
+  {
+    '@timestamp': timestamp,
+    created_at: timestamp,
+    updated_at: timestamp,
+    last_viewed_at: timestamp,
+    users: [
+      {
+        id: 'u_mGBROF_q5bmFCATbLXAcCwKa0k8JvONAwSruelyKA5E_0',
+        name: 'elastic',
+      },
+    ],
+    status: DefendInsightStatus.Enum.succeeded,
+    api_config: {
+      action_type_id: '.bedrock',
+      connector_id: 'ac4e19d1-e2e2-49af-bf4b-59428473101c',
+      model: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+    },
+    endpoint_ids: ['6e09ec1c-644c-4148-a02d-be451c35400d'],
+    insight_type: DefendInsightType.Enum.incompatible_antivirus,
+    insights: [
+      {
+        group: 'windows_defenders',
+        events: [],
+      },
+    ],
+    namespace: 'default',
+    id: '655c52ec-49ee-4d20-87e5-7edd6d8f84e8',
+    generation_intervals: [
+      {
+        date: timestamp,
+        duration_ms: 13113,
+      },
+    ],
+    average_interval_ms: 13113,
+    events_context_count: 100,
+    replacements: [
+      {
+        uuid: '2009c67b-89b8-43d9-b502-2c32f71875a0',
+        value: 'root',
+      },
+      {
+        uuid: '9f7f91b6-6853-48b7-bfb8-403f5efb2364',
+        value: 'joey-dev-default-3539',
+      },
+    ],
+  },
+  {
+    '@timestamp': timestamp,
+    created_at: timestamp,
+    updated_at: timestamp,
+    last_viewed_at: timestamp,
+    users: [
+      {
+        id: '00468e82-e37f-4224-80c1-c62e594c74b1',
+        name: 'ubuntu',
+      },
+    ],
+    status: DefendInsightStatus.Enum.succeeded,
+    api_config: {
+      action_type_id: '.bedrock',
+      connector_id: 'bc5e19d1-e2e2-49af-bf4b-59428473101d',
+      model: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+    },
+    endpoint_ids: ['b557bb12-8206-44b6-b2a5-dbcce5b1e65e'],
+    insight_type: DefendInsightType.Enum.noisy_process_tree,
+    insights: [
+      {
+        group: 'linux_security',
+        events: [],
+      },
+    ],
+    namespace: 'default',
+    id: '7a1b52ec-49ee-4d20-87e5-7edd6d8f84e9',
+    generation_intervals: [
+      {
+        date: timestamp,
+        duration_ms: 13113,
+      },
+    ],
+    average_interval_ms: 13113,
+    events_context_count: 100,
+    replacements: [
+      {
+        uuid: '3119c67b-89b8-43d9-b502-2c32f71875b1',
+        value: 'ubuntu',
+      },
+      {
+        uuid: '8e7f91b6-6853-48b7-bfb8-403f5efb2365',
+        value: 'ubuntu-dev-default-3540',
+      },
+    ],
+  },
+];
+
+export const getRawDefendInsightsMock = (timestamp: string) =>
+  JSON.stringify(getParsedDefendInsightsMock(timestamp));

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request_context.ts

@@ -29,7 +29,7 @@ import {
 } from '../ai_assistant_data_clients/knowledge_base';
 import { defaultAssistantFeatures } from '@kbn/elastic-assistant-common';
 import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
-import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights';
+import { DefendInsightsDataClient } from '../lib/defend_insights/persistence';
 import { authenticatedUser } from './user';
 
 export const createMockClients = () => {

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_service/index.ts

@@ -21,7 +21,7 @@ import { omit } from 'lodash';
 import { InstallationStatus } from '@kbn/product-doc-base-plugin/common/install_status';
 import { TrainedModelsProvider } from '@kbn/ml-plugin/server/shared_services/providers';
 import { attackDiscoveryFieldMap } from '../lib/attack_discovery/persistence/field_maps_configuration/field_maps_configuration';
-import { defendInsightsFieldMap } from '../ai_assistant_data_clients/defend_insights/field_maps_configuration';
+import { defendInsightsFieldMap } from '../lib/defend_insights/persistence/field_maps_configuration';
 import { getDefaultAnonymizationFields } from '../../common/anonymization';
 import { AssistantResourceNames, GetElser } from '../types';
 import {
@@ -49,7 +49,7 @@ import {
   GetAIAssistantKnowledgeBaseDataClientParams,
 } from '../ai_assistant_data_clients/knowledge_base';
 import { AttackDiscoveryDataClient } from '../lib/attack_discovery/persistence';
-import { DefendInsightsDataClient } from '../ai_assistant_data_clients/defend_insights';
+import { DefendInsightsDataClient } from '../lib/defend_insights/persistence';
 import { createGetElserId, ensureProductDocumentationInstalled } from './helpers';
 import { hasAIAssistantLicense } from '../routes/helpers';
 

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,
+      hasZeroEvents: true,
+      maxHallucinationFailuresReached: true,
+      maxRetriesReached: true,
+    });
+
+    expect(result).toEqual('end');
+  });
+
+  it("returns 'refine' if hasUnrefinedResults is true and getShouldEnd returns false", () => {
+    const result = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults: true,
+      hasZeroEvents: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toEqual('refine');
+  });
+
+  it("returns 'generate' if hasUnrefinedResults is false and getShouldEnd returns false", () => {
+    const result = getGenerateOrRefineOrEndDecision({
+      hasUnrefinedResults: false,
+      hasZeroEvents: false,
+      maxHallucinationFailuresReached: false,
+      maxRetriesReached: false,
+    });
+
+    expect(result).toEqual('generate');
+  });
+});