[Security AI] Add tool and chat title prompts to security-ai-prompts (#209163)

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts

@@ -12,6 +12,7 @@ import { invokeGraph, streamGraph } from './helpers';
 import { loggerMock } from '@kbn/logging-mocks';
 import { AgentExecutorParams, AssistantDataClients } from '../../executors/types';
 import { elasticsearchClientMock } from '@kbn/core-elasticsearch-client-server-mocks';
+import { getPrompt, resolveProviderAndModel } from '@kbn/security-ai-prompts';
 import { getFindAnonymizationFieldsResultWithSingleHit } from '../../../../__mocks__/response';
 import {
   createOpenAIToolsAgent,
@@ -20,7 +21,10 @@ import {
 } from 'langchain/agents';
 import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock';
 import { savedObjectsClientMock } from '@kbn/core-saved-objects-api-server-mocks';
-import { resolveProviderAndModel } from '@kbn/security-ai-prompts';
+import { AssistantTool, AssistantToolParams } from '../../../..';
+import { promptGroupId as toolsGroupId } from '../../../prompt/tool_prompts';
+import { promptDictionary } from '../../../prompt';
+import { promptGroupId } from '../../../prompt/local_prompt_object';
 jest.mock('./graph');
 jest.mock('./helpers');
 jest.mock('langchain/agents');
@@ -29,6 +33,7 @@ jest.mock('@kbn/langchain/server/tracers/telemetry');
 jest.mock('@kbn/security-ai-prompts');
 const getDefaultAssistantGraphMock = getDefaultAssistantGraph as jest.Mock;
 const resolveProviderAndModelMock = resolveProviderAndModel as jest.Mock;
+const getPromptMock = getPrompt as jest.Mock;
 describe('callAssistantGraph', () => {
   const mockDataClients = {
     anonymizationFieldsDataClient: {
@@ -98,6 +103,7 @@ describe('callAssistantGraph', () => {
     (mockDataClients?.anonymizationFieldsDataClient?.findDocuments as jest.Mock).mockResolvedValue(
       getFindAnonymizationFieldsResultWithSingleHit()
     );
+    getPromptMock.mockResolvedValue('prompt');
   });
 
   it('calls invokeGraph with correct parameters for non-streaming', async () => {
@@ -173,6 +179,58 @@ describe('callAssistantGraph', () => {
     });
   });
 
+  it('calls getPrompt for each tool and the default system prompt', async () => {
+    const getTool = jest.fn();
+    const mockTool: AssistantTool = {
+      id: 'id',
+      name: 'name',
+      description: 'description',
+      sourceRegister: 'sourceRegister',
+      isSupported: (params: AssistantToolParams) => true,
+      getTool,
+    };
+    const params = {
+      ...defaultParams,
+      assistantTools: [
+        { ...mockTool, name: 'test-tool' },
+        { ...mockTool, name: 'test-tool2' },
+      ],
+    };
+    await callAssistantGraph(params);
+
+    expect(getPromptMock).toHaveBeenCalledTimes(3);
+    expect(getPromptMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'test-model',
+        provider: 'openai',
+        promptId: 'test-tool',
+        promptGroupId: toolsGroupId,
+      })
+    );
+    expect(getPromptMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'test-model',
+        provider: 'openai',
+        promptId: 'test-tool2',
+        promptGroupId: toolsGroupId,
+      })
+    );
+    expect(getPromptMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        model: 'test-model',
+        provider: 'openai',
+        promptId: promptDictionary.systemPrompt,
+        promptGroupId: promptGroupId.aiAssistant,
+      })
+    );
+
+    expect(getTool).toHaveBeenCalledWith(
+      expect.objectContaining({
+        description: 'prompt',
+      })
+    );
+  });
+
   describe('agentRunnable', () => {
     it('creates OpenAIToolsAgent for openai llmType', async () => {
       const params = { ...defaultParams, llmType: 'openai' };

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts

@@ -15,10 +15,11 @@ import {
 import { APMTracer } from '@kbn/langchain/server/tracers/apm';
 import { TelemetryTracer } from '@kbn/langchain/server/tracers/telemetry';
 import { pruneContentReferences, MessageMetadata } from '@kbn/elastic-assistant-common';
-import { resolveProviderAndModel } from '@kbn/security-ai-prompts';
+import { getPrompt, resolveProviderAndModel } from '@kbn/security-ai-prompts';
+import { localToolPrompts, promptGroupId as toolsGroupId } from '../../../prompt/tool_prompts';
 import { promptGroupId } from '../../../prompt/local_prompt_object';
 import { getModelOrOss } from '../../../prompt/helpers';
-import { getPrompt, promptDictionary } from '../../../prompt';
+import { getPrompt as localGetPrompt, promptDictionary } from '../../../prompt';
 import { getLlmClass } from '../../../../routes/utils';
 import { EsAnonymizationFieldsSchema } from '../../../../ai_assistant_data_clients/anonymization_fields/types';
 import { AssistantToolParams } from '../../../../types';
@@ -124,9 +125,33 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     telemetry,
   };
 
-  const tools: StructuredTool[] = assistantTools.flatMap(
-    (tool) => tool.getTool({ ...assistantToolParams, llm: createLlmInstance(), isOssModel }) ?? []
-  );
+  const tools: StructuredTool[] = (
+    await Promise.all(
+      assistantTools.map(async (tool) => {
+        let description: string | undefined;
+        try {
+          description = await getPrompt({
+            actionsClient,
+            connectorId,
+            localPrompts: localToolPrompts,
+            model: getModelOrOss(llmType, isOssModel, request.body.model),
+            promptId: tool.name,
+            promptGroupId: toolsGroupId,
+            provider: llmType,
+            savedObjectsClient,
+          });
+        } catch (e) {
+          logger.error(`Failed to get prompt for tool: ${tool.name}`);
+        }
+        return tool.getTool({
+          ...assistantToolParams,
+          llm: createLlmInstance(),
+          isOssModel,
+          description,
+        });
+      })
+    )
+  ).filter((e) => e != null) as StructuredTool[];
 
   // If KB enabled, fetch for any KB IndexEntries and generate a tool for each
   if (isEnabledKnowledgeBase) {
@@ -139,7 +164,7 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     }
   }
 
-  const defaultSystemPrompt = await getPrompt({
+  const defaultSystemPrompt = await localGetPrompt({
     actionsClient,
     connectorId,
     model: getModelOrOss(llmType, isOssModel, request.body.model),
@@ -176,7 +201,7 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
   const telemetryTracer = telemetryParams
     ? new TelemetryTracer(
         {
-          elasticTools: assistantTools.map(({ name }) => name),
+          elasticTools: tools.map(({ name }) => name),
           totalTools: tools.length,
           telemetry,
           telemetryParams,

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/nodes/generate_chat_title.ts

@@ -8,45 +8,22 @@ import { StringOutputParser } from '@langchain/core/output_parsers';
 
 import { ChatPromptTemplate } from '@langchain/core/prompts';
 import { BaseChatModel } from '@langchain/core/language_models/chat_models';
+import { getPrompt, promptDictionary } from '../../../../prompt';
 import { AgentState, NodeParamsBase } from '../types';
 import { NodeType } from '../constants';
+import { promptGroupId } from '../../../../prompt/local_prompt_object';
 
-export const GENERATE_CHAT_TITLE_PROMPT = (responseLanguage: string, llmType?: string) =>
-  llmType === 'bedrock'
-    ? ChatPromptTemplate.fromMessages([
-        [
-          'system',
-          `You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Please create the title in ${responseLanguage}. Respond with the title only with no other text explaining your response. As an example, for the given MESSAGE, this is the TITLE:
-
-    MESSAGE: I am having trouble with the Elastic Security app.
-    TITLE: Troubleshooting Elastic Security app issues
-    `,
-        ],
-        ['human', '{input}'],
-      ])
-    : llmType === 'gemini'
-    ? ChatPromptTemplate.fromMessages([
-        [
-          'system',
-          `You are a title generator for a helpful assistant for Elastic Security. Assume the following human message is the start of a conversation between you and a human. Generate a relevant conversation title for the human's message in plain text. Make sure the title is formatted for the user, without using quotes or markdown. The title should clearly reflect the content of the message and be appropriate for a list of conversations. Please create the title in ${responseLanguage}. Respond only with the title. As an example, for the given MESSAGE, this is the TITLE:
-
-    MESSAGE: I am having trouble with the Elastic Security app.
-    TITLE: Troubleshooting Elastic Security app issues
-    `,
-        ],
-        ['human', '{input}'],
-      ])
-    : ChatPromptTemplate.fromMessages([
-        [
-          'system',
-          `You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Please create the title in ${responseLanguage}. As an example, for the given MESSAGE, this is the TITLE:
-
-    MESSAGE: I am having trouble with the Elastic Security app.
-    TITLE: Troubleshooting Elastic Security app issues
-    `,
-        ],
-        ['human', '{input}'],
-      ]);
+export const GENERATE_CHAT_TITLE_PROMPT = ({
+  prompt,
+  responseLanguage,
+}: {
+  prompt: string;
+  responseLanguage: string;
+}) =>
+  ChatPromptTemplate.fromMessages([
+    ['system', `${prompt}\nPlease create the title in ${responseLanguage}.`],
+    ['human', '{input}'],
+  ]);
 
 export interface GenerateChatTitleParams extends NodeParamsBase {
   state: AgentState;
@@ -54,7 +31,9 @@ export interface GenerateChatTitleParams extends NodeParamsBase {
 }
 
 export async function generateChatTitle({
+  actionsClient,
   logger,
+  savedObjectsClient,
   state,
   model,
 }: GenerateChatTitleParams): Promise<Partial<AgentState>> {
@@ -64,7 +43,15 @@ export async function generateChatTitle({
     );
 
     const outputParser = new StringOutputParser();
-    const graph = GENERATE_CHAT_TITLE_PROMPT(state.responseLanguage, state.llmType)
+    const prompt = await getPrompt({
+      actionsClient,
+      connectorId: state.connectorId,
+      promptId: promptDictionary.chatTitle,
+      promptGroupId: promptGroupId.aiAssistant,
+      provider: state.llmType,
+      savedObjectsClient,
+    });
+    const graph = GENERATE_CHAT_TITLE_PROMPT({ prompt, responseLanguage: state.responseLanguage })
       .pipe(model)
       .pipe(outputParser);
 

x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/local_prompt_object.ts

@@ -21,6 +21,9 @@ import {
   GEMINI_SYSTEM_PROMPT,
   GEMINI_USER_PROMPT,
   STRUCTURED_SYSTEM_PROMPT,
+  BEDROCK_CHAT_TITLE,
+  GEMINI_CHAT_TITLE,
+  DEFAULT_CHAT_TITLE,
 } from './prompts';
 
 export const promptGroupId = {
@@ -31,6 +34,7 @@ export const promptGroupId = {
 export const promptDictionary = {
   systemPrompt: `systemPrompt`,
   userPrompt: `userPrompt`,
+  chatTitle: `chatTitle`,
   attackDiscoveryDefault: `default`,
   attackDiscoveryRefine: `refine`,
   attackDiscoveryContinue: `continue`,
@@ -154,4 +158,27 @@ export const localPrompts: Prompt[] = [
       default: ATTACK_DISCOVERY_GENERATION_INSIGHTS,
     },
   },
+  {
+    promptId: promptDictionary.chatTitle,
+    promptGroupId: promptGroupId.aiAssistant,
+    prompt: {
+      default: DEFAULT_CHAT_TITLE,
+    },
+  },
+  {
+    promptId: promptDictionary.chatTitle,
+    promptGroupId: promptGroupId.aiAssistant,
+    provider: 'bedrock',
+    prompt: {
+      default: BEDROCK_CHAT_TITLE,
+    },
+  },
+  {
+    promptId: promptDictionary.chatTitle,
+    promptGroupId: promptGroupId.aiAssistant,
+    provider: 'gemini',
+    prompt: {
+      default: GEMINI_CHAT_TITLE,
+    },
+  },
 ];

x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts

@@ -131,3 +131,21 @@ export const ATTACK_DISCOVERY_GENERATION_SUMMARY_MARKDOWN = `A markdown summary
 export const ATTACK_DISCOVERY_GENERATION_TITLE =
   'A short, no more than 7 words, title for the insight, NOT formatted with special syntax or markdown. This must be as brief as possible.';
 export const ATTACK_DISCOVERY_GENERATION_INSIGHTS = `Insights with markdown that always uses special ${SYNTAX} syntax for field names and values from the source data. ${GOOD_SYNTAX_EXAMPLES} ${BAD_SYNTAX_EXAMPLES}`;
+
+export const BEDROCK_CHAT_TITLE = `You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. Respond with the title only with no other text explaining your response. As an example, for the given MESSAGE, this is the TITLE:
+
+MESSAGE: I am having trouble with the Elastic Security app.
+TITLE: Troubleshooting Elastic Security app issues
+`;
+
+export const GEMINI_CHAT_TITLE = `You are a title generator for a helpful assistant for Elastic Security. Assume the following human message is the start of a conversation between you and a human. Generate a relevant conversation title for the human's message in plain text. Make sure the title is formatted for the user, without using quotes or markdown. The title should clearly reflect the content of the message and be appropriate for a list of conversations. Respond only with the title. As an example, for the given MESSAGE, this is the TITLE:
+
+MESSAGE: I am having trouble with the Elastic Security app.
+TITLE: Troubleshooting Elastic Security app issues
+`;
+
+export const DEFAULT_CHAT_TITLE = `You are a helpful assistant for Elastic Security. Assume the following user message is the start of a conversation between you and a user; give this conversation a title based on the content below. DO NOT UNDER ANY CIRCUMSTANCES wrap this title in single or double quotes. This title is shown in a list of conversations to the user, so title it for the user, not for you. As an example, for the given MESSAGE, this is the TITLE:
+
+MESSAGE: I am having trouble with the Elastic Security app.
+TITLE: Troubleshooting Elastic Security app issues
+`;

x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/tool_prompts.ts

@@ -0,0 +1,82 @@
+/*
+ * 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 { Prompt } from '@kbn/security-ai-prompts';
+
+export const promptGroupId = 'security-tools';
+
+// promptId must match tool name
+// cannot import from security_solution because it causes a circular dependency
+export const localToolPrompts: Prompt[] = [
+  {
+    promptId: 'AlertCountsTool',
+    promptGroupId,
+    prompt: {
+      default:
+        'Call this for the counts of last 24 hours of open and acknowledged alerts in the environment, grouped by their severity and workflow status. The response will be JSON and from it you can summarize the information to answer the question.',
+    },
+  },
+  {
+    promptId: 'NaturalLanguageESQLTool',
+    promptGroupId,
+    prompt: {
+      default: `You MUST use the "NaturalLanguageESQLTool" function when the user wants to:
+  - breakdown or filter ES|QL queries that are displayed on the current page
+  - convert queries from another language to ES|QL
+  - asks general questions about ES|QL
+
+  ALWAYS use this tool to generate ES|QL queries or explain anything about the ES|QL query language rather than coming up with your own answer.`,
+    },
+  },
+  {
+    promptId: 'ProductDocumentationTool',
+    promptGroupId,
+    prompt: {
+      default:
+        'Use this tool to retrieve documentation about Elastic products. You can retrieve documentation about the Elastic stack, such as Kibana and Elasticsearch, or for Elastic solutions, such as Elastic Security, Elastic Observability or Elastic Enterprise Search.',
+    },
+  },
+  {
+    promptId: 'KnowledgeBaseRetrievalTool',
+    promptGroupId,
+    prompt: {
+      default:
+        "Call this for fetching details from the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Call this function when the user asks for information about themself, like 'what is my favorite...' or 'using my saved....'. Input must always be the free-text query on a single line, with no other text. You are welcome to re-write the query to be a summary of items/things to search for in the knowledge base, as a vector search will be performed to return similar results when requested. If the results returned do not look relevant, disregard and tell the user you were unable to find the information they were looking for. All requests include a `knowledge history` section which includes some existing knowledge of the user. DO NOT CALL THIS FUNCTION if the `knowledge history` sections appears to be able to answer the user's query.",
+    },
+  },
+  {
+    promptId: 'KnowledgeBaseWriteTool',
+    promptGroupId,
+    prompt: {
+      default:
+        "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, a short UI friendly name for the entry, and whether or not the entry is required.",
+    },
+  },
+  {
+    promptId: 'SecurityLabsKnowledgeBaseTool',
+    promptGroupId,
+    prompt: {
+      default:
+        'Call this for knowledge from Elastic Security Labs content, which contains information on malware, attack techniques, and more.',
+    },
+  },
+  {
+    promptId: 'OpenAndAcknowledgedAlertsTool',
+    promptGroupId,
+    prompt: {
+      default:
+        'Call this for knowledge about the latest n open and acknowledged alerts (sorted by `kibana.alert.risk_score`) in the environment, or when answering questions about open alerts. Do not call this tool for alert count or quantity. The output is an array of the latest n open and acknowledged alerts.',
+    },
+  },
+  {
+    promptId: 'defendInsightsTool',
+    promptGroupId,
+    prompt: {
+      default: 'Call this for Elastic Defend insights.',
+    },
+  },
+];

x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.test.ts

@@ -10,6 +10,7 @@
 import type { AuthenticatedUser } from '@kbn/core-security-common';
 import type { DefendInsightsPostRequestBody } from '@kbn/elastic-assistant-common';
 
+import { getPrompt } from '@kbn/security-ai-prompts';
 import { elasticsearchServiceMock } from '@kbn/core-elasticsearch-server-mocks';
 import { actionsMock } from '@kbn/actions-plugin/server/mocks';
 import { OpenAiProviderType } from '@kbn/stack-connectors-plugin/common/openai/constants';
@@ -29,8 +30,10 @@ import { getAssistantTool, createDefendInsight, isDefendInsightsEnabled } from '
 import { postDefendInsightsRoute } from './post_defend_insights';
 import { licensingMock } from '@kbn/licensing-plugin/public/mocks';
 
+jest.mock('@kbn/security-ai-prompts');
 jest.mock('./helpers');
 
+const getPromptMock = getPrompt as jest.Mock;
 describe('postDefendInsightsRoute', () => {
   let server: ReturnType<typeof serverMock.create>;
   let context: ElasticAssistantRequestHandlerContextMock;
@@ -80,6 +83,7 @@ describe('postDefendInsightsRoute', () => {
       langSmithApiKey: 'langSmithApiKey',
     };
   }
+  const getTool = jest.fn();
 
   beforeEach(() => {
     const tools = requestContextMock.createTools();
@@ -94,7 +98,8 @@ describe('postDefendInsightsRoute', () => {
     mockDataClient = getDefaultDataClient();
     mockApiConfig = getDefaultApiConfig();
     mockRequestBody = getDefaultRequestBody();
-    (getAssistantTool as jest.Mock).mockReturnValue({ getTool: jest.fn() });
+    getPromptMock.mockResolvedValue('prompt');
+    (getAssistantTool as jest.Mock).mockReturnValue({ getTool, name: 'test-tool' });
     (createDefendInsight as jest.Mock).mockResolvedValue({
       currentInsight: mockCurrentInsight,
       defendInsightId: mockCurrentInsight.id,
@@ -196,4 +201,23 @@ describe('postDefendInsightsRoute', () => {
       status_code: 500,
     });
   });
+
+  it('should call getPrompt for tool description', async () => {
+    await server.inject(
+      postDefendInsightsRequest(mockRequestBody),
+      requestContextMock.convertContext(context)
+    );
+    expect(getPromptMock).toHaveBeenCalledWith(
+      expect.objectContaining({
+        connectorId: 'connector-id',
+        promptId: 'test-tool',
+        promptGroupId: 'security-tools',
+      })
+    );
+    expect(getTool).toHaveBeenCalledWith(
+      expect.objectContaining({
+        description: 'prompt',
+      })
+    );
+  });
 });

x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts

@@ -20,6 +20,8 @@ import {
 import { transformError } from '@kbn/securitysolution-es-utils';
 import { IRouter, Logger } from '@kbn/core/server';
 
+import { getPrompt } from '@kbn/security-ai-prompts';
+import { localToolPrompts, promptGroupId } from '../../lib/prompt/tool_prompts';
 import { buildResponse } from '../../lib/build_response';
 import { ElasticAssistantRequestHandlerContext } from '../../types';
 import { DEFAULT_PLUGIN_NAME, getPluginNameFromRequest } from '../helpers';
@@ -71,6 +73,7 @@ export const postDefendInsightsRoute = (router: IRouter<ElasticAssistantRequestH
         const resp = buildResponse(response);
 
         const ctx = await context.resolve(['licensing', 'elasticAssistant']);
+        const savedObjectsClient = ctx.elasticAssistant.savedObjectsClient;
 
         const assistantContext = ctx.elasticAssistant;
 
@@ -159,7 +162,15 @@ export const postDefendInsightsRoute = (router: IRouter<ElasticAssistantRequestH
             request,
           });
 
-          const toolInstance = assistantTool.getTool(assistantToolParams);
+          const description = await getPrompt({
+            actionsClient,
+            connectorId: apiConfig.connectorId,
+            localPrompts: localToolPrompts,
+            promptId: assistantTool.name,
+            promptGroupId,
+            savedObjectsClient,
+          });
+          const toolInstance = assistantTool.getTool({ ...assistantToolParams, description });
 
           const { currentInsight, defendInsightId } = await createDefendInsight(
             endpointIds,

x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts

@@ -13,6 +13,7 @@ import { evaluate } from 'langsmith/evaluation';
 import { v4 as uuidv4 } from 'uuid';
 
 import { getRequestAbortedSignal } from '@kbn/data-plugin/server';
+import { getPrompt } from '@kbn/security-ai-prompts';
 import {
   API_VERSIONS,
   newContentReferencesStore,
@@ -31,6 +32,7 @@ import {
   createToolCallingAgent,
 } from 'langchain/agents';
 import { omit } from 'lodash/fp';
+import { localToolPrompts, promptGroupId as toolsGroupId } from '../../lib/prompt/tool_prompts';
 import { promptGroupId } from '../../lib/prompt/local_prompt_object';
 import { getModelOrOss } from '../../lib/prompt/helpers';
 import { getAttackDiscoveryPrompts } from '../../lib/attack_discovery/graphs/default_attack_discovery_graph/nodes/helpers/prompts';
@@ -38,7 +40,7 @@ import {
   formatPrompt,
   formatPromptStructured,
 } from '../../lib/langchain/graphs/default_assistant_graph/prompts';
-import { getPrompt, promptDictionary } from '../../lib/prompt';
+import { getPrompt as localGetPrompt, promptDictionary } from '../../lib/prompt';
 import { buildResponse } from '../../lib/build_response';
 import { AssistantDataClients } from '../../lib/langchain/executors/types';
 import { AssistantToolParams, ElasticAssistantRequestHandlerContext, GetElser } from '../../types';
@@ -317,11 +319,36 @@ export const postEvaluateRoute = (
                 ...(productDocsAvailable ? { llmTasks: ctx.elasticAssistant.llmTasks } : {}),
               };
 
-              const tools: StructuredTool[] = assistantTools.flatMap(
-                (tool) => tool.getTool(assistantToolParams) ?? []
-              );
+              const tools: StructuredTool[] = (
+                await Promise.all(
+                  assistantTools.map(async (tool) => {
+                    let description: string | undefined;
+                    try {
+                      description = await getPrompt({
+                        actionsClient,
+                        connector,
+                        connectorId: connector.id,
+                        model: getModelOrOss(llmType, isOssModel),
+                        localPrompts: localToolPrompts,
+                        promptId: tool.name,
+                        promptGroupId: toolsGroupId,
+                        provider: llmType,
+                        savedObjectsClient,
+                      });
+                    } catch (e) {
+                      logger.error(`Failed to get prompt for tool: ${tool.name}`);
+                    }
+                    return tool.getTool({
+                      ...assistantToolParams,
+                      llm: createLlmInstance(),
+                      isOssModel,
+                      description,
+                    });
+                  })
+                )
+              ).filter((e) => e != null) as StructuredTool[];
 
-              const defaultSystemPrompt = await getPrompt({
+              const defaultSystemPrompt = await localGetPrompt({
                 actionsClient,
                 connector,
                 connectorId: connector.id,

x-pack/solutions/security/plugins/elastic_assistant/server/types.ts

@@ -240,6 +240,7 @@ export interface AssistantToolParams {
   isEnabledKnowledgeBase: boolean;
   connectorId?: string;
   contentReferencesStore: ContentReferencesStore | undefined;
+  description?: string;
   esClient: ElasticsearchClient;
   kbDataClient?: AIAssistantKnowledgeBaseDataClient;
   langChainTimeout?: number;

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts

@@ -23,6 +23,9 @@ export const ALERT_COUNTS_TOOL_DESCRIPTION =
 export const ALERT_COUNTS_TOOL: AssistantTool = {
   id: 'alert-counts-tool',
   name: 'AlertCountsTool',
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description: ALERT_COUNTS_TOOL_DESCRIPTION,
   sourceRegister: APP_UI_ID,
   isSupported: (params: AssistantToolParams): params is AlertCountsToolParams => {
@@ -35,7 +38,7 @@ export const ALERT_COUNTS_TOOL: AssistantTool = {
       params as AlertCountsToolParams;
     return new DynamicStructuredTool({
       name: 'AlertCountsTool',
-      description: ALERT_COUNTS_TOOL_DESCRIPTION,
+      description: params.description || ALERT_COUNTS_TOOL_DESCRIPTION,
       schema: z.object({}),
       func: async () => {
         const query = getAlertsCountQuery(alertsIndexPattern);

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/defend_insights/index.ts

@@ -41,6 +41,9 @@ export interface DefendInsightsToolParams extends AssistantToolParams {
 export const DEFEND_INSIGHTS_TOOL: AssistantTool = Object.freeze({
   id: DEFEND_INSIGHTS_TOOL_ID,
   name: 'defendInsightsTool',
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description: DEFEND_INSIGHTS_TOOL_DESCRIPTION,
   sourceRegister: APP_UI_ID,
 
@@ -67,7 +70,7 @@ export const DEFEND_INSIGHTS_TOOL: AssistantTool = Object.freeze({
 
     return new DynamicTool({
       name: 'DefendInsightsTool',
-      description: DEFEND_INSIGHTS_TOOL_DESCRIPTION,
+      description: params.description || DEFEND_INSIGHTS_TOOL_DESCRIPTION,
       func: async () => {
         if (llm == null) {
           throw new Error('LLM is required for Defend Insights');

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts

@@ -21,6 +21,9 @@ const TOOL_NAME = 'NaturalLanguageESQLTool';
 const toolDetails = {
   id: 'nl-to-esql-tool',
   name: TOOL_NAME,
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description: `You MUST use the "${TOOL_NAME}" function when the user wants to:
   - breakdown or filter ES|QL queries that are displayed on the current page
   - convert queries from another language to ES|QL
@@ -57,7 +60,8 @@ export const NL_TO_ESQL_TOOL: AssistantTool = {
     return new DynamicStructuredTool({
       name: toolDetails.name,
       description:
-        toolDetails.description + (isOssModel ? getPromptSuffixForOssModel(TOOL_NAME) : ''),
+        (params.description || toolDetails.description) +
+        (isOssModel ? getPromptSuffixForOssModel(TOOL_NAME) : ''),
       schema: z.object({
         question: z.string().describe(`The user's exact question about ESQL`),
       }),

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts

@@ -19,6 +19,9 @@ export interface KnowledgeBaseRetrievalToolParams extends AssistantToolParams {
 }
 
 const toolDetails = {
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description:
     "Call this for fetching details from the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Call this function when the user asks for information about themself, like 'what is my favorite...' or 'using my saved....'. Input must always be the free-text query on a single line, with no other text. You are welcome to re-write the query to be a summary of items/things to search for in the knowledge base, as a vector search will be performed to return similar results when requested. If the results returned do not look relevant, disregard and tell the user you were unable to find the information they were looking for. All requests include a `knowledge history` section which includes some existing knowledge of the user. DO NOT CALL THIS FUNCTION if the `knowledge history` sections appears to be able to answer the user's query.",
   id: 'knowledge-base-retrieval-tool',
@@ -40,7 +43,7 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = {
 
     return new DynamicStructuredTool({
       name: toolDetails.name,
-      description: toolDetails.description,
+      description: params.description || toolDetails.description,
       schema: z.object({
         query: z.string().describe(`Summary of items/things to search for in the knowledge base`),
       }),

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts

@@ -20,6 +20,9 @@ export interface KnowledgeBaseWriteToolParams extends AssistantToolParams {
 }
 
 const toolDetails = {
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description:
     "Call this for writing details to the user's knowledge base. The knowledge base contains useful information the user wants to store between conversation contexts. Input will be the summarized knowledge base entry to store, a short UI friendly name for the entry, and whether or not the entry is required.",
   id: 'knowledge-base-write-tool',
@@ -40,7 +43,7 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = {
 
     return new DynamicStructuredTool({
       name: toolDetails.name,
-      description: toolDetails.description,
+      description: params.description || toolDetails.description,
       schema: z.object({
         name: z
           .string()

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts

@@ -37,6 +37,9 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION =
 export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = {
   id: 'open-and-acknowledged-alerts-tool',
   name: 'OpenAndAcknowledgedAlertsTool',
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION,
   sourceRegister: APP_UI_ID,
   isSupported: (params: AssistantToolParams): params is OpenAndAcknowledgedAlertsToolParams => {
@@ -62,7 +65,7 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = {
     } = params as OpenAndAcknowledgedAlertsToolParams;
     return new DynamicStructuredTool({
       name: 'OpenAndAcknowledgedAlertsTool',
-      description: OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION,
+      description: params.description || OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL_DESCRIPTION,
       schema: z.object({}),
       func: async () => {
         const query = getOpenAndAcknowledgedAlertsQuery({

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts

@@ -18,6 +18,9 @@ import type { RetrieveDocumentationResultDoc } from '@kbn/llm-tasks-plugin/serve
 import { APP_UI_ID } from '../../../../common';
 
 const toolDetails = {
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description:
     'Use this tool to retrieve documentation about Elastic products. You can retrieve documentation about the Elastic stack, such as Kibana and Elasticsearch, or for Elastic solutions, such as Elastic Security, Elastic Observability or Elastic Enterprise Search.',
   id: 'product-documentation-tool',
@@ -40,7 +43,7 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = {
 
     return new DynamicStructuredTool({
       name: toolDetails.name,
-      description: toolDetails.description,
+      description: params.description || toolDetails.description,
       schema: z.object({
         query: z.string().describe(
           `The query to use to retrieve documentation

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts

@@ -14,6 +14,9 @@ import { knowledgeBaseReference, contentReferenceString } from '@kbn/elastic-ass
 import { APP_UI_ID } from '../../../../common';
 
 const toolDetails = {
+  // note: this description is overwritten when `getTool` is called
+  // local definitions exist ../elastic_assistant/server/lib/prompt/tool_prompts.ts
+  // local definitions can be overwritten by security-ai-prompt integration definitions
   description:
     'Call this for knowledge from Elastic Security Labs content, which contains information on malware, attack techniques, and more.',
   id: 'security-labs-knowledge-base-tool',
@@ -34,7 +37,7 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = {
 
     return new DynamicStructuredTool({
       name: toolDetails.name,
-      description: toolDetails.description,
+      description: params.description || toolDetails.description,
       schema: z.object({
         question: z
           .string()