[SIEM Rule Migration] Minor prompt improvements

x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/graph.ts

@@ -20,7 +20,7 @@ export function getRuleMigrationAgent({
   connectorId,
   logger,
 }: MigrateRuleGraphParams) {
-  const matchPrebuiltRuleNode = getMatchPrebuiltRuleNode({ model, prebuiltRulesMap, logger });
+  const matchPrebuiltRuleNode = getMatchPrebuiltRuleNode({ model, prebuiltRulesMap });
   const translationSubGraph = getTranslateRuleGraph({
     model,
     inferenceClient,

x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/match_prebuilt_rule.ts

@@ -5,17 +5,19 @@
  * 2.0.
  */
 
-import type { Logger } from '@kbn/core/server';
-import { StringOutputParser } from '@langchain/core/output_parsers';
+import { JsonOutputParser } from '@langchain/core/output_parsers';
 import type { ChatModel } from '../../../util/actions_client_chat';
-import type { GraphNode } from '../../types';
 import { filterPrebuiltRules, type PrebuiltRulesMapByName } from '../../../util/prebuilt_rules';
+import type { GraphNode } from '../../types';
 import { MATCH_PREBUILT_RULE_PROMPT } from './prompts';
 
 interface GetMatchPrebuiltRuleNodeParams {
   model: ChatModel;
   prebuiltRulesMap: PrebuiltRulesMapByName;
-  logger: Logger;
+}
+
+interface GetMatchedRuleResponse {
+  match: string;
 }
 
 export const getMatchPrebuiltRuleNode =
@@ -25,34 +27,32 @@ export const getMatchPrebuiltRuleNode =
     if (!mitreAttackIds?.length) {
       return {};
     }
+
     const filteredPrebuiltRulesMap = filterPrebuiltRules(prebuiltRulesMap, mitreAttackIds);
     if (filteredPrebuiltRulesMap.size === 0) {
       return {};
     }
 
-    const outputParser = new StringOutputParser();
+    const outputParser = new JsonOutputParser();
     const matchPrebuiltRule = MATCH_PREBUILT_RULE_PROMPT.pipe(model).pipe(outputParser);
 
-    const elasticSecurityRules = Array(filteredPrebuiltRulesMap.keys()).join('\n');
-    const response = await matchPrebuiltRule.invoke({
+    const elasticSecurityRules = [...filteredPrebuiltRulesMap.keys()].join('\n');
+    const response = (await matchPrebuiltRule.invoke({
       elasticSecurityRules,
       ruleTitle: state.original_rule.title,
-    });
-    const cleanResponse = response.trim();
-    if (cleanResponse === 'no_match') {
-      return {};
-    }
-
-    const result = filteredPrebuiltRulesMap.get(cleanResponse);
-    if (result != null) {
-      return {
-        elastic_rule: {
-          title: result.rule.name,
-          description: result.rule.description,
-          prebuilt_rule_id: result.rule.rule_id,
-          id: result.installedRuleId,
-        },
-      };
+    })) as GetMatchedRuleResponse;
+    if (response.match) {
+      const result = filteredPrebuiltRulesMap.get(response.match);
+      if (result != null) {
+        return {
+          elastic_rule: {
+            title: result.rule.name,
+            description: result.rule.description,
+            prebuilt_rule_id: result.rule.rule_id,
+            id: result.installedRuleId,
+          },
+        };
+      }
     }
 
     return {};

x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/nodes/match_prebuilt_rule/prompts.ts

@@ -11,25 +11,39 @@ export const MATCH_PREBUILT_RULE_PROMPT = ChatPromptTemplate.fromMessages([
     'system',
     `You are an expert assistant in Cybersecurity, your task is to help migrating a SIEM detection rule, from Splunk Security to Elastic Security.
 You will be provided with a Splunk Detection Rule name by the user, your goal is to try find an Elastic Detection Rule that covers the same threat, if any.
-The list of Elastic Detection Rules suggested is provided in the context below.
+Here are some context for you to reference for your task, read it carefully as you will get questions about it later:
 
-Guidelines:
-If there is no Elastic rule in the list that covers the same threat, answer only with the string: no_match
-If there is one Elastic rule in the list that covers the same threat, answer only with its name without any further explanation.
-If there are multiple rules in the list that cover the same threat, answer with the most specific of them, for example: "Linux User Account Creation" is more specific than "User Account Creation".
-
-<ELASTIC_DETECTION_RULE_NAMES>
+<context>
+<elastic_detection_rule_names>
 {elasticSecurityRules}
-</ELASTIC_DETECTION_RULE_NAMES>
+</elastic_detection_rule_names>
+</context>
 `,
   ],
   [
     'human',
-    `The Splunk Detection Rule is:
-<<SPLUNK_RULE_TITLE>>
+    `See the below description of the relevant splunk rule and try to match it with any of the elastic detection rules with similar names.     
+<splunk_rule_name>
 {ruleTitle}
-<<SPLUNK_RULE_TITLE>>
+</splunk_rule_name>
+
+<guidelines>
+- Always reply with a JSON object with the key "match" and the value being the most relevant matched elastic detection rule name. Do not reply with anything else.
+- Only reply with exact matches, if you are unsure or do not find a very confident match, always reply with an empty string value in the match key, do not guess or reply with anything else.
+- If there is one Elastic rule in the list that covers the same threat, set the name of the matching rule as a value of the match key. Do not reply with anything else.
+- If there are multiple rules in the list that cover the same threat, answer with the most specific of them, for example: "Linux User Account Creation" is more specific than "User Account Creation".
+</guidelines>
+
+<example_response>
+U: <splunk_rule_name>
+Linux Auditd Add User Account Type
+</splunk_rule_name>
+A: Please find the match JSON object below:
+\`\`\`json
+{{"match": "Linux User Account Creation"}}
+\`\`\`
+</example_response>
 `,
   ],
-  ['ai', 'Please find the answer below:'],
+  ['ai', 'Please find the match JSON object below:'],
 ]);

x-pack/plugins/security_solution/server/lib/siem_migrations/rules/task/agent/sub_graphs/translate_rule/nodes/process_query/process_query.ts

@@ -10,7 +10,7 @@ import { isEmpty } from 'lodash/fp';
 import type { ChatModel } from '../../../../../util/actions_client_chat';
 import type { RuleResourceRetriever } from '../../../../../util/rule_resource_retriever';
 import type { GraphNode } from '../../types';
-import { getReplaceQueryResourcesPrompt } from './prompts';
+import { REPLACE_QUERY_RESOURCE_PROMPT, getResourcesContext } from './prompts';
 
 interface GetProcessQueryNodeParams {
   model: ChatModel;
@@ -25,9 +25,15 @@ export const getProcessQueryNode = ({
     let query = state.original_rule.query;
     const resources = await resourceRetriever.getResources(state.original_rule);
     if (!isEmpty(resources)) {
-      const replaceQueryResourcesPrompt = getReplaceQueryResourcesPrompt(state, resources);
-      const stringParser = new StringOutputParser();
-      query = await model.pipe(stringParser).invoke(replaceQueryResourcesPrompt);
+      const replaceQueryParser = new StringOutputParser();
+      const replaceQueryResourcePrompt =
+        REPLACE_QUERY_RESOURCE_PROMPT.pipe(model).pipe(replaceQueryParser);
+      const resourceContext = getResourcesContext(resources);
+      query = await replaceQueryResourcePrompt.invoke({
+        query: state.original_rule.query,
+        macros: resourceContext.macros,
+        lookup_tables: resourceContext.lists,
+      });
     }
     return { inline_query: query };
   };