Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Security Assistant] Fix ESQL tool availability #195827

Merged
merged 2 commits into from
Oct 10, 2024
Merged

[Security Assistant] Fix ESQL tool availability #195827

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a37e905
fix
stephmilovic Oct 10, 2024
08b032d
rm from test files
stephmilovic Oct 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion ...ck/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,6 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
isEnabledKnowledgeBase,
kbDataClient: dataClients?.kbDataClient,
logger,
modelExists: isEnabledKnowledgeBase,
onNewReplacements,
replacements,
request,
Expand Down
1 change: 0 additions & 1 deletion x-pack/plugins/elastic_assistant/server/routes/attack_discovery/helpers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,6 @@ const formatAssistantToolParams = ({
langChainTimeout,
llm,
logger,
modelExists: false, // not required for attack discovery
onNewReplacements,
replacements: latestReplacements,
request,
Expand Down
1 change: 0 additions & 1 deletion x-pack/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -236,7 +236,6 @@ export const postEvaluateRoute = (
llm,
isOssModel,
logger,
modelExists: isEnabledKnowledgeBase,
request: skeletonRequest,
alertsIndexPattern,
// onNewReplacements,
Expand Down
1 change: 0 additions & 1 deletion x-pack/plugins/elastic_assistant/server/types.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,6 @@ export interface AssistantToolParams {
llm?: ActionsClientLlm | AssistantToolLlm;
isOssModel?: boolean;
logger: Logger;
modelExists: boolean;
onNewReplacements?: (newReplacements: Replacements) => void;
replacements?: Replacements;
request: KibanaRequest<
Expand Down
52 changes: 4 additions & 48 deletions x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,61 +40,17 @@ describe('NaturalLanguageESQLTool', () => {
request,
inference,
connectorId,
isEnabledKnowledgeBase: true,
modelExists: true,
stephmilovic marked this conversation as resolved.
Show resolved Hide resolved
};

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,
Expand Down
5 changes: 3 additions & 2 deletions x-pack/plugins/security_solution/server/assistant/tools/esql/nl_to_esql_tool.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions .../security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions ...gins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_write_tool.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 2 additions & 2 deletions x-pack/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down