[Security Solution][Detection Engine] update query automatically in rule create form through AI assistant (#190963)

## Summary - addresses https://github.com/elastic/kibana/issues/187270 ### UX Introduced button in code block <img width="1218" alt="Screenshot 2024-08-21 at 16 35 51" src="https://github.com/user-attachments/assets/69c82d7c-7305-41a6-9a29-5f27755727a6"> ### DEMO https://github.com/user-attachments/assets/32419edc-4bfa-4f4e-892b-2a6abb3c0f27 ### 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 - [x] Any UI touched in this PR is usable by keyboard only (learn more about [keyboard accessibility](https://webaim.org/techniques/keyboard/)) - [x] Any UI touched in this PR does not create any new axe failures (run axe in browser: [FF](https://addons.mozilla.org/en-US/firefox/addon/axe-devtools/), [Chrome](https://chrome.google.com/webstore/detail/axe-web-accessibility-tes/lhdoppojpmngadmnindnejefpokejbdd?hl=en-US)) - [x] This was checked for [cross-browser compatibility](https://www.elastic.co/support/matrix#matrix_browsers) --------- Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
2025-04-24 09:48:58 -04:00 · 2024-09-03 18:37:36 +01:00 · 2024-09-03 18:37:36 +01:00 · 3c9198abeb
commit 3c9198abeb
parent bcb030e558
15 changed files with 266 additions and 41 deletions
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant/index.tsx
@ -263,7 +263,7 @@ const AssistantComponent: React.FC<Props> = ({
  // Add min-height to all codeblocks so timeline icon doesn't overflow
  const codeBlockContainers = [...document.getElementsByClassName('euiCodeBlock')];
  // @ts-ignore-expect-error
-  codeBlockContainers.forEach((e) => (e.style.minHeight = '75px'));
+  codeBlockContainers.forEach((e) => (e.style.minHeight = '85px'));
  ////

  const onToggleShowAnonymizedValues = useCallback(() => {
--- a/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/assistant_context/index.tsx
@ -8,7 +8,7 @@
 import { EuiCommentProps } from '@elastic/eui';
 import type { HttpSetup } from '@kbn/core-http-browser';
 import { omit } from 'lodash/fp';
-import React, { useCallback, useMemo, useState } from 'react';
+import React, { useCallback, useMemo, useState, useRef } from 'react';
 import type { IToasts } from '@kbn/core-notifications-browser';
 import { ActionTypeRegistryContract } from '@kbn/triggers-actions-ui-plugin/public';
 import { useLocalStorage, useSessionStorage } from 'react-use';
@ -137,6 +137,7 @@ export interface UseAssistantContext {
  basePromptContexts: PromptContextTemplate[];
  unRegisterPromptContext: UnRegisterPromptContext;
  currentAppId: string;
+  codeBlockRef: React.MutableRefObject<(codeBlock: string) => void>;
 }

 const AssistantContext = React.createContext<UseAssistantContext | undefined>(undefined);
@ -237,6 +238,11 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
   */
  const [selectedSettingsTab, setSelectedSettingsTab] = useState<SettingsTabs | null>(null);

+  /**
+   * Setting code block ref that can be used to store callback from parent components
+   */
+  const codeBlockRef = useRef(() => {});
+
  const getLastConversationId = useCallback(
    // if a conversationId has been provided, use that
    // if not, check local storage
@ -284,6 +290,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
      setLastConversationId: setLocalStorageLastConversationId,
      baseConversations,
      currentAppId,
+      codeBlockRef,
    }),
    [
      actionTypeRegistry,
@ -316,6 +323,7 @@ export const AssistantProvider: React.FC<AssistantProviderProps> = ({
      setLocalStorageLastConversationId,
      baseConversations,
      currentAppId,
+      codeBlockRef,
    ]
  );

--- a/x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.test.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.test.tsx
@ -18,6 +18,11 @@ jest.mock('../assistant/use_assistant_overlay', () => ({
  useAssistantOverlay: () => mockUseAssistantOverlay,
 }));

+let mockUseAssistantContext = { codeBlockRef: { current: null } };
+jest.mock('../..', () => ({
+  useAssistantContext: () => mockUseAssistantContext,
+}));
+
 const defaultProps: Props = {
  category: 'alert',
  description: 'Test description',
@ -27,6 +32,9 @@ const defaultProps: Props = {
 };

 describe('NewChat', () => {
+  beforeEach(() => {
+    mockUseAssistantContext = { codeBlockRef: { current: null } };
+  });
  afterEach(() => {
    jest.clearAllMocks();
  });
@ -118,4 +126,17 @@ describe('NewChat', () => {

    expect(onShowOverlaySpy).toHaveBeenCalled();
  });
+
+  it('assigns onExportCodeBlock callback to context codeBlock reference', () => {
+    const onExportCodeBlock = jest.fn();
+    render(<NewChat {...defaultProps} onExportCodeBlock={onExportCodeBlock} />);
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toBe(onExportCodeBlock);
+  });
+
+  it('does not change assigns context codeBlock reference if onExportCodeBlock not defined', () => {
+    render(<NewChat {...defaultProps} />);
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toBe(null);
+  });
 });
--- a/x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.tsx
+++ b/x-pack/packages/kbn-elastic-assistant/impl/new_chat/index.tsx
@ -6,7 +6,8 @@
 */

 import { EuiButtonEmpty, EuiLink } from '@elastic/eui';
-import React, { useCallback, useMemo } from 'react';
+import React, { useCallback, useMemo, useEffect } from 'react';
+import { useAssistantContext } from '../..';

 import { PromptContext } from '../assistant/prompt_context/types';
 import { useAssistantOverlay } from '../assistant/use_assistant_overlay';
@ -29,6 +30,8 @@ export type Props = Omit<PromptContext, 'id'> & {
  asLink?: boolean;
  /** Optional callback when overlay shows */
  onShowOverlay?: () => void;
+  /** Optional callback that returns copied code block */
+  onExportCodeBlock?: (codeBlock: string) => void;
 };

 const NewChatComponent: React.FC<Props> = ({
@ -45,6 +48,7 @@ const NewChatComponent: React.FC<Props> = ({
  isAssistantEnabled,
  asLink = false,
  onShowOverlay,
+  onExportCodeBlock,
 }) => {
  const { showAssistantOverlay } = useAssistantOverlay(
    category,
@ -56,12 +60,25 @@ const NewChatComponent: React.FC<Props> = ({
    tooltip,
    isAssistantEnabled
  );
+  const { codeBlockRef } = useAssistantContext();

  const showOverlay = useCallback(() => {
    showAssistantOverlay(true);
    onShowOverlay?.();
  }, [showAssistantOverlay, onShowOverlay]);

+  useEffect(() => {
+    if (onExportCodeBlock) {
+      codeBlockRef.current = onExportCodeBlock;
+    }
+
+    return () => {
+      if (onExportCodeBlock) {
+        codeBlockRef.current = () => {};
+      }
+    };
+  }, [codeBlockRef, onExportCodeBlock]);
+
  const icon = useMemo(() => {
    if (iconType === null) {
      return undefined;
--- a/x-pack/plugins/security_solution/public/assistant/content/conversations/index.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/content/conversations/index.tsx
@ -10,6 +10,7 @@ import type { Conversation } from '@kbn/elastic-assistant';
 import { DATA_QUALITY_DASHBOARD_CONVERSATION_ID } from '@kbn/ecs-data-quality-dashboard';

 import { DETECTION_RULES_CONVERSATION_ID } from '../../../detections/pages/detection_engine/rules/translations';
+import { DETECTION_RULES_CREATE_FORM_CONVERSATION_ID } from '../../../detections/pages/detection_engine/translations';
 import {
  ALERT_SUMMARY_CONVERSATION_ID,
  EVENT_SUMMARY_CONVERSATION_ID,
@ -41,6 +42,14 @@ export const BASE_SECURITY_CONVERSATIONS: Record<string, Conversation> = {
    messages: [],
    replacements: {},
  },
+  [DETECTION_RULES_CREATE_FORM_CONVERSATION_ID]: {
+    id: '',
+    title: DETECTION_RULES_CREATE_FORM_CONVERSATION_ID,
+    category: 'assistant',
+    isDefault: true,
+    messages: [],
+    replacements: {},
+  },
  [EVENT_SUMMARY_CONVERSATION_ID]: {
    id: '',
    title: EVENT_SUMMARY_CONVERSATION_ID,
--- a/x-pack/plugins/security_solution/public/assistant/helpers.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/helpers.tsx
@ -13,9 +13,9 @@ import { replaceAnonymizedValuesWithOriginalValues } from '@kbn/elastic-assistan
 import type { TimelineEventsDetailsItem } from '../../common/search_strategy';
 import type { Rule } from '../detection_engine/rule_management/logic';
 import { SendToTimelineButton } from './send_to_timeline';
-
+import { DETECTION_RULES_CREATE_FORM_CONVERSATION_ID } from '../detections/pages/detection_engine/translations';
 export const LOCAL_STORAGE_KEY = `securityAssistant`;
-
+import { UpdateQueryInFormButton } from './update_query_in_form';
 export interface QueryField {
  field: string;
  values: string;
@ -84,30 +84,37 @@ export const augmentMessageCodeBlocks = (
          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: [],
-              },
-            ]}
-            keepDataView={true}
-          >
-            <EuiIcon type="timeline" />
-          </SendToTimelineButton>
-        ) : null,
+        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: [],
+                  },
+                ]}
+                keepDataView={true}
+              >
+                <EuiIcon type="timeline" />
+              </SendToTimelineButton>
+            ) : null}
+            {DETECTION_RULES_CREATE_FORM_CONVERSATION_ID === currentConversation.title ? (
+              <UpdateQueryInFormButton query={codeBlock.content ?? ''} />
+            ) : null}
+          </>
+        ),
      };
    })
  );
--- a/x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.test.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.test.tsx
@ -0,0 +1,32 @@
+/*
+ * 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 React from 'react';
+import { render, screen } from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+
+import { UpdateQueryInFormButton } from '.';
+
+const mockUseAssistantContext = { codeBlockRef: { current: jest.fn() } };
+jest.mock('@kbn/elastic-assistant', () => ({
+  useAssistantContext: () => mockUseAssistantContext,
+}));
+
+describe('UpdateQueryInFormButton', () => {
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('calls codeBlockRef callback on click', () => {
+    const testQuery = 'from auditbeat* | limit 10';
+    render(<UpdateQueryInFormButton query={testQuery} />);
+
+    userEvent.click(screen.getByTestId('update-query-in-form-button'));
+
+    expect(mockUseAssistantContext.codeBlockRef.current).toHaveBeenCalledWith(testQuery);
+  });
+});
--- a/x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.tsx
+++ b/x-pack/plugins/security_solution/public/assistant/update_query_in_form/index.tsx
@ -0,0 +1,44 @@
+/*
+ * 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 { FC, PropsWithChildren } from 'react';
+import React from 'react';
+import { EuiButtonEmpty, EuiToolTip, EuiIcon } from '@elastic/eui';
+import { useAssistantContext } from '@kbn/elastic-assistant';
+
+import { UPDATE_QUERY_IN_FORM_TOOLTIP } from './translations';
+
+export interface UpdateQueryInFormButtonProps {
+  query: string;
+}
+
+export const UpdateQueryInFormButton: FC<PropsWithChildren<UpdateQueryInFormButtonProps>> = ({
+  query,
+}) => {
+  const { codeBlockRef } = useAssistantContext();
+
+  const handleClick = () => {
+    codeBlockRef?.current?.(query);
+  };
+
+  return (
+    <EuiButtonEmpty
+      data-test-subj="update-query-in-form-button"
+      aria-label={UPDATE_QUERY_IN_FORM_TOOLTIP}
+      onClick={handleClick}
+      color="text"
+      flush="both"
+      size="xs"
+    >
+      <EuiToolTip position="right" content={UPDATE_QUERY_IN_FORM_TOOLTIP}>
+        <EuiIcon type="documentEdit" />
+      </EuiToolTip>
+    </EuiButtonEmpty>
+  );
+};
+
+UpdateQueryInFormButton.displayName = 'UpdateQueryInFormButton';
--- a/x-pack/plugins/security_solution/public/assistant/update_query_in_form/translations.ts
+++ b/x-pack/plugins/security_solution/public/assistant/update_query_in_form/translations.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 { i18n } from '@kbn/i18n';
+
+export const UPDATE_QUERY_IN_FORM_TOOLTIP = i18n.translate(
+  'xpack.securitySolution.assistant.updateQueryInFormTooltip',
+  {
+    defaultMessage: 'Update query in form',
+  }
+);
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/ai_assistant/index.test.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/ai_assistant/index.test.tsx
@ -26,14 +26,14 @@ describe('AiAssistant', () => {
  it('does not render chat component when does not have hasAssistantPrivilege', () => {
    useAssistantAvailabilityMock.mockReturnValue({ hasAssistantPrivilege: false });

-    const { container } = render(<AiAssistant getFields={jest.fn()} />, {
+    const { container } = render(<AiAssistant getFields={jest.fn()} setFieldValue={jest.fn()} />, {
      wrapper: TestProviders,
    });

    expect(container).toBeEmptyDOMElement();
  });
  it('renders chat component when has hasAssistantPrivilege', () => {
-    render(<AiAssistant getFields={jest.fn()} />, {
+    render(<AiAssistant getFields={jest.fn()} setFieldValue={jest.fn()} />, {
      wrapper: TestProviders,
    });

--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/ai_assistant/index.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/ai_assistant/index.tsx
@ -13,7 +13,7 @@ import { NewChat, AssistantAvatar } from '@kbn/elastic-assistant';

 import { METRIC_TYPE, TELEMETRY_EVENT, track } from '../../../../common/lib/telemetry';
 import { useAssistantAvailability } from '../../../../assistant/use_assistant_availability';
-import * as i18nAssistant from '../../../../detections/pages/detection_engine/rules/translations';
+import * as i18nAssistant from '../../../../detections/pages/detection_engine/translations';
 import type { DefineStepRule } from '../../../../detections/pages/detection_engine/rules/types';
 import type { FormHook, ValidationError } from '../../../../shared_imports';

@ -38,10 +38,15 @@ const retrieveErrorMessages = (errors: ValidationError[]): string =>

 interface AiAssistantProps {
  getFields: FormHook<DefineStepRule>['getFields'];
+  setFieldValue: FormHook<DefineStepRule>['setFieldValue'];
  language?: string | undefined;
 }

-const AiAssistantComponent: React.FC<AiAssistantProps> = ({ getFields, language }) => {
+const AiAssistantComponent: React.FC<AiAssistantProps> = ({
+  getFields,
+  setFieldValue,
+  language,
+}) => {
  const { hasAssistantPrivilege, isAssistantEnabled } = useAssistantAvailability();

  const languageName = getLanguageName(language);
@ -68,6 +73,23 @@ Proposed solution should be valid and must not contain new line symbols (\\n)`;
    track(METRIC_TYPE.COUNT, TELEMETRY_EVENT.OPEN_ASSISTANT_ON_RULE_QUERY_ERROR);
  }, []);

+  const handleOnExportCodeBlock = useCallback(
+    (codeBlock) => {
+      const queryField = getFields().queryBar;
+      const queryBar = queryField.value as DefineStepRule['queryBar'];
+
+      // sometimes AI assistant include redundant backtick symbols in code block
+      const newQuery = codeBlock.replaceAll('`', '');
+      if (queryBar.query.query !== newQuery) {
+        setFieldValue('queryBar', {
+          ...queryBar,
+          query: { ...queryBar.query, query: newQuery },
+        });
+      }
+    },
+    [getFields, setFieldValue]
+  );
+
  if (!hasAssistantPrivilege) {
    return null;
  }
@ -84,7 +106,7 @@ Proposed solution should be valid and must not contain new line symbols (\\n)`;
            <NewChat
              asLink={true}
              category="detection-rules"
-              conversationId={i18nAssistant.DETECTION_RULES_CONVERSATION_ID}
+              conversationId={i18nAssistant.DETECTION_RULES_CREATE_FORM_CONVERSATION_ID}
              description={i18n.ASK_ASSISTANT_DESCRIPTION}
              getPromptContext={getPromptContext}
              suggestedUserPrompt={i18n.ASK_ASSISTANT_USER_PROMPT(languageName)}
@ -92,6 +114,7 @@ Proposed solution should be valid and must not contain new line symbols (\\n)`;
              iconType={null}
              onShowOverlay={onShowOverlay}
              isAssistantEnabled={isAssistantEnabled}
+              onExportCodeBlock={handleOnExportCodeBlock}
            >
              <AssistantAvatar size="xxs" /> {i18n.ASK_ASSISTANT_ERROR_BUTTON}
            </NewChat>
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.test.tsx
@ -625,14 +625,48 @@ describe('StepDefineRule', () => {
  });

  describe('AI assistant', () => {
-    it('renders assistant when query is not valid', () => {
-      render(<TestForm formProps={{ isQueryBarValid: false, ruleType: 'query' }} />, {
-        wrapper: TestProviders,
-      });
+    it('renders assistant when query is not valid and not empty', () => {
+      const initialState = {
+        queryBar: {
+          query: { query: '*:*', language: 'kuery' },
+          filters: [],
+          saved_id: null,
+        },
+      };
+      render(
+        <TestForm
+          formProps={{ isQueryBarValid: false, ruleType: 'query' }}
+          initialState={initialState}
+        />,
+        {
+          wrapper: TestProviders,
+        }
+      );

      expect(screen.getByTestId('ai-assistant')).toBeInTheDocument();
    });

+    it('does not render assistant when query is not valid and empty', () => {
+      const initialState = {
+        queryBar: {
+          query: { query: '', language: 'kuery' },
+          filters: [],
+          saved_id: null,
+        },
+      };
+      render(
+        <TestForm
+          formProps={{ isQueryBarValid: false, ruleType: 'query' }}
+          initialState={initialState}
+        />,
+        {
+          wrapper: TestProviders,
+        }
+      );
+
+      expect(screen.queryByTestId('ai-assistant')).toBe(null);
+    });
+
    it('does not render assistant when query is valid', () => {
      render(<TestForm formProps={{ isQueryBarValid: true, ruleType: 'query' }} />, {
        wrapper: TestProviders,
--- a/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx
+++ b/x-pack/plugins/security_solution/public/detection_engine/rule_creation_ui/components/step_define_rule/index.tsx
@ -953,8 +953,12 @@ const StepDefineRuleComponent: FC<StepDefineRuleProps> = ({
            </>
          </RuleTypeEuiFormRow>

-          {!isMlRule(ruleType) && !isQueryBarValid && (
-            <AiAssistant getFields={form.getFields} language={queryBar?.query?.language} />
+          {!isMlRule(ruleType) && !isQueryBarValid && queryBar?.query?.query && (
+            <AiAssistant
+              getFields={form.getFields}
+              setFieldValue={form.setFieldValue}
+              language={queryBar?.query?.language}
+            />
          )}

          {isQueryRule(ruleType) && (
--- a/x-pack/plugins/security_solution/public/detections/pages/detection_engine/translations.ts
+++ b/x-pack/plugins/security_solution/public/detections/pages/detection_engine/translations.ts
@ -95,3 +95,10 @@ export const ML_RULES_UNAVAILABLE = (totalRules: number) =>
    defaultMessage:
      '{totalRules} {totalRules, plural, =1 {rule requires} other {rules require}} Machine Learning to enable.',
  });
+
+export const DETECTION_RULES_CREATE_FORM_CONVERSATION_ID = i18n.translate(
+  'xpack.securitySolution.detectionEngine.ruleManagement.detectionRulesCreateEditFormConversationId',
+  {
+    defaultMessage: 'Detection Rules Create form',
+  }
+);
--- a/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/prompts.cy.ts
+++ b/x-pack/test/security_solution_cypress/cypress/e2e/ai_assistant/prompts.cy.ts
@ -88,7 +88,11 @@ describe('AI Assistant Prompts', { tags: ['@ess', '@serverless'] }, () => {
    it('Add prompt from system prompt selector and set multiple conversations (including current) as default conversation', () => {
      visitGetStartedPage();
      openAssistant();
-      createSystemPrompt(testPrompt.title, testPrompt.prompt, ['Welcome', 'Alert summary']);
+      createSystemPrompt(testPrompt.title, testPrompt.prompt, [
+        'Welcome',
+        'Alert summary',
+        'Data Quality Dashboard',
+      ]);
      assertSystemPrompt(testPrompt.title);
      typeAndSendMessage('hello');
      assertMessageSent('hello', true, testPrompt.prompt);