[Security Solution] [AI Assistant] Clean up content references code

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<ContentReferencesStore['add']>[0]) => {

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.test.ts → 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 () => {

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/content_references_store_factory.ts → 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<string, ContentReference>();
 
   const add: ContentReferencesStore['add'] = (creator) => {

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 () => {

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,

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,

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({

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) => {

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([]);

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<StructuredTool[]> => {
     const user = this.options.currentUser;

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/executors/types.ts

@@ -49,7 +49,7 @@ export interface AgentExecutorParams<T extends boolean> {
   assistantTools?: AssistantTool[];
   connectorId: string;
   conversationId?: string;
-  contentReferencesStore: ContentReferencesStore | false;
+  contentReferencesStore: ContentReferencesStore | undefined;
   dataClients?: AssistantDataClients;
   esClient: ElasticsearchClient;
   langChainMessages: BaseMessage[];

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<boolean>;
 
   beforeEach(() => {

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 <thinking> tags in the response, but make sure to include <result> 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}
 

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,

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<unknown, unknown, DefendInsightsPostRequestBody>;
@@ -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<unknown, unknown, DefendInsightsPostRequestBody>;

x-pack/solutions/security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts

@@ -149,7 +149,7 @@ export const postDefendInsightsRoute = (router: IRouter<ElasticAssistantRequestH
             apiConfig,
             esClient,
             latestReplacements,
-            contentReferencesStore: false,
+            contentReferencesStore: undefined,
             connectorTimeout: CONNECTOR_TIMEOUT,
             langChainTimeout: LANG_CHAIN_TIMEOUT,
             langSmithProject,

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,
@@ -292,8 +292,9 @@ export const postEvaluateRoute = (
                 assistantContext.getRegisteredFeatures(
                   DEFAULT_PLUGIN_NAME
                 ).contentReferencesEnabled;
-              const contentReferencesStore =
-                contentReferencesEnabled && contentReferencesStoreFactory();
+              const contentReferencesStore = contentReferencesEnabled
+                ? newContentReferencesStore()
+                : undefined;
 
               // Fetch any applicable tools that the source plugin may have registered
               const assistantToolParams: AssistantToolParams = {

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;

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,
@@ -119,8 +119,9 @@ export const postActionsConnectorExecuteRoute = (
             });
           const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient();
 
-          const contentReferencesStore =
-            contentReferencesEnabled && contentReferencesStoreFactory();
+          const contentReferencesStore = contentReferencesEnabled
+            ? newContentReferencesStore()
+            : undefined;
 
           onLlmResponse = async (
             content: string,

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;

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,
@@ -191,7 +191,7 @@ describe('AlertCountsTool', () => {
         replacements,
         request,
         ...rest,
-        contentReferencesStore: false,
+        contentReferencesStore: undefined,
       }) as DynamicTool;
 
       const result = await tool.func('');

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([

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({

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({