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 d38bcfef0e755..5145efdbec8f2 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 @@ -20,7 +20,7 @@ import { EsqlContentReference, 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'; // Mock dependencies jest.mock('@elastic/elasticsearch'); @@ -149,7 +149,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 7d612cd53ae74..174fd1c6d923b 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 @@ -148,7 +148,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 9bd6730417b1d..4755953876b06 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 @@ -18,7 +18,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 { resolveProviderAndModel } from '@kbn/security-ai-prompts'; jest.mock('./graph'); @@ -79,7 +79,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/lib/prompt/prompts.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts index fc3d7d68ecce0..526d65640b887 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/lib/prompt/prompts.ts @@ -23,7 +23,7 @@ export const GEMINI_SYSTEM_PROMPT = `${BASE_GEMINI_PROMPT} ${KB_CATCH} {include_ export const BEDROCK_SYSTEM_PROMPT = `Use tools as often as possible, as they have access to the latest data and syntax. Never return tags in the response, but make sure to include tags content in the response. Do not reflect on the quality of the returned search results in your response. ALWAYS return the exact response from NaturalLanguageESQLTool verbatim in the final response, without adding further description.`; export const GEMINI_USER_PROMPT = `Now, always using the tools at your disposal, step by step, come up with a response to this request:\n\n`; -export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. ${KNOWLEDGE_HISTORY} You have access to the following tools: +export const STRUCTURED_SYSTEM_PROMPT = `Respond to the human as helpfully and accurately as possible. ${KNOWLEDGE_HISTORY} {include_citations_prompt_placeholder} You have access to the following tools: {tools} 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 fb5adfa05bee0..44995d2112732 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 @@ -149,7 +149,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' });