[Security AI] Add Kibana Support for Security AI Prompts Integration (#…

…207138)

packages/kbn-check-mappings-update-cli/current_fields.json

@@ -914,6 +914,15 @@
     "username"
   ],
   "search-telemetry": [],
+  "security-ai-prompt": [
+    "description",
+    "model",
+    "prompt",
+    "prompt.default",
+    "promptGroupId",
+    "promptId",
+    "provider"
+  ],
   "security-rule": [
     "rule_id",
     "version"

packages/kbn-check-mappings-update-cli/current_mappings.json

@@ -3023,6 +3023,33 @@
     "dynamic": false,
     "properties": {}
   },
+  "security-ai-prompt": {
+    "dynamic": false,
+    "properties": {
+      "description": {
+        "type": "text"
+      },
+      "model": {
+        "type": "keyword"
+      },
+      "prompt": {
+        "properties": {
+          "default": {
+            "type": "text"
+          }
+        }
+      },
+      "promptGroupId": {
+        "type": "keyword"
+      },
+      "promptId": {
+        "type": "keyword"
+      },
+      "provider": {
+        "type": "keyword"
+      }
+    }
+  },
   "security-rule": {
     "dynamic": false,
     "properties": {

src/core/packages/saved-objects/server-internal/src/object_types/index.ts

@@ -11,4 +11,4 @@ export { registerCoreObjectTypes } from './registration';
 
 // set minimum number of registered saved objects to ensure no object types are removed after 8.8
 // declared in internal implementation exclicilty to prevent unintended changes.
-export const SAVED_OBJECT_TYPES_COUNT = 127 as const;
+export const SAVED_OBJECT_TYPES_COUNT = 128 as const;

src/core/server/integration_tests/ci_checks/saved_objects/check_registered_types.test.ts

@@ -157,6 +157,7 @@ describe('checking migration metadata changes on all registered SO types', () =>
         "search": "0aa6eefb37edd3145be340a8b67779c2ca578b22",
         "search-session": "b2fcd840e12a45039ada50b1355faeafa39876d1",
         "search-telemetry": "b568601618744720b5662946d3103e3fb75fe8ee",
+        "security-ai-prompt": "cc8ee5aaa9d001e89c131bbd5af6bc80bc271046",
         "security-rule": "07abb4d7e707d91675ec0495c73816394c7b521f",
         "security-solution-signals-migration": "9d99715fe5246f19de2273ba77debd2446c36bb1",
         "siem-detection-engine-rule-actions": "54f08e23887b20da7c805fab7c60bc67c428aff9",

src/core/server/integration_tests/saved_objects/registration/type_registrations.test.ts

@@ -123,6 +123,7 @@ const previouslyRegisteredTypes = [
   'search',
   'search-session',
   'search-telemetry',
+  'security-ai-prompt',
   'security-rule',
   'security-solution-signals-migration',
   'risk-engine-configuration',

x-pack/platform/packages/shared/ai-infra/inference-common/index.ts

@@ -96,6 +96,7 @@ export {
   isInferenceRequestError,
   isInferenceRequestAbortedError,
 } from './src/errors';
+export { elasticModelDictionary } from './src/const';
 
 export { truncateList } from './src/truncate_list';
 export {

x-pack/platform/packages/shared/ai-infra/inference-common/src/const.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.
+ */
+
+import { ElasticModelDictionary } from './types';
+
+export const elasticModelDictionary: ElasticModelDictionary = {
+  'rainbow-sprinkles': {
+    provider: 'bedrock',
+    model: 'us.anthropic.claude-3-5-sonnet-20240620-v1:0',
+  },
+};

x-pack/platform/packages/shared/ai-infra/inference-common/src/types.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.
+ */
+
+export interface ElasticModelDictionary {
+  [key: string]: {
+    provider: string;
+    model: string;
+  };
+}

x-pack/platform/plugins/shared/fleet/.storybook/context/fixtures/integration.nginx.ts

@@ -251,6 +251,7 @@ export const item: GetInfoResponse['item'] = {
       index_pattern: [],
       lens: [],
       map: [],
+      security_ai_prompt: [],
       security_rule: [],
       csp_rule_template: [],
       tag: [],

x-pack/platform/plugins/shared/fleet/.storybook/context/fixtures/integration.okta.ts

@@ -106,6 +106,7 @@ export const item: GetInfoResponse['item'] = {
       ml_module: [],
       osquery_pack_asset: [],
       osquery_saved_query: [],
+      security_ai_prompt: [],
       security_rule: [],
       csp_rule_template: [],
       tag: [],

x-pack/platform/plugins/shared/fleet/common/services/package_to_package_policy.test.ts

@@ -33,6 +33,7 @@ describe('Fleet - packageToPackagePolicy', () => {
         map: [],
         lens: [],
         ml_module: [],
+        security_ai_prompt: [],
         security_rule: [],
         tag: [],
         osquery_pack_asset: [],

x-pack/platform/plugins/shared/fleet/common/types/models/epm.ts

@@ -59,6 +59,7 @@ export enum KibanaAssetType {
   indexPattern = 'index_pattern',
   map = 'map',
   mlModule = 'ml_module',
+  securityAIPrompt = 'security_ai_prompt',
   securityRule = 'security_rule',
   cloudSecurityPostureRuleTemplate = 'csp_rule_template',
   osqueryPackAsset = 'osquery_pack_asset',
@@ -77,6 +78,7 @@ export enum KibanaSavedObjectType {
   indexPattern = 'index-pattern',
   map = 'map',
   mlModule = 'ml-module',
+  securityAIPrompt = 'security-ai-prompt',
   securityRule = 'security-rule',
   cloudSecurityPostureRuleTemplate = 'csp-rule-template',
   osqueryPackAsset = 'osquery-pack-asset',

x-pack/platform/plugins/shared/fleet/public/applications/integrations/sections/epm/constants.tsx

@@ -47,6 +47,12 @@ export const AssetTitleMap: Record<
   map: i18n.translate('xpack.fleet.epm.assetTitles.maps', {
     defaultMessage: 'Maps',
   }),
+  'security-ai-prompt': i18n.translate('xpack.fleet.epm.assetTitles.securityAIPrompt', {
+    defaultMessage: 'Security AI prompt',
+  }),
+  security_ai_prompt: i18n.translate('xpack.fleet.epm.assetTitles.securityAIPrompt', {
+    defaultMessage: 'Security AI prompt',
+  }),
   'security-rule': i18n.translate('xpack.fleet.epm.assetTitles.securityRules', {
     defaultMessage: 'Security rules',
   }),

x-pack/platform/plugins/shared/fleet/server/routes/epm/index.test.ts

@@ -193,6 +193,7 @@ describe('schema validation', () => {
         map: [],
         index_pattern: [],
         ml_module: [],
+        security_ai_prompt: [],
         security_rule: [],
         tag: [],
         csp_rule_template: [],

x-pack/platform/plugins/shared/fleet/server/services/agent_policies/package_policies_to_agent_permissions.test.ts

@@ -38,6 +38,7 @@ packageInfoCache.set('test_package-0.0.0', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],
@@ -122,6 +123,7 @@ packageInfoCache.set('osquery_manager-0.3.0', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],
@@ -172,6 +174,7 @@ packageInfoCache.set('profiler_symbolizer-8.8.0-preview', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],
@@ -222,6 +225,7 @@ packageInfoCache.set('profiler_collector-8.9.0-preview', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],
@@ -264,6 +268,7 @@ packageInfoCache.set('apm-8.9.0-preview', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],
@@ -306,6 +311,7 @@ packageInfoCache.set('elastic_connectors-1.0.0', {
       index_pattern: [],
       map: [],
       lens: [],
+      security_ai_prompt: [],
       security_rule: [],
       ml_module: [],
       tag: [],

x-pack/platform/plugins/shared/fleet/server/services/epm/kibana/assets/install.ts

@@ -33,9 +33,10 @@ import { deleteKibanaSavedObjectsAssets } from '../../packages/remove';
 import { FleetError, KibanaSOReferenceError } from '../../../../errors';
 import { withPackageSpan } from '../../packages/utils';
 
+import { appContextService } from '../../..';
+
 import { tagKibanaAssets } from './tag_assets';
 import { getSpaceAwareSaveobjectsClients } from './saved_objects';
-import { appContextService } from '../../..';
 
 const MAX_ASSETS_TO_INSTALL_IN_PARALLEL = 1000;
 
@@ -70,6 +71,7 @@ export const KibanaSavedObjectTypeMapping: Record<KibanaAssetType, KibanaSavedOb
   [KibanaAssetType.visualization]: KibanaSavedObjectType.visualization,
   [KibanaAssetType.lens]: KibanaSavedObjectType.lens,
   [KibanaAssetType.mlModule]: KibanaSavedObjectType.mlModule,
+  [KibanaAssetType.securityAIPrompt]: KibanaSavedObjectType.securityAIPrompt,
   [KibanaAssetType.securityRule]: KibanaSavedObjectType.securityRule,
   [KibanaAssetType.cloudSecurityPostureRuleTemplate]:
     KibanaSavedObjectType.cloudSecurityPostureRuleTemplate,

x-pack/solutions/security/plugins/elastic_assistant/common/constants.ts

@@ -31,3 +31,6 @@ export const CAPABILITIES = `${BASE_PATH}/capabilities`;
  Licensing requirements
  */
 export const MINIMUM_AI_ASSISTANT_LICENSE = 'enterprise' as const;
+
+// Saved Objects
+export const promptSavedObjectType = 'security-ai-prompt';

x-pack/solutions/security/plugins/elastic_assistant/scripts/draw_graph_script.ts

@@ -18,6 +18,19 @@ import type { Logger } from '@kbn/logging';
 import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { FakeLLM } from '@langchain/core/utils/testing';
 import { createOpenAIFunctionsAgent } from 'langchain/agents';
+import { actionsClientMock } from '@kbn/actions-plugin/server/actions_client/actions_client.mock';
+import { savedObjectsClientMock } from '@kbn/core/server/mocks';
+import {
+  ATTACK_DISCOVERY_GENERATION_DETAILS_MARKDOWN,
+  ATTACK_DISCOVERY_GENERATION_ENTITY_SUMMARY_MARKDOWN,
+  ATTACK_DISCOVERY_GENERATION_INSIGHTS,
+  ATTACK_DISCOVERY_GENERATION_MITRE_ATTACK_TACTICS,
+  ATTACK_DISCOVERY_GENERATION_SUMMARY_MARKDOWN,
+  ATTACK_DISCOVERY_GENERATION_TITLE,
+  ATTACK_DISCOVERY_CONTINUE,
+  ATTACK_DISCOVERY_DEFAULT,
+  ATTACK_DISCOVERY_REFINE,
+} from '../server/lib/prompt/prompts';
 import { getDefaultAssistantGraph } from '../server/lib/langchain/graphs/default_assistant_graph/graph';
 import { getDefaultAttackDiscoveryGraph } from '../server/lib/attack_discovery/graphs/default_attack_discovery_graph';
 
@@ -49,11 +62,13 @@ async function getAssistantGraph(logger: Logger): Promise<Drawable> {
     streamRunnable: false,
   });
   const graph = getDefaultAssistantGraph({
+    actionsClient: actionsClientMock.create(),
     agentRunnable,
     logger,
     createLlmInstance,
     tools: [],
     replacements: {},
+    savedObjectsClient: savedObjectsClientMock.create(),
   });
   return graph.getGraph();
 }
@@ -67,6 +82,17 @@ async function getAttackDiscoveryGraph(logger: Logger): Promise<Drawable> {
     llm: mockLlm as unknown as ActionsClientLlm,
     logger,
     replacements: {},
+    prompts: {
+      default: ATTACK_DISCOVERY_DEFAULT,
+      refine: ATTACK_DISCOVERY_REFINE,
+      continue: ATTACK_DISCOVERY_CONTINUE,
+      detailsMarkdown: ATTACK_DISCOVERY_GENERATION_DETAILS_MARKDOWN,
+      entitySummaryMarkdown: ATTACK_DISCOVERY_GENERATION_ENTITY_SUMMARY_MARKDOWN,
+      mitreAttackTactics: ATTACK_DISCOVERY_GENERATION_MITRE_ATTACK_TACTICS,
+      summaryMarkdown: ATTACK_DISCOVERY_GENERATION_SUMMARY_MARKDOWN,
+      title: ATTACK_DISCOVERY_GENERATION_TITLE,
+      insights: ATTACK_DISCOVERY_GENERATION_INSIGHTS,
+    },
     size: 20,
   });
 

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

@@ -54,8 +54,8 @@ export const createMockClients = () => {
       getCurrentUser: jest.fn(),
       inference: jest.fn(),
       llmTasks: jest.fn(),
+      savedObjectsClient: core.savedObjects.client,
     },
-    savedObjectsClient: core.savedObjects.client,
 
     licensing: {
       ...licensingMock.createRequestHandlerContext({ license }),
@@ -148,6 +148,7 @@ const createElasticAssistantRequestContextMock = (
     inference: { getClient: jest.fn() },
     llmTasks: { retrieveDocumentationAvailable: jest.fn(), retrieveDocumentation: jest.fn() },
     core: clients.core,
+    savedObjectsClient: clients.elasticAssistant.savedObjectsClient,
     telemetry: clients.elasticAssistant.telemetry,
   };
 };

x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.test.ts

@@ -4,7 +4,6 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-
 import type { ActionsClient } from '@kbn/actions-plugin/server';
 import type { Connector } from '@kbn/actions-plugin/server/application/connector/types';
 import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
@@ -46,7 +45,6 @@ jest.mock('./helpers/get_evaluator_llm', () => {
     getEvaluatorLlm: jest.fn().mockResolvedValue(mockLlm),
   };
 });
-
 const actionsClient = {
   get: jest.fn(),
 } as unknown as ActionsClient;
@@ -61,7 +59,22 @@ const logger = loggerMock.create();
 const mockEsClient = elasticsearchServiceMock.createElasticsearchClient();
 const runName = 'test-run-name';
 
-const connectors = [mockExperimentConnector];
+const connectors = [
+  {
+    ...mockExperimentConnector,
+    prompts: {
+      default: '',
+      refine: '',
+      continue: '',
+      detailsMarkdown: '',
+      entitySummaryMarkdown: '',
+      mitreAttackTactics: '',
+      summaryMarkdown: '',
+      title: '',
+      insights: '',
+    },
+  },
+];
 
 const projectName = 'test-lang-smith-project';
 

x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/evaluation/index.ts

@@ -16,12 +16,16 @@ import { getLangSmithTracer } from '@kbn/langchain/server/tracers/langsmith';
 import { asyncForEach } from '@kbn/std';
 import { PublicMethodsOf } from '@kbn/utility-types';
 
+import { CombinedPrompts } from '../graphs/default_attack_discovery_graph/nodes/helpers/prompts';
 import { DEFAULT_EVAL_ANONYMIZATION_FIELDS } from './constants';
 import { AttackDiscoveryGraphMetadata } from '../../langchain/graphs';
 import { DefaultAttackDiscoveryGraph } from '../graphs/default_attack_discovery_graph';
 import { getLlmType } from '../../../routes/utils';
 import { runEvaluations } from './run_evaluations';
 
+interface ConnectorWithPrompts extends Connector {
+  prompts: CombinedPrompts;
+}
 export const evaluateAttackDiscovery = async ({
   actionsClient,
   attackDiscoveryGraphs,
@@ -43,7 +47,7 @@ export const evaluateAttackDiscovery = async ({
   attackDiscoveryGraphs: AttackDiscoveryGraphMetadata[];
   alertsIndexPattern: string;
   anonymizationFields?: AnonymizationFieldResponse[];
-  connectors: Connector[];
+  connectors: ConnectorWithPrompts[];
   connectorTimeout: number;
   datasetName: string;
   esClient: ElasticsearchClient;
@@ -96,6 +100,7 @@ export const evaluateAttackDiscovery = async ({
         esClient,
         llm,
         logger,
+        prompts: connector.prompts,
         size,
       });
 

x-pack/solutions/security/plugins/elastic_assistant/server/lib/attack_discovery/graphs/default_attack_discovery_graph/edges/generate_or_end/index.test.ts

@@ -35,6 +35,7 @@ const graphState: GraphState = {
   ],
   combinedGenerations: 'generations',
   combinedRefinements: 'refinements',
+  continuePrompt: 'continue',
   errors: [],
   generationAttempts: 0,
   generations: [],