[8.x] [Security Assistant] Fix ESQL tool availability (#195827) (#195838

)

# Backport

This will backport the following commits from `main` to `8.x`:
- [[Security Assistant] Fix ESQL tool availability
(#195827)](#195827)

<!--- Backport version: 9.4.3 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Steph
Milovic","email":"[email protected]"},"sourceCommit":{"committedDate":"2024-10-10T22:49:22Z","message":"[Security
Assistant] Fix ESQL tool availability
(#195827)","sha":"2b995fa86eb44a2bd54c44a74eb47a2a26ec0ed2","branchLabelMapping":{"^v9.0.0$":"main","^v8.16.0$":"8.x","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","v9.0.0","Team:
SecuritySolution","Team:Security Generative
AI","v8.16.0","backport:version"],"title":"[Security Assistant] Fix ESQL
tool
availability","number":195827,"url":"https://github.com/elastic/kibana/pull/195827","mergeCommit":{"message":"[Security
Assistant] Fix ESQL tool availability
(#195827)","sha":"2b995fa86eb44a2bd54c44a74eb47a2a26ec0ed2"}},"sourceBranch":"main","suggestedTargetBranches":["8.x"],"targetPullRequestStates":[{"branch":"main","label":"v9.0.0","branchLabelMappingKey":"^v9.0.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/195827","number":195827,"mergeCommit":{"message":"[Security
Assistant] Fix ESQL tool availability
(#195827)","sha":"2b995fa86eb44a2bd54c44a74eb47a2a26ec0ed2"}},{"branch":"8.x","label":"v8.16.0","branchLabelMappingKey":"^v8.16.0$","isSourceBranch":false,"state":"NOT_CREATED"}]}]
BACKPORT-->

Co-authored-by: Steph Milovic <[email protected]>

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

@@ -103,7 +103,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     isEnabledKnowledgeBase,
     kbDataClient: dataClients?.kbDataClient,
     logger,
-    modelExists: isEnabledKnowledgeBase,
     onNewReplacements,
     replacements,
     request,

x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.test.ts

@@ -196,7 +196,6 @@ describe('helpers', () => {
         langChainTimeout,
         llm,
         logger: mockLogger,
-        modelExists: false,
         onNewReplacements,
         replacements: latestReplacements,
         request: mockRequest,
@@ -231,7 +230,6 @@ describe('helpers', () => {
         langChainTimeout,
         llm,
         logger: mockLogger,
-        modelExists: false,
         onNewReplacements,
         replacements: latestReplacements,
         request: mockRequest,

x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts

@@ -157,7 +157,6 @@ const formatAssistantToolParams = ({
   langChainTimeout,
   llm,
   logger,
-  modelExists: false, // not required for attack discovery
   onNewReplacements,
   replacements: latestReplacements,
   request,

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

@@ -236,7 +236,6 @@ export const postEvaluateRoute = (
                 llm,
                 isOssModel,
                 logger,
-                modelExists: isEnabledKnowledgeBase,
                 request: skeletonRequest,
                 alertsIndexPattern,
                 // onNewReplacements,

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

@@ -244,7 +244,6 @@ export interface AssistantToolParams {
   llm?: ActionsClientLlm | AssistantToolLlm;
   isOssModel?: boolean;
   logger: Logger;
-  modelExists: boolean;
   onNewReplacements?: (newReplacements: Replacements) => void;
   replacements?: Replacements;
   request: KibanaRequest<

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

@@ -29,13 +29,11 @@ describe('AlertCountsTool', () => {
   } as unknown as KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
   const isEnabledKnowledgeBase = true;
   const chain = {} as unknown as RetrievalQAChain;
-  const modelExists = true;
   const logger = loggerMock.create();
   const rest = {
     isEnabledKnowledgeBase,
     chain,
     logger,
-    modelExists,
   };
 
   beforeEach(() => {

x-pack/plugins/security_solution/server/assistant/tools/attack_discovery/attack_discovery_tool.test.ts

@@ -75,7 +75,6 @@ describe('AttackDiscoveryTool', () => {
     isEnabledKnowledgeBase: false,
     llm,
     logger,
-    modelExists: false,
     onNewReplacements: jest.fn(),
     size,
   };

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

@@ -40,65 +40,18 @@ describe('NaturalLanguageESQLTool', () => {
     request,
     inference,
     connectorId,
+    isEnabledKnowledgeBase: true,
   };
 
   describe('isSupported', () => {
-    it('returns false if isEnabledKnowledgeBase is false', () => {
-      const params = {
-        isEnabledKnowledgeBase: false,
-        modelExists: true,
-        ...rest,
-      };
-
-      expect(NL_TO_ESQL_TOOL.isSupported(params)).toBe(false);
-    });
-
-    it('returns false if modelExists is false (the ELSER model is not installed)', () => {
-      const params = {
-        isEnabledKnowledgeBase: true,
-        modelExists: false, // <-- ELSER model is not installed
-        ...rest,
-      };
-
-      expect(NL_TO_ESQL_TOOL.isSupported(params)).toBe(false);
-    });
-
-    it('returns true if isEnabledKnowledgeBase and modelExists are true', () => {
-      const params = {
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
-        ...rest,
-      };
-
-      expect(NL_TO_ESQL_TOOL.isSupported(params)).toBe(true);
+    it('returns true if connectorId and inference have values', () => {
+      expect(NL_TO_ESQL_TOOL.isSupported(rest)).toBe(true);
     });
   });
 
   describe('getTool', () => {
-    it('returns null if isEnabledKnowledgeBase is false', () => {
-      const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: false,
-        modelExists: true,
-        ...rest,
-      });
-
-      expect(tool).toBeNull();
-    });
-
-    it('returns null if modelExists is false (the ELSER model is not installed)', () => {
-      const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: false, // <-- ELSER model is not installed
-        ...rest,
-      });
-
-      expect(tool).toBeNull();
-    });
-
     it('returns null if inference plugin is not provided', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         ...rest,
         inference: undefined,
       });
@@ -108,19 +61,15 @@ describe('NaturalLanguageESQLTool', () => {
 
     it('returns null if connectorId is not provided', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         ...rest,
         connectorId: undefined,
       });
 
       expect(tool).toBeNull();
     });
 
-    it('should return a Tool instance if isEnabledKnowledgeBase and modelExists are true', () => {
+    it('should return a Tool instance when given required properties', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         ...rest,
       });
 
@@ -129,8 +78,6 @@ describe('NaturalLanguageESQLTool', () => {
 
     it('should return a tool with the expected tags', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         ...rest,
       }) as DynamicTool;
 
@@ -139,8 +86,6 @@ describe('NaturalLanguageESQLTool', () => {
 
     it('should return tool with the expected description for OSS model', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         isOssModel: true,
         ...rest,
       }) as DynamicTool;
@@ -150,8 +95,6 @@ describe('NaturalLanguageESQLTool', () => {
 
     it('should return tool with the expected description for non-OSS model', () => {
       const tool = NL_TO_ESQL_TOOL.getTool({
-        isEnabledKnowledgeBase: true,
-        modelExists: true,
         isOssModel: false,
         ...rest,
       }) as DynamicTool;

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

@@ -13,6 +13,7 @@ import { naturalLanguageToEsql } from '@kbn/inference-plugin/server';
 import { APP_UI_ID } from '../../../../common';
 import { getPromptSuffixForOssModel } from './common';
 
+// select only some properties of AssistantToolParams
 export type ESQLToolParams = AssistantToolParams;
 
 const TOOL_NAME = 'NaturalLanguageESQLTool';
@@ -32,8 +33,8 @@ export const NL_TO_ESQL_TOOL: AssistantTool = {
   ...toolDetails,
   sourceRegister: APP_UI_ID,
   isSupported: (params: ESQLToolParams): params is ESQLToolParams => {
-    const { inference, connectorId, isEnabledKnowledgeBase, modelExists } = params;
-    return isEnabledKnowledgeBase && modelExists && inference != null && connectorId != null;
+    const { inference, connectorId } = params;
+    return inference != null && connectorId != null;
   },
   getTool(params: ESQLToolParams) {
     if (!this.isSupported(params)) return null;

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

@@ -25,8 +25,8 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = {
   ...toolDetails,
   sourceRegister: APP_UI_ID,
   isSupported: (params: AssistantToolParams): params is KnowledgeBaseRetrievalToolParams => {
-    const { kbDataClient, isEnabledKnowledgeBase, modelExists } = params;
-    return isEnabledKnowledgeBase && modelExists && kbDataClient != null;
+    const { kbDataClient, isEnabledKnowledgeBase } = params;
+    return isEnabledKnowledgeBase && kbDataClient != null;
   },
   getTool(params: AssistantToolParams) {
     if (!this.isSupported(params)) return null;

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

@@ -28,8 +28,8 @@ export const KNOWLEDGE_BASE_WRITE_TOOL: AssistantTool = {
   ...toolDetails,
   sourceRegister: APP_UI_ID,
   isSupported: (params: AssistantToolParams): params is KnowledgeBaseWriteToolParams => {
-    const { isEnabledKnowledgeBase, kbDataClient, modelExists } = params;
-    return isEnabledKnowledgeBase && modelExists && kbDataClient != null;
+    const { isEnabledKnowledgeBase, kbDataClient } = params;
+    return isEnabledKnowledgeBase && kbDataClient != null;
   },
   getTool(params: AssistantToolParams) {
     if (!this.isSupported(params)) return null;

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

@@ -32,14 +32,12 @@ describe('OpenAndAcknowledgedAlertsTool', () => {
   } as unknown as KibanaRequest<unknown, unknown, ExecuteConnectorRequestBody>;
   const isEnabledKnowledgeBase = true;
   const chain = {} as unknown as RetrievalQAChain;
-  const modelExists = true;
   const logger = loggerMock.create();
   const rest = {
     isEnabledKnowledgeBase,
     esClient,
     chain,
     logger,
-    modelExists,
   };
 
   const anonymizationFields = [

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

@@ -22,8 +22,8 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = {
   ...toolDetails,
   sourceRegister: APP_UI_ID,
   isSupported: (params: AssistantToolParams): params is AssistantToolParams => {
-    const { kbDataClient, isEnabledKnowledgeBase, modelExists } = params;
-    return isEnabledKnowledgeBase && modelExists && kbDataClient != null;
+    const { kbDataClient, isEnabledKnowledgeBase } = params;
+    return isEnabledKnowledgeBase && kbDataClient != null;
   },
   getTool(params: AssistantToolParams) {
     if (!this.isSupported(params)) return null;