remove content references from post defend endpoint

elastic · Feb 25, 2025 · 5be647b · 5be647b
1 parent e2fe023
commit 5be647b
Show file tree

Hide file tree

Showing 9 changed files with 28 additions and 20 deletions.
diff --git a/.../packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts b/.../packages/shared/kbn-elastic-assistant-common/impl/content_references/references/utils.ts
@@ -26,8 +26,11 @@ export const getContentReferenceId = (
  * @returns ContentReferenceBlock
  */
 export const contentReferenceBlock = (
-  contentReference: ContentReference
-): ContentReferenceBlock => {
+  contentReference: ContentReference | undefined
+): ContentReferenceBlock | "" => {
+  if (!contentReference) {
+    return '';
+  }
   return `{reference(${contentReference.id})}`;
 };
 
@@ -36,7 +39,10 @@ export const contentReferenceBlock = (
  * @param contentReference A ContentReference
  * @returns the string: `Reference: <contentReferenceBlock>`
  */
-export const contentReferenceString = (contentReference: ContentReference) => {
+export const contentReferenceString = (contentReference: ContentReference | undefined) => {
+  if (!contentReference) {
+    return '';
+  }
   return `Citation: ${contentReferenceBlock(contentReference)}` as const;
 };
 

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
@@ -124,7 +124,7 @@ export function getAssistantToolParams({
   langSmithProject?: string;
   langSmithApiKey?: string;
   logger: Logger;
-  contentReferencesStore: ContentReferencesStore;
+  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;
+  contentReferencesStore: ContentReferencesStore | undefined;
   replacements: Replacements;
   onNewReplacements: (newReplacements: Replacements) => void;
   request: KibanaRequest<unknown, unknown, DefendInsightsPostRequestBody>;

diff --git a/.../security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts b/.../security/plugins/elastic_assistant/server/routes/defend_insights/post_defend_insights.ts
@@ -153,7 +153,7 @@ export const postDefendInsightsRoute = (router: IRouter<ElasticAssistantRequestH
             apiConfig,
             esClient,
             latestReplacements,
-            contentReferencesStore: newContentReferencesStore(),
+            contentReferencesStore: undefined,
             connectorTimeout: CONNECTOR_TIMEOUT,
             langChainTimeout: LANG_CHAIN_TIMEOUT,
             langSmithProject,

diff --git 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;
+  contentReferencesStore: ContentReferencesStore | undefined;
   description?: string;
   esClient: ElasticsearchClient;
   kbDataClient?: AIAssistantKnowledgeBaseDataClient;

diff --git a/...curity/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts b/...curity/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts
@@ -43,7 +43,7 @@ export const ALERT_COUNTS_TOOL: AssistantTool = {
       func: async () => {
         const query = getAlertsCountQuery(alertsIndexPattern);
         const result = await esClient.search<SearchResponse>(query);
-        const alertsCountReference = contentReferencesStore.add((p) =>
+        const alertsCountReference = contentReferencesStore?.add((p) =>
           securityAlertsPageReference(p.id)
         );
 

diff --git a/.../security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts b/.../security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts
@@ -66,9 +66,9 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = {
   },
 };
 
-function enrichDocument(contentReferencesStore: ContentReferencesStore) {
+function enrichDocument(contentReferencesStore: ContentReferencesStore | undefined) {
   return (document: Document<Record<string, string>>) => {
-    if (document.id == null) {
+    if (document.id == null || contentReferencesStore == null) {
       return document;
     }
     const documentId = document.id;

diff --git a/.../server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts b/.../server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts
@@ -86,21 +86,20 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = {
         };
 
         return JSON.stringify(
-          result.hits?.hits?.map((x) => {
+          result.hits?.hits?.map((hit) => {
             const transformed = transformRawData({
               anonymizationFields,
               currentReplacements: localReplacements, // <-- the latest local replacements
               getAnonymizedValue,
               onNewReplacements: localOnNewReplacements, // <-- the local callback
-              rawData: getRawDataOrDefault(x.fields),
+              rawData: getRawDataOrDefault(hit.fields),
             });
 
-            const hitId = x._id;
+            const hitId = hit._id;
+            const reference = hitId ? contentReferencesStore?.add((p) => securityAlertReference(p.id, hitId)) : undefined
             const citation =
-              hitId &&
-              `\nCitation,${contentReferenceBlock(
-                contentReferencesStore.add((p) => securityAlertReference(p.id, hitId))
-              )}`;
+            reference &&
+              `\nCitation,${contentReferenceBlock(reference)}`;
 
             return `${transformed}${citation ?? ''}`;
           })

diff --git a/...ugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts b/...ugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts
@@ -90,11 +90,14 @@ export const PRODUCT_DOCUMENTATION_TOOL: AssistantTool = {
 };
 
 type EnrichedDocument = RetrieveDocumentationResultDoc & {
-  citation: string;
+  citation?: string;
 };
 
-const enrichDocument = (contentReferencesStore: ContentReferencesStore) => {
+const enrichDocument = (contentReferencesStore: ContentReferencesStore | undefined) => {
   return (document: RetrieveDocumentationResultDoc): EnrichedDocument => {
+    if (contentReferencesStore == null) {
+      return document
+    }
     const reference = contentReferencesStore.add((p) =>
       productDocumentationReference(p.id, document.title, document.url)
     );

diff --git a/...rity/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts b/...rity/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts
@@ -51,7 +51,7 @@ export const SECURITY_LABS_KNOWLEDGE_BASE_TOOL: AssistantTool = {
           query: input.question,
         });
 
-        const reference = contentReferencesStore.add((p) =>
+        const reference = contentReferencesStore?.add((p) =>
           knowledgeBaseReference(p.id, 'Elastic Security Labs content', 'securityLabsId')
         );