From d5f54ac9ef82fbd7a46abc4f83a0146f023c312f Mon Sep 17 00:00:00 2001 From: Kenneth Kreindler Date: Thu, 30 Jan 2025 10:24:49 +0000 Subject: [PATCH 1/5] rename contentReferencesStoreFactory to better name. Fix typing of optional content references store @e40pud --- .../__mocks__/content_references_store.mock.ts | 2 +- ...e_factory.test.ts => content_references_store.test.ts} | 6 +++--- ...ences_store_factory.ts => content_references_store.ts} | 2 +- .../prune_content_references.test.ts | 4 ++-- .../impl/content_references/index.ts | 2 +- .../packages/shared/kbn-elastic-assistant-common/index.ts | 2 +- .../knowledge_base/helpers.test.tsx | 4 ++-- .../ai_assistant_data_clients/knowledge_base/helpers.ts | 2 +- .../knowledge_base/index.test.ts | 8 ++++---- .../ai_assistant_data_clients/knowledge_base/index.ts | 2 +- .../server/lib/langchain/executors/types.ts | 2 +- .../graphs/default_assistant_graph/index.test.ts | 4 ++-- .../server/routes/chat/chat_complete_route.ts | 4 ++-- .../server/routes/defend_insights/helpers.ts | 4 ++-- .../server/routes/evaluate/post_evaluate.ts | 4 ++-- .../plugins/elastic_assistant/server/routes/helpers.ts | 2 +- .../server/routes/post_actions_connector_execute.ts | 4 ++-- .../tools/alert_counts/alert_counts_tool.test.ts | 4 ++-- .../knowledge_base/knowledge_base_retrieval_tool.test.ts | 4 ++-- .../open_and_acknowledged_alerts_tool.test.ts | 4 ++-- .../tools/product_docs/product_documentation_tool.test.ts | 4 ++-- .../tools/security_labs/security_labs_tool.test.ts | 4 ++-- 22 files changed, 39 insertions(+), 39 deletions(-) rename x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/{content_references_store_factory.test.ts => content_references_store.test.ts} (90%) rename x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/{content_references_store_factory.ts => content_references_store.ts} (94%) 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 f2327130b6fe6..1e2b0df9efe31 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'; jest.mock('./graph'); jest.mock('./helpers'); @@ -76,7 +76,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..a1f1763ca8aee 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'; @@ -187,7 +187,7 @@ export const chatCompleteRoute = ( } const contentReferencesStore = - contentReferencesEnabled && contentReferencesStoreFactory(); + 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/evaluate/post_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index 9a9342f60baa9..d8112c5ba218e 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -15,7 +15,7 @@ import { v4 as uuidv4 } from 'uuid'; import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { API_VERSIONS, - contentReferencesStoreFactory, + newContentReferencesStore, ELASTIC_AI_ASSISTANT_EVALUATE_URL, ExecuteConnectorRequestBody, INTERNAL_API_ACCESS, @@ -293,7 +293,7 @@ export const postEvaluateRoute = ( DEFAULT_PLUGIN_NAME ).contentReferencesEnabled; const contentReferencesStore = - contentReferencesEnabled && contentReferencesStoreFactory(); + contentReferencesEnabled ? newContentReferencesStore() : undefined; // Fetch any applicable tools that the source plugin may have registered const assistantToolParams: AssistantToolParams = { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts index 9adc52ea620e1..d97a6b77f93b3 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/helpers.ts @@ -232,7 +232,7 @@ export interface LangChainExecuteParams { telemetry: AnalyticsServiceSetup; actionTypeId: string; connectorId: string; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; llmTasks?: LlmTasksPluginStart; inference: InferenceServerStart; isOssModel?: boolean; diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index e6df1f2ba3d08..ba217fa7be647 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -12,7 +12,7 @@ import { getRequestAbortedSignal } from '@kbn/data-plugin/server'; import { schema } from '@kbn/config-schema'; import { API_VERSIONS, - contentReferencesStoreFactory, + newContentReferencesStore, ExecuteConnectorRequestBody, Message, Replacements, @@ -120,7 +120,7 @@ export const postActionsConnectorExecuteRoute = ( const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient(); const contentReferencesStore = - contentReferencesEnabled && contentReferencesStoreFactory(); + contentReferencesEnabled ? newContentReferencesStore() : undefined; onLlmResponse = async ( content: string, diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts index 946447e2815e0..fffce77b9251e 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -13,7 +13,7 @@ import { ALERT_COUNTS_TOOL } from './alert_counts_tool'; import type { RetrievalQAChain } from 'langchain/chains'; import type { ExecuteConnectorRequestBody } from '@kbn/elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen'; import type { ContentReferencesStore } 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('AlertCountsTool', () => { 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, 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..da58b8c045da4 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 = { 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..67ba84cee52f2 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, 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..f797a306ddc66 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, 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..21c3af0ccfd4f 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 = { From 9d8c84fec3bbcf012ef8313ebfb1b939ec927df0 Mon Sep 17 00:00:00 2001 From: Kenneth Kreindler Date: Thu, 30 Jan 2025 11:41:39 +0000 Subject: [PATCH 2/5] lint --- .../plugins/elastic_assistant/server/lib/prompt/prompts.ts | 2 +- .../server/routes/chat/chat_complete_route.ts | 5 +++-- .../server/routes/evaluate/post_evaluate.ts | 5 +++-- .../server/routes/post_actions_connector_execute.ts | 5 +++-- .../security/plugins/elastic_assistant/server/types.ts | 2 +- 5 files changed, 11 insertions(+), 8 deletions(-) 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 a9a2f166399e8..c356b3508c3b2 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. Always return value from NaturalLanguageESQLTool as is. 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.`; 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 a1f1763ca8aee..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 @@ -186,8 +186,9 @@ export const chatCompleteRoute = ( })); } - const contentReferencesStore = - contentReferencesEnabled ? newContentReferencesStore() : undefined; + const contentReferencesStore = contentReferencesEnabled + ? newContentReferencesStore() + : undefined; const onLlmResponse = async ( content: string, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts index d8112c5ba218e..6f6897db0c6d6 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/evaluate/post_evaluate.ts @@ -292,8 +292,9 @@ export const postEvaluateRoute = ( assistantContext.getRegisteredFeatures( DEFAULT_PLUGIN_NAME ).contentReferencesEnabled; - const contentReferencesStore = - contentReferencesEnabled ? newContentReferencesStore() : undefined; + const contentReferencesStore = contentReferencesEnabled + ? newContentReferencesStore() + : undefined; // Fetch any applicable tools that the source plugin may have registered const assistantToolParams: AssistantToolParams = { diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts index ba217fa7be647..685863cfa97b1 100644 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/routes/post_actions_connector_execute.ts @@ -119,8 +119,9 @@ export const postActionsConnectorExecuteRoute = ( }); const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient(); - const contentReferencesStore = - contentReferencesEnabled ? newContentReferencesStore() : undefined; + const contentReferencesStore = contentReferencesEnabled + ? newContentReferencesStore() + : undefined; onLlmResponse = async ( content: string, diff --git a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts index 981b00d68e643..3e5d3eb383df3 100755 --- a/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts +++ b/x-pack/solutions/security/plugins/elastic_assistant/server/types.ts @@ -239,7 +239,7 @@ export interface AssistantToolParams { inference?: InferenceServerStart; isEnabledKnowledgeBase: boolean; connectorId?: string; - contentReferencesStore: ContentReferencesStore | false; + contentReferencesStore: ContentReferencesStore | undefined; esClient: ElasticsearchClient; kbDataClient?: AIAssistantKnowledgeBaseDataClient; langChainTimeout?: number; From 39d8027a515122fc91dc7845f878a985f5e4793e Mon Sep 17 00:00:00 2001 From: Kenneth Kreindler Date: Thu, 30 Jan 2025 12:08:49 +0000 Subject: [PATCH 3/5] type --- .../server/routes/defend_insights/post_defend_insights.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 Date: Thu, 30 Jan 2025 13:09:42 +0000 Subject: [PATCH 4/5] type --- .../tools/knowledge_base/knowledge_base_retrieval_tool.test.ts | 2 +- .../open_and_acknowledged_alerts_tool.test.ts | 2 +- .../tools/product_docs/product_documentation_tool.test.ts | 2 +- .../assistant/tools/security_labs/security_labs_tool.test.ts | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) 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 da58b8c045da4..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 @@ -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 67ba84cee52f2..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 @@ -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 f797a306ddc66..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 @@ -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 21c3af0ccfd4f..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 @@ -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' }); From b27748cc16fce48b4229917f205087d031e382ce Mon Sep 17 00:00:00 2001 From: Kenneth Kreindler Date: Thu, 30 Jan 2025 14:01:33 +0000 Subject: [PATCH 5/5] type --- .../assistant/tools/alert_counts/alert_counts_tool.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts index fffce77b9251e..61ddceab607cf 100644 --- a/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts +++ b/x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.test.ts @@ -191,7 +191,7 @@ describe('AlertCountsTool', () => { replacements, request, ...rest, - contentReferencesStore: false, + contentReferencesStore: undefined, }) as DynamicTool; const result = await tool.func('');