[Security Solution][Elastic AI Assistant] ESQL query generation tweaks (#169984)

## Summary

After testing in the latest BC these three tweaks were identified:

- [X] Set k value to 10 so more documents are returned (since ESQL KB
docs are broken up by function and can be small)
- [X] Adds Quick Prompt to aid in ESQL Query Generation that only shows
if KB is enabled

x-pack/packages/kbn-elastic-assistant/impl/assistant/quick_prompts/quick_prompts.test.tsx

@@ -25,6 +25,9 @@ const mockUseAssistantContext = {
   setSelectedSettingsTab,
   promptContexts: {},
   allQuickPrompts: MOCK_QUICK_PROMPTS,
+  knowledgeBase: {
+    assistantLangChain: true,
+  },
 };
 
 const testTitle = 'SPL_QUERY_CONVERSION_TITLE';

x-pack/packages/kbn-elastic-assistant/impl/assistant/quick_prompts/quick_prompts.tsx

@@ -19,6 +19,8 @@ const QuickPromptsFlexGroup = styled(EuiFlexGroup)`
   margin: 16px;
 `;
 
+export const KNOWLEDGE_BASE_CATEGORY = 'knowledge-base';
+
 const COUNT_BEFORE_OVERFLOW = 5;
 interface QuickPromptsProps {
   setInput: (input: string) => void;
@@ -33,10 +35,15 @@ interface QuickPromptsProps {
  */
 export const QuickPrompts: React.FC<QuickPromptsProps> = React.memo(
   ({ setInput, setIsSettingsModalVisible, trackPrompt }) => {
-    const { allQuickPrompts, promptContexts, setSelectedSettingsTab } = useAssistantContext();
+    const { allQuickPrompts, knowledgeBase, promptContexts, setSelectedSettingsTab } =
+      useAssistantContext();
 
     const contextFilteredQuickPrompts = useMemo(() => {
       const registeredPromptContextTitles = Object.values(promptContexts).map((pc) => pc.category);
+      // If KB is enabled, include KNOWLEDGE_BASE_CATEGORY so KB dependent quick prompts are shown
+      if (knowledgeBase.assistantLangChain) {
+        registeredPromptContextTitles.push(KNOWLEDGE_BASE_CATEGORY);
+      }
       return allQuickPrompts.filter((quickPrompt) => {
         // Return quick prompt as match if it has no categories, otherwise ensure category exists in registered prompt contexts
         if (quickPrompt.categories == null || quickPrompt.categories.length === 0) {
@@ -47,7 +54,7 @@ export const QuickPrompts: React.FC<QuickPromptsProps> = React.memo(
           });
         }
       });
-    }, [allQuickPrompts, promptContexts]);
+    }, [allQuickPrompts, knowledgeBase.assistantLangChain, promptContexts]);
 
     // Overflow state
     const [isOverflowPopoverOpen, setIsOverflowPopoverOpen] = useState(false);

x-pack/plugins/elastic_assistant/server/lib/langchain/elasticsearch_store/elasticsearch_store.ts

@@ -222,6 +222,14 @@ export class ElasticsearchStore extends VectorStore {
         return getFlattenedHits(maybeEsqlMsearchResponse);
       });
 
+      this.logger.debug(
+        `Similarity search metadata source:\n${JSON.stringify(
+          results.map((r) => r?.metadata?.source ?? '(missing metadata.source)'),
+          null,
+          2
+        )}`
+      );
+
       return results;
     } catch (e) {
       this.logger.error(e);

x-pack/plugins/elastic_assistant/server/lib/langchain/execute_custom_llm_chain/index.ts

@@ -55,7 +55,8 @@ export const callAgentExecutor = async ({
     );
   }
 
-  const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever());
+  // Create a chain that uses the ELSER backed ElasticsearchStore, override k=10 for esql query generation for now
+  const chain = RetrievalQAChain.fromLLM(llm, esStore.asRetriever(10));
 
   const tools: Tool[] = [
     new ChainTool({

x-pack/plugins/security_solution/public/assistant/content/prompt_contexts/index.tsx

@@ -16,6 +16,7 @@ export const PROMPT_CONTEXT_ALERT_CATEGORY = 'alert';
 export const PROMPT_CONTEXT_EVENT_CATEGORY = 'event';
 export const PROMPT_CONTEXT_DETECTION_RULES_CATEGORY = 'detection-rules';
 export const DATA_QUALITY_DASHBOARD_CATEGORY = 'data-quality-dashboard';
+export const KNOWLEDGE_BASE_CATEGORY = 'knowledge-base';
 
 /**
  * Global list of PromptContexts intended to be used throughout Security Solution.

x-pack/plugins/security_solution/public/assistant/content/quick_prompts/index.tsx

@@ -8,6 +8,7 @@
 import type { QuickPrompt } from '@kbn/elastic-assistant';
 import * as i18n from './translations';
 import {
+  KNOWLEDGE_BASE_CATEGORY,
   PROMPT_CONTEXT_ALERT_CATEGORY,
   PROMPT_CONTEXT_DETECTION_RULES_CATEGORY,
   PROMPT_CONTEXT_EVENT_CATEGORY,
@@ -26,6 +27,13 @@ export const BASE_SECURITY_QUICK_PROMPTS: QuickPrompt[] = [
     categories: [PROMPT_CONTEXT_ALERT_CATEGORY],
     isDefault: true,
   },
+  {
+    title: i18n.ESQL_QUERY_GENERATION_TITLE,
+    prompt: i18n.ESQL_QUERY_GENERATION_PROMPT,
+    color: '#9170B8',
+    categories: [KNOWLEDGE_BASE_CATEGORY],
+    isDefault: true,
+  },
   {
     title: i18n.RULE_CREATION_TITLE,
     prompt: i18n.RULE_CREATION_PROMPT,

x-pack/plugins/security_solution/public/assistant/content/quick_prompts/translations.ts

@@ -22,6 +22,21 @@ export const ALERT_SUMMARIZATION_PROMPT = i18n.translate(
   }
 );
 
+export const ESQL_QUERY_GENERATION_TITLE = i18n.translate(
+  'xpack.securitySolution.assistant.quickPrompts.esqlQueryGenerationTitle',
+  {
+    defaultMessage: 'ES|QL Query Generation',
+  }
+);
+
+export const ESQL_QUERY_GENERATION_PROMPT = i18n.translate(
+  'xpack.securitySolution.assistant.quickPrompts.esqlQueryGenerationPrompt',
+  {
+    defaultMessage:
+      "As an expert user of Elastic Security, please generate an accurate and valid ESQL query to detect the use case below. Your response should be formatted to be able to use immediately in an Elastic Security timeline or detection rule. Take your time with the answer, check your knowledge really well on all the functions I am asking for. For ES|QL answers specifically, you should only ever answer with what's available in your private knowledge. I cannot afford for queries to be inaccurate. Assume I am using the Elastic Common Schema and Elastic Agent.\n\nEnsure the answers are formatted in a way which is easily copyable as a separate code block in markdown.",
+  }
+);
+
 export const RULE_CREATION_TITLE = i18n.translate(
   'xpack.securitySolution.assistant.quickPrompts.ruleCreationTitle',
   {