Skip to content

Commit

Permalink
[Security Assistant] Fix ESQL tool availability (elastic#195827)
Browse files Browse the repository at this point in the history
  • Loading branch information
stephmilovic authored Oct 10, 2024
1 parent 0004217 commit 2b995fa
Show file tree
Hide file tree
Showing 13 changed files with 13 additions and 80 deletions.
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
Original file line number Diff line number Diff line change
Expand Up @@ -196,7 +196,6 @@ describe('helpers', () => {
langChainTimeout,
llm,
logger: mockLogger,
modelExists: false,
onNewReplacements,
replacements: latestReplacements,
request: mockRequest,
Expand Down Expand Up @@ -231,7 +230,6 @@ describe('helpers', () => {
langChainTimeout,
llm,
logger: mockLogger,
modelExists: false,
onNewReplacements,
replacements: latestReplacements,
request: mockRequest,
Expand Down
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
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
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
Original file line number Diff line number Diff line change
Expand Up @@ -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(() => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -75,7 +75,6 @@ describe('AttackDiscoveryTool', () => {
isEnabledKnowledgeBase: false,
llm,
logger,
modelExists: false,
onNewReplacements: jest.fn(),
size,
};
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
});
Expand All @@ -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,
});

Expand All @@ -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;

Expand All @@ -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;
Expand All @@ -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;
Expand Down
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
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
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
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down
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

0 comments on commit 2b995fa

Please sign in to comment.