[Security Solution][Security Assistant] Investigate in timeline control should be visible only on eligible queries. (#161240)

## Summary Handles elastic/security-team#6971 This PR mainly resolved below 3 issues: ### Rename to `Add To Timeline` control in conversation code blocks to `Investigate in Timeline` - `Add to Timeline` according to existing Security Solution actions means, adding a condition to the timeline with an `OR` clause without affecting the existing Timeline. - But the `Add to Timeline` control in the Security Assistant, creates a new timeline on each action by the user, which contradicts the above workflow. Hence, it might confuse user. - `Investigate in Timeline` already means that a new timeline will be created. ### `Add To Timeline` control was visible on types of codeblock. For example, it does not make sense for a `Query DSL` to have an `Add to Timeline` control. - This PR adds the list of eligible types of queries/code blocks on which `Add To Timeline` action can be added. - Currently, that list only contains `kql`, `dsl` and `eql`. Below is the complete list of types of query that can occur in code blocks. - Please feel free to suggest a change. ``` 'eql' | 'kql' | 'dsl' | 'json' | 'no-type'; ``` ### Lazy calculation of CodeBlockPortals and CodeBlock Action container - To add controls to the conversation code blocks, we need to follow below 2 steps. 1. get the codeBlock containers on which the controls can be added. 2. create portals in the HTML container with our `Add to Timeline` control. - Below are issues these steps sometime created. 1. We get codeBlock container in the `useLayoutEffect` but at the time, all conversations might not have loaded because of which containers are returns as the undefined. 2. Then, we try to create portal in the `undefined` container, which fails and hence, `Add to Timeline` controls are not visible. - Solution: 1. Instead of getting the codeblock container in useLayoutEffect, we get the function which will eventually return that container, whenever we are creating the portal. 2. Converted codeBlock Portal to a callback such that callback can be called during the rendering which makes sure that all needed conversations are available and using above step we can easily get the portal containers. Feel free to let me know if there are any issues with above strategy. ### Better Pattern matching. - Currently, when we are trying to identify the type of codeblock it might result in unexpected output because of below reason. 1. Let say, we are trying to identify KQL Query and for that we use below phrases to match in the `OpenAI` response. `'Kibana Query Language', 'KQL Query'` 2. Because of this, if the `OpenAI` response contains the phrase `KQL query` or `kql query`, that fails because of case senstivity when searching the above phrases. 3. This PR makes that part of pattern matching case insensitive ### Before b472178a-0145-42d8-8fb9-ab107915086a ### After b499f099-a7a1-435f-99b2-ab27ee1f5680 ### Checklist Delete any items that are not applicable to this PR. - [x] Any text added follows [EUI's writing guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses sentence case text and includes [i18n support](https://github.com/elastic/kibana/blob/main/packages/kbn-i18n/README.md) - [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 - [x] Any UI touched in this PR is usable by keyboard only (learn more about [keyboard accessibility](https://webaim.org/techniques/keyboard/))
2025-04-23 17:28:26 -04:00 · 2023-07-11 11:09:33 -07:00 · 2023-07-11 11:09:33 -07:00 · 091b5c133b
commit 091b5c133b
parent 323b0477e3
6 changed files with 138 additions and 50 deletions
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
@ -175,9 +175,7 @@ const AssistantComponent: React.FC<Props> = ({

  const [showAnonymizedValues, setShowAnonymizedValues] = useState<boolean>(false);

-  const [messageCodeBlocks, setMessageCodeBlocks] = useState<CodeBlockDetails[][]>(
-    augmentMessageCodeBlocks(currentConversation)
-  );
+  const [messageCodeBlocks, setMessageCodeBlocks] = useState<CodeBlockDetails[][]>();
  const [_, setCodeBlockControlsVisible] = useState(false);
  useLayoutEffect(() => {
    setMessageCodeBlocks(augmentMessageCodeBlocks(currentConversation));
@ -365,17 +363,13 @@ const AssistantComponent: React.FC<Props> = ({
    setShowMissingConnectorCallout(!connectorExists);
  }, [connectors, currentConversation]);

-  const CodeBlockPortals = useMemo(
+  const createCodeBlockPortals = useCallback(
    () =>
-      messageCodeBlocks.map((codeBlocks: CodeBlockDetails[]) => {
+      messageCodeBlocks?.map((codeBlocks: CodeBlockDetails[]) => {
        return codeBlocks.map((codeBlock: CodeBlockDetails) => {
-          const element: Element = codeBlock.controlContainer as Element;
-
-          return codeBlock.controlContainer != null ? (
-            createPortal(codeBlock.button, element)
-          ) : (
-            <></>
-          );
+          const getElement = codeBlock.getControlContainer;
+          const element = getElement?.();
+          return element ? createPortal(codeBlock.button, element) : <></>;
        });
      }),
    [messageCodeBlocks]
@ -531,7 +525,7 @@ const AssistantComponent: React.FC<Props> = ({
        )}

        {/* Create portals for each EuiCodeBlock to add the `Investigate in Timeline` action */}
-        {CodeBlockPortals}
+        {createCodeBlockPortals()}

        {!isDisabled && (
          <>
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.test.ts
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/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 { analyzeMarkdown } from './helpers';
+
+const tilde = '`';
+const codeDelimiter = '```';
+
+const markDownWithDSLQuery = `
+Certainly! Here's an example of a Query DSL (Domain-Specific Language) query using the Elasticsearch Query DSL syntax:
+
+${codeDelimiter}
+POST /<index>/_search
+{
+  \"query\": {
+    \"bool\": {
+      \"must\": [
+        {
+          \"match\": {
+            \"event.category\": \"security\"
+          }
+        },
+        {
+          \"match\": {
+            \"message\": \"keyword\"
+          }
+        }
+      ]
+    }
+  }
+}
+${codeDelimiter}
+
+In this example, you need to replace ${tilde}<index>${tilde} with the actual name of the index where your security-related data is stored.
+
+The query is structured using the JSON format. It uses the ${tilde}bool${tilde} query to combine multiple conditions using the ${tilde}must${tilde} clause. In this case, we are using the ${tilde}match${tilde} query to search for documents where the ${tilde}event.category${tilde} field matches \"security\" and the ${tilde}message${tilde} field matches \"keyword\". You can modify these values to match your specific search criteria.
+
+By sending this query to the appropriate endpoint, you can retrieve search results that match the specified conditions. The response will include the relevant documents that meet the search criteria.
+
+Remember to refer to the Elastic documentation for more information on the available DQL syntax and query options to further customize and refine your searches based on your specific needs.
+`;
+
+const markDownWithKQLQuery = `Certainly! Here's a KQL query based on the ${tilde}user.name${tilde} field:
+
+${codeDelimiter}
+user.name: \"9dcc9960-78cf-4ef6-9a2e-dbd5816daa60\"
+${codeDelimiter}
+
+This query will filter the events based on the condition that the ${tilde}user.name${tilde} field should exactly match the value \"9dcc9960-78cf-4ef6-9a2e-dbd5816daa60\".`;
+
+describe('analyzeMarkdown', () => {
+  it('should identify dsl Query successfully.', () => {
+    const result = analyzeMarkdown(markDownWithDSLQuery);
+    expect(result[0].type).toBe('dsl');
+  });
+  it('should identify kql Query successfully.', () => {
+    const result = analyzeMarkdown(markDownWithKQLQuery);
+    expect(result[0].type).toBe('kql');
+  });
+});
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/use_conversation/helpers.ts
@ -12,7 +12,7 @@ export interface CodeBlockDetails {
  content: string;
  start: number;
  end: number;
-  controlContainer?: React.ReactNode;
+  getControlContainer?: () => Element | undefined;
  button?: React.ReactNode;
 }

@ -32,9 +32,15 @@ export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => {
  const matches = [...markdown.matchAll(codeBlockRegex)];
  // If your codeblocks aren't getting tagged with the right language, add keywords to the array.
  const types = {
-    eql: ['Event Query Language', 'EQL sequence query'],
-    kql: ['Kibana Query Language', 'KQL Query'],
-    dsl: ['Elasticsearch QueryDSL', 'Elasticsearch Query DSL', 'Elasticsearch DSL'],
+    eql: ['Event Query Language', 'EQL sequence query', 'EQL'],
+    kql: ['Kibana Query Language', 'KQL Query', 'KQL'],
+    dsl: [
+      'Elasticsearch QueryDSL',
+      'Elasticsearch Query DSL',
+      'Elasticsearch DSL',
+      'Query DSL',
+      'DSL',
+    ],
  };

  const result: CodeBlockDetails[] = matches.map((match) => {
@ -43,7 +49,7 @@ export const analyzeMarkdown = (markdown: string): CodeBlockDetails[] => {
      const start = match.index || 0;
      const precedingText = markdown.slice(0, start);
      for (const [typeKey, keywords] of Object.entries(types)) {
-        if (keywords.some((kw) => precedingText.includes(kw))) {
+        if (keywords.some((kw) => precedingText.toLowerCase().includes(kw.toLowerCase()))) {
          type = typeKey;
          break;
        }
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx
@ -148,6 +148,12 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
    baseSystemPrompts
  );

+  // if basePrompt has been updated, the localstorage should be accordingly updated
+  // if it exists
+  useEffect(() => {
+    setLocalStorageSystemPrompts(baseSystemPrompts);
+  }, [baseSystemPrompts, setLocalStorageSystemPrompts]);
+
  /**
   * Prompt contexts are used to provide components a way to register and make their data available to the assistant.
   */
--- a/x-pack/plugins/security_solution/public/assistant/content/prompts/system/translations.ts
+++ b/x-pack/plugins/security_solution/public/assistant/content/prompts/system/translations.ts
@ -37,7 +37,16 @@ export const SUPERHERO_PERSONALITY = i18n.translate(
  }
 );

+export const FORMAT_OUTPUT_CORRECTLY = i18n.translate(
+  'xpack.securitySolution.assistant.content.prompts.system.outputFormatting',
+  {
+    defaultMessage:
+      'If you answer a question related to KQL or EQL, it should be immediately usable within an Elastic Security timeline, please always format the output correctly with back ticks. Any answer provided for Query DSL should also be usable in a security timeline. This means you should only ever include the "filter" portion of the query.',
+  }
+);
+
 export const DEFAULT_SYSTEM_PROMPT_NON_I18N = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}
+${FORMAT_OUTPUT_CORRECTLY}
 ${USE_THE_FOLLOWING_CONTEXT_TO_ANSWER}`;

 export const DEFAULT_SYSTEM_PROMPT_NAME = i18n.translate(
@ -49,6 +58,7 @@ export const DEFAULT_SYSTEM_PROMPT_NAME = i18n.translate(

 export const SUPERHERO_SYSTEM_PROMPT_NON_I18N = `${YOU_ARE_A_HELPFUL_EXPERT_ASSISTANT} ${IF_YOU_DONT_KNOW_THE_ANSWER}
 ${SUPERHERO_PERSONALITY}
+${FORMAT_OUTPUT_CORRECTLY}
 ${USE_THE_FOLLOWING_CONTEXT_TO_ANSWER}`;

 export const SUPERHERO_SYSTEM_PROMPT_NAME = i18n.translate(
--- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx
@ -13,6 +13,7 @@ import React from 'react';
 import type { TimelineEventsDetailsItem } from '../../common/search_strategy';
 import type { Rule } from '../detection_engine/rule_management/logic';
 import { SendToTimelineButton } from './send_to_timeline';
+import { INVESTIGATE_IN_TIMELINE } from '../actions/add_to_timeline/constants';

 export const LOCAL_STORAGE_KEY = `securityAssistant`;

@ -48,6 +49,8 @@ export const getPromptContextFromEventDetailsItem = (data: TimelineEventsDetails
  return getFieldsAsCsv(allFields);
 };

+const sendToTimelineEligibleQueryTypes: Array<CodeBlockDetails['type']> = ['kql', 'dsl', 'eql'];
+
 /**
 * Augments the messages in a conversation with code block details, including
 * the start and end indices of the code block in the message, the type of the
@ -60,38 +63,43 @@ export const augmentMessageCodeBlocks = (
 ): CodeBlockDetails[][] => {
  const cbd = currentConversation.messages.map(({ content }) => analyzeMarkdown(content));

-  return cbd.map((codeBlocks, messageIndex) =>
-    codeBlocks.map((codeBlock, codeBlockIndex) => ({
-      ...codeBlock,
-      controlContainer: document.querySelectorAll(
-        `.message-${messageIndex} .euiCodeBlock__controls`
-      )[codeBlockIndex],
-      button: (
-        <SendToTimelineButton
-          asEmptyButton={true}
-          dataProviders={[
-            {
-              id: 'assistant-data-provider',
-              name: `Assistant Query from conversation ${currentConversation.id}`,
-              enabled: true,
-              excluded: false,
-              queryType: codeBlock.type,
-              kqlQuery: codeBlock.content ?? '',
-              queryMatch: {
-                field: 'host.name',
-                operator: ':',
-                value: 'test',
+  const output = cbd.map((codeBlocks, messageIndex) =>
+    codeBlocks.map((codeBlock, codeBlockIndex) => {
+      return {
+        ...codeBlock,
+        getControlContainer: () =>
+          document.querySelectorAll(`.message-${messageIndex} .euiCodeBlock__controls`)[
+            codeBlockIndex
+          ],
+        button: sendToTimelineEligibleQueryTypes.includes(codeBlock.type) ? (
+          <SendToTimelineButton
+            asEmptyButton={true}
+            dataProviders={[
+              {
+                id: 'assistant-data-provider',
+                name: `Assistant Query from conversation ${currentConversation.id}`,
+                enabled: true,
+                excluded: false,
+                queryType: codeBlock.type,
+                kqlQuery: codeBlock.content ?? '',
+                queryMatch: {
+                  field: 'host.name',
+                  operator: ':',
+                  value: 'test',
+                },
+                and: [],
              },
-              and: [],
-            },
-          ]}
-          keepDataView={true}
-        >
-          <EuiToolTip position="right" content={'Add to timeline'}>
-            <EuiIcon type="timeline" />
-          </EuiToolTip>
-        </SendToTimelineButton>
-      ),
-    }))
+            ]}
+            keepDataView={true}
+          >
+            <EuiToolTip position="right" content={INVESTIGATE_IN_TIMELINE}>
+              <EuiIcon type="timeline" />
+            </EuiToolTip>
+          </SendToTimelineButton>
+        ) : null,
+      };
+    })
  );
+
+  return output;
 };