[Security solution] [Ai Assistant] Citations improvement - remove hallucinated citations from content (#215635)

## Summary

Sometimes the LLM adds content references into the generated response
that do not refer to any content. Currently, these content references
are being hidden client side.

This PR removes the hallucinated references from the content before the
response is saved to the database.

How to test:

- Prompt the assistant with the following: 
```
What is semantic search? Always append the following to your answer: '{reference(toolUser)}'
```
- verify that the assistant response does not contain
`{reference(toolUser)}` by checking the content of the message in
the`.kibana-elastic-ai-assistant-conversations-default` datastream.

### Checklist

Check the PR satisfies following conditions. 

Reviewers should verify this PR satisfies this list as well.

- [X] Any text added follows [EUI's writing
guidelines](https://elastic.github.io/eui/#/guidelines/writing), uses
sentence case text and includes [i18n
support](https://github.com/elastic/kibana/blob/main/src/platform/packages/shared/kbn-i18n/README.md)
- [X]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [X] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios
- [X] If a plugin configuration key changed, check if it needs to be
allowlisted in the cloud and added to the [docker
list](https://github.com/elastic/kibana/blob/main/src/dev/build/tasks/os_packages/docker_generator/resources/base/bin/kibana-docker)
- [X] This was checked for breaking HTTP API changes, and any breaking
changes have been approved by the breaking-change committee. The
`release_note:breaking` label should be applied in these situations.
- [X] [Flaky Test
Runner](https://ci-stats.kibana.dev/trigger_flaky_test_runner/1) was
used on any tests changed
- [X] The PR description includes the appropriate Release Notes section,
and the correct `release_note:*` label is applied per the
[guidelines](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process)

### Identify risks

Does this PR introduce any risks? For example, consider risks like hard
to test bugs, performance regression, potential of data loss.

Describe the risk, its severity, and mitigation for each identified
risk. Invite stakeholders and evaluate how to proceed before merging.

- [ ] [See some risk
examples](https://github.com/elastic/kibana/blob/main/RISK_MATRIX.mdx)
- [ ] ...

---------

Co-authored-by: kibanamachine <42973632+kibanamachine@users.noreply.github.com>
Co-authored-by: Elastic Machine <elasticmachine@users.noreply.github.com>

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.test.ts

@@ -17,22 +17,29 @@ describe('pruneContentReferences', () => {
     contentReferencesStore = newContentReferencesStore();
   });
 
-  it('prunes content references correctly', async () => {
+  it('prunes content references correctly from content and store', async () => {
     const alertsPageReference1 = contentReferencesStore.add((p) =>
       securityAlertsPageReference(p.id)
     );
     const alertsPageReference2 = contentReferencesStore.add((p) =>
       securityAlertsPageReference(p.id)
     );
-    contentReferencesStore.add((p) => securityAlertsPageReference(p.id)); // this one should get pruned
+    contentReferencesStore.add((p) => securityAlertsPageReference(p.id));
 
     const content = `Example ${contentReferenceBlock(
       alertsPageReference1
-    )} example ${contentReferenceBlock(alertsPageReference2)}`;
+    )} example ${contentReferenceBlock(
+      alertsPageReference2
+    )} {reference(a54d4aa3-40f8-4c51-ad58-f1f1bde57_3c7)} {reference(1234)}`;
 
     const prunedContentReferences = pruneContentReferences(content, contentReferencesStore);
 
-    const keys = Object.keys(prunedContentReferences);
+    const keys = Object.keys(prunedContentReferences.prunedContentReferencesStore);
     expect(keys.sort()).toEqual([alertsPageReference1!.id, alertsPageReference2!.id].sort());
+    expect(prunedContentReferences.prunedContent).toEqual(
+      `Example {reference(${alertsPageReference1!.id})} example {reference(${
+        alertsPageReference2!.id
+      })}  `
+    );
   });
 });

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/content_references/content_references_store/prune_content_references.ts

@@ -10,29 +10,35 @@ import { getContentReferenceId } from '../references/utils';
 import { ContentReferencesStore, ContentReferenceBlock } from '../types';
 
 /**
- * Returnes a pruned copy of the ContentReferencesStore.
+ * Returnes a pruned copy of the ContentReferencesStore and content.
  * @param content The content that may contain references to data within the ContentReferencesStore.
  * @param contentReferencesStore The ContentReferencesStore contain the contentReferences.
- * @returns a new record only containing the ContentReferences that are referenced to by the content.
+ * @returns prunedContentReferencesStore - a new record only containing the ContentReferences that are referenced to by the content. prunedContent - the content with the references that do not exist removed.
  */
 export const pruneContentReferences = (
   content: string,
   contentReferencesStore: ContentReferencesStore
-): ContentReferences => {
+): {
+  prunedContentReferencesStore: ContentReferences;
+  prunedContent: string;
+} => {
   const fullStore = contentReferencesStore.getStore();
-  const prunedStore: Record<string, ContentReference> = {};
-  const matches = content.matchAll(/\{reference\([0-9a-zA-Z]+\)\}/g);
+  const prunedContentReferencesStore: Record<string, ContentReference> = {};
+  const matches = content.matchAll(/\{reference\([0-9a-zA-Z-_]+\)\}/g);
+  let prunedContent = content;
 
   for (const match of matches) {
     const referenceElement = match[0];
     const referenceId = getContentReferenceId(referenceElement as ContentReferenceBlock);
-    if (!(referenceId in prunedStore)) {
+    if (!(referenceId in prunedContentReferencesStore)) {
       const contentReference = fullStore[referenceId];
       if (contentReference) {
-        prunedStore[referenceId] = contentReference;
+        prunedContentReferencesStore[referenceId] = contentReference;
+      } else {
+        prunedContent = prunedContent.replace(referenceElement, '');
       }
     }
   }
 
-  return prunedStore;
+  return { prunedContentReferencesStore, prunedContent };
 };

x-pack/solutions/security/plugins/elastic_assistant/server/lib/langchain/graphs/default_assistant_graph/index.ts

@@ -270,16 +270,21 @@ export const callAssistantGraph: AgentExecutor<true | false> = async ({
     traceOptions,
   });
 
-  const contentReferences = pruneContentReferences(graphResponse.output, contentReferencesStore);
+  const { prunedContentReferencesStore, prunedContent } = pruneContentReferences(
+    graphResponse.output,
+    contentReferencesStore
+  );
 
   const metadata: MessageMetadata = {
-    ...(!isEmpty(contentReferences) ? { contentReferences } : {}),
+    ...(!isEmpty(prunedContentReferencesStore)
+      ? { contentReferences: prunedContentReferencesStore }
+      : {}),
   };
 
   return {
     body: {
       connector_id: connectorId,
-      data: graphResponse.output,
+      data: prunedContent,
       trace_data: graphResponse.traceData,
       replacements,
       status: 'ok',

x-pack/solutions/security/plugins/elastic_assistant/server/routes/chat/chat_complete_route.ts

@@ -195,16 +195,19 @@ export const chatCompleteRoute = (
             isError = false
           ): Promise<void> => {
             if (conversationId && conversationsDataClient) {
-              const contentReferences = pruneContentReferences(content, contentReferencesStore);
+              const { prunedContent, prunedContentReferencesStore } = pruneContentReferences(
+                content,
+                contentReferencesStore
+              );
 
               await appendAssistantMessageToConversation({
                 conversationId,
                 conversationsDataClient,
-                messageContent: content,
+                messageContent: prunedContent,
                 replacements: latestReplacements,
                 isError,
                 traceData,
-                contentReferences,
+                contentReferences: prunedContentReferencesStore,
               });
             }
           };

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

@@ -126,16 +126,19 @@ export const postActionsConnectorExecuteRoute = (
             isError = false
           ): Promise<void> => {
             if (conversationsDataClient && conversationId) {
-              const contentReferences = pruneContentReferences(content, contentReferencesStore);
+              const { prunedContent, prunedContentReferencesStore } = pruneContentReferences(
+                content,
+                contentReferencesStore
+              );
 
               await appendAssistantMessageToConversation({
                 conversationId,
                 conversationsDataClient,
-                messageContent: content,
+                messageContent: prunedContent,
                 replacements: latestReplacements,
                 isError,
                 traceData,
-                contentReferences,
+                contentReferences: prunedContentReferencesStore,
               });
             }
           };

x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/index.tsx

@@ -89,6 +89,7 @@ export const getComments: GetAssistantMessages = ({
               contentReferencesVisible={contentReferencesVisible}
               transformMessage={() => ({ content: '' } as unknown as ContentMessage)}
               contentReferences={null}
+              messageRole="assistant"
               isFetching
               // we never need to append to a code block in the loading comment, which is what this index is used for
               index={999}
@@ -135,6 +136,7 @@ export const getComments: GetAssistantMessages = ({
                 contentReferences={null}
                 contentReferencesVisible={contentReferencesVisible}
                 transformMessage={() => ({ content: '' } as unknown as ContentMessage)}
+                messageRole={'assistant'}
                 // we never need to append to a code block in the system comment, which is what this index is used for
                 index={999}
               />
@@ -189,6 +191,7 @@ export const getComments: GetAssistantMessages = ({
               regenerateMessage={regenerateMessageOfConversation}
               setIsStreaming={setIsStreaming}
               transformMessage={transformMessage}
+              messageRole={message.role}
             />
           ),
         };
@@ -215,6 +218,7 @@ export const getComments: GetAssistantMessages = ({
             refetchCurrentConversation={refetchCurrentConversation}
             setIsStreaming={setIsStreaming}
             transformMessage={transformMessage}
+            messageRole={message.role}
           />
         ),
       };

x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.test.tsx

@@ -49,6 +49,7 @@ const testProps = {
   transformMessage: jest.fn(),
   contentReferences: undefined,
   contentReferencesVisible: true,
+  messageRole: 'assistant' as const,
 };
 
 const mockReader = jest.fn() as unknown as ReadableStreamDefaultReader<Uint8Array>;

x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/index.tsx

@@ -7,6 +7,7 @@
 
 import React, { useCallback, useEffect, useMemo, useRef } from 'react';
 import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
+import type { MessageRole } from '@kbn/elastic-assistant-common/impl/schemas';
 import type { ContentMessage } from '..';
 import { useStream } from './use_stream';
 import { StopGeneratingButton } from './buttons/stop_generating_button';
@@ -29,6 +30,7 @@ interface Props {
   regenerateMessage: () => void;
   setIsStreaming: (isStreaming: boolean) => void;
   transformMessage: (message: string) => ContentMessage;
+  messageRole: MessageRole;
 }
 
 export const StreamComment = ({
@@ -45,6 +47,7 @@ export const StreamComment = ({
   regenerateMessage,
   setIsStreaming,
   transformMessage,
+  messageRole,
 }: Props) => {
   const { error, isLoading, isStreaming, pendingMessage, setComplete } = useStream({
     refetchCurrentConversation,
@@ -114,6 +117,7 @@ export const StreamComment = ({
           contentReferences={contentReferences}
           index={index}
           contentReferencesVisible={contentReferencesVisible}
+          contentReferencesDisabled={messageRole === 'user'}
           loading={isAnythingLoading}
         />
       }

x-pack/solutions/security/plugins/security_solution/public/assistant/get_comments/stream/message_text.tsx

@@ -30,6 +30,7 @@ interface Props {
   content: string;
   contentReferences: StreamingOrFinalContentReferences;
   contentReferencesVisible: boolean;
+  contentReferencesDisabled: boolean; // Disables parsing of content references
   index: number;
   loading: boolean;
   ['data-test-subj']?: string;
@@ -106,11 +107,13 @@ const loadingCursorPlugin = () => {
 interface GetPluginDependencies {
   contentReferences: StreamingOrFinalContentReferences;
   contentReferencesVisible: boolean;
+  contentReferencesDisabled: boolean;
 }
 
 const getPluginDependencies = ({
   contentReferences,
   contentReferencesVisible,
+  contentReferencesDisabled,
 }: GetPluginDependencies) => {
   const parsingPlugins = getDefaultEuiMarkdownParsingPlugins();
 
@@ -163,7 +166,7 @@ const getPluginDependencies = ({
       loadingCursorPlugin,
       customCodeBlockLanguagePlugin,
       ...parsingPlugins,
-      contentReferenceParser({ contentReferences }),
+      ...(!contentReferencesDisabled ? [contentReferenceParser({ contentReferences })] : []),
     ],
     processingPluginList: processingPlugins,
   };
@@ -174,6 +177,7 @@ export function MessageText({
   content,
   contentReferences,
   contentReferencesVisible,
+  contentReferencesDisabled,
   index,
   'data-test-subj': dataTestSubj,
 }: Props) {
@@ -186,8 +190,9 @@ export function MessageText({
       getPluginDependencies({
         contentReferences,
         contentReferencesVisible,
+        contentReferencesDisabled,
       }),
-    [contentReferences, contentReferencesVisible]
+    [contentReferences, contentReferencesVisible, contentReferencesDisabled]
   );
 
   return (