[Security solution] [Ai Assistant] Allow user to disabled content references by using the content_references_disabled query arg (#215818)

## Summary

Some users using the security AI assistant through the API may want to
disable content references (aka citations) programmatically. This PR
allows users to disable content references using the
`content_references_disabled` query arg in the following routes:
- POST
/internal/elastic_assistant/actions/connector/gpt4oAzure/_execute?content_references_disabled=true
- POST
/api/security_ai_assistant/chat/complete?content_references_disabled=true


## How to test:
Bellow are 2 example curl requests (with
content_references_disabled=true). When executed, you will notice that
the response stream does not contain the string `{reference(...)}`
anywhere. If you remove the query arg, the reference will be visible.

####
/internal/elastic_assistant/actions/connector/gpt4oAzure/_execute?content_references_disabled=true
```curl
curl --location 'http://localhost:5601/internal/elastic_assistant/actions/connector/gpt4oAzure/_execute?content_references_disabled=true' \
--header 'Accept: */*' \
--header 'Accept-Language: en-GB,en-US;q=0.9,en;q=0.8' \
--header 'Connection: keep-alive' \
--header 'Content-Type: application/json' \
--header 'Origin: http://localhost:5601' \
--header 'Referer: http://localhost:5601/app/security/dashboards?sourcerer=(default:(id:security-solution-default,selectedPatterns:!()))&timeline=(activeTab:query,graphEventId:%27%27,isOpen:!f)' \
--header 'Sec-Fetch-Dest: empty' \
--header 'Sec-Fetch-Mode: cors' \
--header 'Sec-Fetch-Site: same-origin' \
--header 'User-Agent: Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36' \
--header 'elastic-api-version: 1' \
--header 'kbn-build-number: 9007199254740991' \
--header 'kbn-version: 9.1.0' \
--header 'sec-ch-ua: "Chromium";v="134", "Not:A-Brand";v="24", "Google Chrome";v="134"' \
--header 'sec-ch-ua-mobile: ?0' \
--header 'sec-ch-ua-platform: "macOS"' \
--header 'x-elastic-internal-origin: Kibana' \
--header 'x-kbn-context: %7B%22type%22%3A%22application%22%2C%22name%22%3A%22securitySolutionUI%22%2C%22url%22%3A%22%2Fapp%2Fsecurity%2Fdashboards%22%2C%22page%22%3A%22list%22%7D' \
--header 'Authorization: Basic ZWxhc3RpYzpjaGFuZ2VtZQ==' \
--data '{
    "subAction": "invokeStream",
    "actionTypeId": ".gen-ai",
    "replacements": {},
    "message": "What is elastic security? Include citations",
    "screenContext": {
        "timeZone": "Asia/Calcutta"
    },
    "alertsIndexPattern": ".alerts-security.alerts-default",
    "size": 100
}'
```

####
/api/security_ai_assistant/chat/complete?content_references_disabled=true

```curl
curl --location 'http://localhost:5601/api/security_ai_assistant/chat/complete?content_references_disabled=true' \
--header 'Content-Type: application/json' \
--header 'Accept: application/octet-stream' \
--header 'kbn-xsrf: 123' \
--header 'Authorization: Basic ZWxhc3RpYzpjaGFuZ2VtZQ==' \
--data '{
  "messages": [
    {
      "role": "user",
      "content": "what is semantic text in esql. Include references to the product documentation"
    }
  ],
  "persist": false,
  "connectorId": "gpt4oAzure",
  "isStream": false
}'
```




### 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>

oas_docs/output/kibana.serverless.yaml

@@ -41366,6 +41366,14 @@ paths:
     post:
       description: Create a model response for the given chat conversation.
       operationId: ChatComplete
+      parameters:
+        - description: If true, the response will not include content references.
+          in: query
+          name: content_references_disabled
+          required: false
+          schema:
+            default: false
+            type: boolean
       requestBody:
         content:
           application/json:

oas_docs/output/kibana.yaml

@@ -44329,6 +44329,14 @@ paths:
     post:
       description: Create a model response for the given chat conversation.
       operationId: ChatComplete
+      parameters:
+        - description: If true, the response will not include content references.
+          in: query
+          name: content_references_disabled
+          required: false
+          schema:
+            default: false
+            type: boolean
       requestBody:
         content:
           application/json:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/ess/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -158,6 +158,14 @@ paths:
     post:
       description: Create a model response for the given chat conversation.
       operationId: ChatComplete
+      parameters:
+        - description: If true, the response will not include content references.
+          in: query
+          name: content_references_disabled
+          required: false
+          schema:
+            default: false
+            type: boolean
       requestBody:
         content:
           application/json:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/docs/openapi/serverless/elastic_assistant_api_2023_10_31.bundled.schema.yaml

@@ -158,6 +158,14 @@ paths:
     post:
       description: Create a model response for the given chat conversation.
       operationId: ChatComplete
+      parameters:
+        - description: If true, the response will not include content references.
+          in: query
+          name: content_references_disabled
+          required: false
+          schema:
+            default: false
+            type: boolean
       requestBody:
         content:
           application/json:

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

@@ -7,15 +7,11 @@
 
 import { newContentReferencesStore } from './content_references_store';
 import { securityAlertsPageReference } from '../references';
-import { ContentReferencesStore } from '../types';
 
 describe('newContentReferencesStore', () => {
-  let contentReferencesStore: ContentReferencesStore;
-  beforeEach(() => {
-    contentReferencesStore = newContentReferencesStore();
-  });
-
   it('adds multiple content reference', async () => {
+    const contentReferencesStore = newContentReferencesStore();
+
     const alertsPageReference1 = contentReferencesStore.add((p) =>
       securityAlertsPageReference(p.id)
     );
@@ -31,18 +27,35 @@ describe('newContentReferencesStore', () => {
     const keys = Object.keys(store);
 
     expect(keys.length).toEqual(3);
-    expect(store[alertsPageReference1.id]).toEqual(alertsPageReference1);
-    expect(store[alertsPageReference2.id]).toEqual(alertsPageReference2);
-    expect(store[alertsPageReference3.id]).toEqual(alertsPageReference3);
+    expect(store[alertsPageReference1!.id]).toEqual(alertsPageReference1);
+    expect(store[alertsPageReference2!.id]).toEqual(alertsPageReference2);
+    expect(store[alertsPageReference3!.id]).toEqual(alertsPageReference3);
+  });
+
+  it('disabled content reference store', async () => {
+    const contentReferencesStore = newContentReferencesStore({
+      disabled: true,
+    });
+    const alertsPageReference1 = contentReferencesStore.add((p) =>
+      securityAlertsPageReference(p.id)
+    );
+
+    const store = contentReferencesStore.getStore();
+
+    const keys = Object.keys(store);
+
+    expect(keys.length).toEqual(0);
+    expect(alertsPageReference1).toBeUndefined();
   });
 
   it('referenceIds are unique', async () => {
+    const contentReferencesStore = newContentReferencesStore();
     const numberOfReferencesToCreate = 50;
 
     const referenceIds = new Set(
       [...new Array(numberOfReferencesToCreate)]
         .map(() => contentReferencesStore.add((p) => securityAlertsPageReference(p.id)))
-        .map((alertsPageReference) => alertsPageReference.id)
+        .map((alertsPageReference) => alertsPageReference!.id)
     );
 
     const store = contentReferencesStore.getStore();

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

@@ -11,13 +11,22 @@ import { ContentReferencesStore } from '../types';
 const CONTENT_REFERENCE_ID_ALPHABET =
   'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789';
 
+interface Options {
+  disabled?: boolean;
+}
+
 /**
  * Creates a new ContentReferencesStore used for storing references (also known as citations)
  */
-export const newContentReferencesStore: () => ContentReferencesStore = () => {
+export const newContentReferencesStore: (options?: Options) => ContentReferencesStore = (
+  options?: Options
+) => {
   const store = new Map<string, ContentReference>();
 
   const add: ContentReferencesStore['add'] = (creator) => {
+    if (options?.disabled) {
+      return undefined;
+    }
     const entry = creator({ id: generateUnsecureId() });
     store.set(entry.id, entry);
     return entry;

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

@@ -32,7 +32,7 @@ describe('pruneContentReferences', () => {
 
     const prunedContentReferences = pruneContentReferences(content, contentReferencesStore);
 
-    const keys = Object.keys(prunedContentReferences!);
-    expect(keys.sort()).toEqual([alertsPageReference1.id, alertsPageReference2.id].sort());
+    const keys = Object.keys(prunedContentReferences);
+    expect(keys.sort()).toEqual([alertsPageReference1!.id, alertsPageReference2!.id].sort());
   });
 });

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

@@ -19,7 +19,9 @@ export interface ContentReferencesStore {
    * @param generator.params.id An ID that is guaranteed to not exist in the store. Intended to be used as the Id of the ContentReference but not required.
    * @returns the new ContentReference
    */
-  add: <T extends ContentReference>(generator: (params: { id: ContentReferenceId }) => T) => T;
+  add: <T extends ContentReference>(
+    generator: (params: { id: ContentReferenceId }) => T
+  ) => T | undefined;
 
   /**
    * Used to read the content reference store.

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.gen.ts

@@ -15,10 +15,20 @@
  */
 
 import { z } from '@kbn/zod';
+import { BooleanFromString } from '@kbn/zod-helpers';
 
 import { NonEmptyString, ScreenContext } from '../common_attributes.gen';
 import { Replacements } from '../conversations/common_attributes.gen';
 
+export type ExecuteConnectorRequestQuery = z.infer<typeof ExecuteConnectorRequestQuery>;
+export const ExecuteConnectorRequestQuery = z.object({
+  /**
+   * If true, the response will not include content references.
+   */
+  content_references_disabled: BooleanFromString.optional().default(false),
+});
+export type ExecuteConnectorRequestQueryInput = z.input<typeof ExecuteConnectorRequestQuery>;
+
 export type ExecuteConnectorRequestParams = z.infer<typeof ExecuteConnectorRequestParams>;
 export const ExecuteConnectorRequestParams = z.object({
   /**

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/actions_connector/post_actions_connector_execute_route.schema.yaml

@@ -19,6 +19,13 @@ paths:
           description: The connector's `id` value.
           schema:
             type: string
+        - name: content_references_disabled
+          in: query
+          required: false
+          description: If true, the response will not include content references.
+          schema:
+            type: boolean
+            default: false
       requestBody:
         required: true
         content:

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/chat/post_chat_complete_route.gen.ts

@@ -15,6 +15,7 @@
  */
 
 import { z } from '@kbn/zod';
+import { BooleanFromString } from '@kbn/zod-helpers';
 
 import { NonEmptyString } from '../common_attributes.gen';
 
@@ -66,6 +67,15 @@ export const ChatCompleteProps = z.object({
   messages: z.array(ChatMessage),
 });
 
+export type ChatCompleteRequestQuery = z.infer<typeof ChatCompleteRequestQuery>;
+export const ChatCompleteRequestQuery = z.object({
+  /**
+   * If true, the response will not include content references.
+   */
+  content_references_disabled: BooleanFromString.optional().default(false),
+});
+export type ChatCompleteRequestQueryInput = z.input<typeof ChatCompleteRequestQuery>;
+
 export type ChatCompleteRequestBody = z.infer<typeof ChatCompleteRequestBody>;
 export const ChatCompleteRequestBody = ChatCompleteProps;
 export type ChatCompleteRequestBodyInput = z.input<typeof ChatCompleteRequestBody>;

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/chat/post_chat_complete_route.schema.yaml

@@ -12,6 +12,14 @@ paths:
       summary: Create a model response
       tags:
         - Chat Complete API
+      parameters:
+        - name: content_references_disabled
+          in: query
+          required: false
+          description: If true, the response will not include content references.
+          schema:
+            type: boolean
+            default: false
       requestBody:
         required: true
         content:

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

@@ -137,6 +137,7 @@ const mockRequest = {
   events: {
     aborted$: NEVER,
   },
+  query: {},
 };
 
 const mockResponse = {

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

@@ -18,6 +18,7 @@ import {
   ConversationResponse,
   newContentReferencesStore,
   pruneContentReferences,
+  ChatCompleteRequestQuery,
 } from '@kbn/elastic-assistant-common';
 import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
 import { getRequestAbortedSignal } from '@kbn/data-plugin/server';
@@ -60,12 +61,14 @@ export const chatCompleteRoute = (
         validate: {
           request: {
             body: buildRouteValidationWithZod(ChatCompleteProps),
+            query: buildRouteValidationWithZod(ChatCompleteRequestQuery),
           },
         },
       },
       async (context, request, response) => {
         const abortSignal = getRequestAbortedSignal(request.events.aborted$);
         const assistantResponse = buildResponse(response);
+        const { content_references_disabled: contentReferencesDisabled } = request.query;
         let telemetry;
         let actionTypeId;
         const ctx = await context.resolve(['core', 'elasticAssistant', 'licensing']);
@@ -182,7 +185,9 @@ export const chatCompleteRoute = (
             ? existingConversationId ?? newConversation?.id
             : undefined;
 
-          const contentReferencesStore = newContentReferencesStore();
+          const contentReferencesStore = newContentReferencesStore({
+            disabled: contentReferencesDisabled ?? false,
+          });
 
           const onLlmResponse = async (
             content: string,

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

@@ -117,6 +117,7 @@ const mockRequest = {
   events: {
     aborted$: NEVER,
   },
+  query: {},
 };
 
 const mockResponse = {

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

@@ -17,6 +17,7 @@ import {
   Message,
   Replacements,
   pruneContentReferences,
+  ExecuteConnectorRequestQuery,
 } from '@kbn/elastic-assistant-common';
 import { buildRouteValidationWithZod } from '@kbn/elastic-assistant-common/impl/schemas/common';
 import { INVOKE_ASSISTANT_ERROR_EVENT } from '../lib/telemetry/event_based_telemetry';
@@ -55,6 +56,7 @@ export const postActionsConnectorExecuteRoute = (
             params: schema.object({
               connectorId: schema.string(),
             }),
+            query: buildRouteValidationWithZod(ExecuteConnectorRequestQuery),
           },
         },
       },
@@ -114,7 +116,9 @@ export const postActionsConnectorExecuteRoute = (
             await assistantContext.getAIAssistantConversationsDataClient();
           const promptsDataClient = await assistantContext.getAIAssistantPromptsDataClient();
 
-          const contentReferencesStore = newContentReferencesStore();
+          const contentReferencesStore = newContentReferencesStore({
+            disabled: request.query.content_references_disabled,
+          });
 
           onLlmResponse = async (
             content: string,

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

@@ -247,7 +247,7 @@ export interface AssistantToolParams {
   inference?: InferenceServerStart;
   isEnabledKnowledgeBase: boolean;
   connectorId?: string;
-  contentReferencesStore: ContentReferencesStore | undefined;
+  contentReferencesStore: ContentReferencesStore;
   description?: string;
   esClient: ElasticsearchClient;
   kbDataClient?: AIAssistantKnowledgeBaseDataClient;

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/alert_counts/alert_counts_tool.ts

@@ -39,7 +39,7 @@ export const ALERT_COUNTS_TOOL: AssistantTool = {
       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)
         );
 

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/knowledge_base/knowledge_base_retrieval_tool.ts

@@ -67,9 +67,9 @@ export const KNOWLEDGE_BASE_RETRIEVAL_TOOL: AssistantTool = {
   },
 };
 
-function enrichDocument(contentReferencesStore: ContentReferencesStore | undefined) {
+function enrichDocument(contentReferencesStore: ContentReferencesStore) {
   return (document: Document<Record<string, string>>) => {
-    if (document.id == null || contentReferencesStore == null) {
+    if (document.id == null) {
       return document;
     }
     const documentId = document.id;

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/open_and_acknowledged_alerts/open_and_acknowledged_alerts_tool.ts

@@ -93,7 +93,7 @@ export const OPEN_AND_ACKNOWLEDGED_ALERTS_TOOL: AssistantTool = {
 
             const hitId = hit._id;
             const reference = hitId
-              ? contentReferencesStore?.add((p) => securityAlertReference(p.id, hitId))
+              ? contentReferencesStore.add((p) => securityAlertReference(p.id, hitId))
               : undefined;
             const citation = reference && `\nCitation,${contentReferenceBlock(reference)}`;
 

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/product_docs/product_documentation_tool.ts

@@ -94,11 +94,8 @@ type EnrichedDocument = RetrieveDocumentationResultDoc & {
   citation?: string;
 };
 
-const enrichDocument = (contentReferencesStore: ContentReferencesStore | undefined) => {
+const enrichDocument = (contentReferencesStore: ContentReferencesStore) => {
   return (document: RetrieveDocumentationResultDoc): EnrichedDocument => {
-    if (contentReferencesStore == null) {
-      return document;
-    }
     const reference = contentReferencesStore.add((p) =>
       productDocumentationReference(p.id, document.title, document.url)
     );

x-pack/solutions/security/plugins/security_solution/server/assistant/tools/security_labs/security_labs_tool.ts

@@ -42,7 +42,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')
         );