diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts index bbc19d0566d70..65beab7a9edb3 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock.ts @@ -7,7 +7,7 @@ import { ContentReferencesStore } from '../../types'; -export const contentReferencesStoreFactoryMock: () => ContentReferencesStore = jest +export const newContentReferencesStoreMock: () => ContentReferencesStore = jest .fn() .mockReturnValue({ add: jest.fn().mockImplementation((creator: Parameters[0]) => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.test.ts similarity index 90% rename from x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts rename to x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.test.ts index c3e848ee66d5c..7346ee7f28765 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.test.ts @@ -5,14 +5,14 @@ * 2.0. */ -import { contentReferencesStoreFactory } from './content_references_store_factory'; +import { newContentReferencesStore } from './content_references_store'; import { securityAlertsPageReference } from '../references'; import { ContentReferencesStore } from '../types'; -describe('contentReferencesStoreFactory', () => { +describe('newContentReferencesStore', () => { let contentReferencesStore: ContentReferencesStore; beforeEach(() => { - contentReferencesStore = contentReferencesStoreFactory(); + contentReferencesStore = newContentReferencesStore(); }); it('adds multiple content reference', async () => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.ts similarity index 94% rename from x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts rename to x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.ts index b1b78263a31f3..e68df0f39467b 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store.ts @@ -14,7 +14,7 @@ const CONTENT_REFERENCE_ID_ALPHABET = /** * Creates a new ContentReferencesStore used for storing references (also known as citations) */ -export const contentReferencesStoreFactory: () => ContentReferencesStore = () => { +export const newContentReferencesStore: () => ContentReferencesStore = () => { const store = new Map(); const add: ContentReferencesStore['add'] = (creator) => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts index 5f6184144315c..7c47c7767e12a 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts @@ -9,12 +9,12 @@ import { pruneContentReferences } from './prune_content_references'; import { securityAlertsPageReference } from '../references'; import { contentReferenceBlock } from '../references/utils'; import { ContentReferencesStore } from '../types'; -import { contentReferencesStoreFactory } from './content_references_store_factory'; +import { newContentReferencesStore } from './content_references_store'; describe('pruneContentReferences', () => { let contentReferencesStore: ContentReferencesStore; beforeEach(() => { - contentReferencesStore = contentReferencesStoreFactory(); + contentReferencesStore = newContentReferencesStore(); }); it('prunes content references correctly', async () => { diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts index ca9afcbb76bbe..fa9872f316966 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/index.ts @@ -5,7 +5,7 @@ * 2.0. */ -export { contentReferencesStoreFactory } from './content_references_store/content_references_store_factory'; +export { newContentReferencesStore } from './content_references_store/content_references_store'; export { pruneContentReferences } from './content_references_store/prune_content_references'; export { securityAlertReference, diff --git a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts index 6c5e52df671e1..57ff1a69830cc 100644 --- a/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts +++ b/x-pack/platform/packages/shared/kbn-elastic-assistant-common/index.ts @@ -23,7 +23,7 @@ export { } from './impl/data_anonymization/helpers'; export { - contentReferencesStoreFactory, + newContentReferencesStore, securityAlertReference, knowledgeBaseReference, securityAlertsPageReference, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx index fbf3d310cbe80..fe4f1cd562162 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.test.tsx @@ -16,7 +16,7 @@ import { import { authenticatedUser } from '../../__mocks__/user'; import { getCreateKnowledgeBaseEntrySchemaMock } from '../../__mocks__/knowledge_base_entry_schema.mock'; import { ContentReferencesStore, IndexEntry } from '@kbn/elastic-assistant-common'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; import { isString } from 'lodash'; // Mock dependencies @@ -146,7 +146,7 @@ describe('getStructuredToolForIndexEntry', () => { const mockEsClient = {} as ElasticsearchClient; const mockIndexEntry = getCreateKnowledgeBaseEntrySchemaMock({ type: 'index' }) as IndexEntry; - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); it('should return a DynamicStructuredTool with correct name and schema', () => { const tool = getStructuredToolForIndexEntry({ diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts index fc282b9a9f611..ab03e525c27fe 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/helpers.ts @@ -153,7 +153,7 @@ export const getStructuredToolForIndexEntry = ({ }: { indexEntry: IndexEntry; esClient: ElasticsearchClient; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; logger: Logger; }): DynamicStructuredTool => { const inputSchema = indexEntry.inputSchema?.reduce((prev, input) => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts index 6357224cc7d85..b21ac92170e73 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.test.ts @@ -27,7 +27,7 @@ import { getSecurityLabsDocsCount, } from '../../lib/langchain/content_loaders/security_labs_loader'; import { DynamicStructuredTool } from '@langchain/core/tools'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; jest.mock('../../lib/langchain/content_loaders/security_labs_loader'); jest.mock('p-retry'); const date = '2023-03-28T22:27:28.159Z'; @@ -522,7 +522,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, - contentReferencesStore: contentReferencesStoreFactoryMock(), + contentReferencesStore: newContentReferencesStoreMock(), }); expect(result).toHaveLength(1); @@ -537,7 +537,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, - contentReferencesStore: contentReferencesStoreFactoryMock(), + contentReferencesStore: newContentReferencesStoreMock(), }); expect(result).toEqual([]); @@ -550,7 +550,7 @@ describe('AIAssistantKnowledgeBaseDataClient', () => { const result = await client.getAssistantTools({ esClient: esClientMock, - contentReferencesStore: contentReferencesStoreFactoryMock(), + contentReferencesStore: newContentReferencesStoreMock(), }); expect(result).toEqual([]); diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts index 11edd3ad0bf02..561467f2256ea 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts @@ -817,7 +817,7 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient { contentReferencesStore, esClient, }: { - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; esClient: ElasticsearchClient; }): Promise => { const user = this.options.currentUser; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts index 0ccc7b6453684..e1133aa0d9bb0 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts @@ -49,7 +49,7 @@ export interface AgentExecutorParams { assistantTools?: AssistantTool[]; connectorId: string; conversationId?: string; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; dataClients?: AssistantDataClients; esClient: ElasticsearchClient; langChainMessages: BaseMessage[]; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts index f224212f54ce5..5bcf61ec07a08 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.test.ts @@ -19,7 +19,7 @@ import { createStructuredChatAgent, createToolCallingAgent, } from 'langchain/agents'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +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 { AssistantTool, AssistantToolParams } from '../../../..'; import { promptGroupId as toolsGroupId } from '../../../prompt/tool_prompts'; @@ -84,7 +84,7 @@ describe('callAssistantGraph', () => { telemetryParams: {}, traceOptions: {}, responseLanguage: 'English', - contentReferencesStore: contentReferencesStoreFactoryMock(), + contentReferencesStore: newContentReferencesStoreMock(), } as unknown as AgentExecutorParams; beforeEach(() => { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts index 2623f3ab60b75..1dc44448f3e38 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts @@ -16,7 +16,7 @@ import { transformRawData, getAnonymizedValue, ConversationResponse, - contentReferencesStoreFactory, + newContentReferencesStore, pruneContentReferences, } from '@kbn/elastic-assistant-common'; import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common'; @@ -186,8 +186,9 @@ export const chatCompleteRoute = ( })); } - const contentReferencesStore = - contentReferencesEnabled && contentReferencesStoreFactory(); + const contentReferencesStore = contentReferencesEnabled + ? newContentReferencesStore() + : undefined; const onLlmResponse = async ( content: string, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts index db7c2d058069f..88ba28c06b8c7 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/helpers.ts @@ -124,7 +124,7 @@ export function getAssistantToolParams({ langSmithProject?: string; langSmithApiKey?: string; logger: Logger; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; latestReplacements: Replacements; onNewReplacements: (newReplacements: Replacements) => void; request: KibanaRequest; @@ -136,7 +136,7 @@ export function getAssistantToolParams({ langChainTimeout: number; llm: ActionsClientLlm; logger: Logger; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; replacements: Replacements; onNewReplacements: (newReplacements: Replacements) => void; request: KibanaRequest; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts index 2adbfe8e18fe5..c8ba9e9819b19 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts @@ -152,7 +152,7 @@ export const postDefendInsightsRoute = (router: IRouter { const alertsIndexPattern = 'alerts-index'; @@ -32,7 +32,7 @@ describe('AlertCountsTool', () => { const isEnabledKnowledgeBase = true; const chain = {} as unknown as RetrievalQAChain; const logger = loggerMock.create(); - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); const rest = { isEnabledKnowledgeBase, chain, @@ -191,7 +191,7 @@ describe('AlertCountsTool', () => { replacements, request, ...rest, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicTool; const result = await tool.func(''); diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts index 4a6683624b980..6218c4c441a44 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.test.ts @@ -12,13 +12,13 @@ import type { ContentReferencesStore, KnowledgeBaseEntryContentReference, } from '@kbn/elastic-assistant-common'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; import { loggerMock } from '@kbn/logging-mocks'; import { Document } from 'langchain/document'; describe('KnowledgeBaseRetievalTool', () => { const logger = loggerMock.create(); - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); const getKnowledgeBaseDocumentEntries = jest.fn(); const kbDataClient = { getKnowledgeBaseDocumentEntries }; const defaultArgs = { @@ -68,7 +68,7 @@ describe('KnowledgeBaseRetievalTool', () => { it('does not include citations if contentReferenceStore is false', async () => { const tool = KNOWLEDGE_BASE_RETRIEVAL_TOOL.getTool({ ...defaultArgs, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicStructuredTool; getKnowledgeBaseDocumentEntries.mockResolvedValue([ diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts index 5cdc59a5fca94..aeda113abc2f0 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.test.ts @@ -18,7 +18,7 @@ import type { ContentReferencesStore, SecurityAlertContentReference, } from '@kbn/elastic-assistant-common'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; const MAX_SIZE = 10000; @@ -44,7 +44,7 @@ describe('OpenAndAcknowledgedAlertsTool', () => { const isEnabledKnowledgeBase = true; const chain = {} as unknown as RetrievalQAChain; const logger = loggerMock.create(); - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); const rest = { isEnabledKnowledgeBase, esClient, @@ -274,7 +274,7 @@ describe('OpenAndAcknowledgedAlertsTool', () => { request, size: request.body.size, ...rest, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicTool; (esClient.search as jest.Mock).mockResolvedValue({ diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts index 11a95bd29b594..8d803f3bbaf25 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.test.ts @@ -20,7 +20,7 @@ import type { ContentReferencesStore, ProductDocumentationContentReference, } from '@kbn/elastic-assistant-common'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; describe('ProductDocumentationTool', () => { const chain = {} as RetrievalQAChain; @@ -35,7 +35,7 @@ describe('ProductDocumentationTool', () => { retrieveDocumentationAvailable: jest.fn(), } as LlmTasksPluginStart; const connectorId = 'fake-connector'; - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); const defaultArgs = { chain, esClient, @@ -144,7 +144,7 @@ describe('ProductDocumentationTool', () => { it('does not include citations if contentReferencesStore is false', async () => { const tool = PRODUCT_DOCUMENTATION_TOOL.getTool({ ...defaultArgs, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicStructuredTool; (retrieveDocumentation as jest.Mock).mockResolvedValue({ diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts index c8bba6e3a7113..4e7df37d4f904 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.test.ts @@ -11,11 +11,11 @@ import type { ContentReferencesStore, KnowledgeBaseEntryContentReference, } from '@kbn/elastic-assistant-common'; -import { contentReferencesStoreFactoryMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; +import { newContentReferencesStoreMock } from '@kbn/elastic-assistant-common/impl/content_references/content_references_store/__mocks__/content_references_store.mock'; import type { AssistantToolParams } from '@kbn/elastic-assistant-plugin/server'; describe('SecurityLabsTool', () => { - const contentReferencesStore = contentReferencesStoreFactoryMock(); + const contentReferencesStore = newContentReferencesStoreMock(); const getKnowledgeBaseDocumentEntries = jest.fn().mockResolvedValue([]); const kbDataClient = { getKnowledgeBaseDocumentEntries }; const defaultArgs = { @@ -54,7 +54,7 @@ describe('SecurityLabsTool', () => { it('does not include citations when contentReferencesStore is false', async () => { const tool = SECURITY_LABS_KNOWLEDGE_BASE_TOOL.getTool({ ...defaultArgs, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicStructuredTool; const result = await tool.func({ query: 'What is Kibana Security?', product: 'kibana' });