Add dedicated routes for update and delete

x-pack/platform/packages/shared/kbn-elastic-assistant-common/constants.ts

@@ -44,6 +44,8 @@ export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL =
   `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/{resource?}` as const;
 export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL =
   `${ELASTIC_AI_ASSISTANT_URL}/knowledge_base/entries` as const;
+export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BY_ID =
+  `${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/{id}` as const;
 export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND =
   `${ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL}/_find` as const;
 export const ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION =

x-pack/platform/packages/shared/kbn-elastic-assistant-common/impl/schemas/knowledge_base/entries/common_attributes.gen.ts

@@ -242,6 +242,28 @@ export const KnowledgeBaseEntryCreateProps = z.discriminatedUnion('type', [
   IndexEntryCreateFields,
 ]);
 
+export type KnowledgeBaseEntryDeleteRequestParams = z.infer<
+  typeof KnowledgeBaseEntryDeleteRequestParams
+>;
+
+export const KnowledgeBaseEntryDeleteRequestParams = z.object({
+  id: NonEmptyString,
+});
+
+export type KnowledgeBaseEntryUpdateRequestParams = z.infer<
+  typeof KnowledgeBaseEntryUpdateRequestParams
+>;
+
+export const KnowledgeBaseEntryUpdateRequestParams = z.object({
+  id: NonEmptyString,
+});
+
+export type KnowledgeBaseEntryDeleteResponse = z.infer<typeof KnowledgeBaseEntryDeleteResponse>;
+
+export const KnowledgeBaseEntryDeleteResponse = z.object({
+  id: NonEmptyString,
+});
+
 export type KnowledgeBaseEntryUpdateProps = z.infer<typeof KnowledgeBaseEntryUpdateProps>;
 export const KnowledgeBaseEntryUpdateProps = z.discriminatedUnion('type', [
   DocumentEntryUpdateFields,

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/data_clients.mock.ts

@@ -60,6 +60,8 @@ const createKnowledgeBaseDataClientMock = () => {
     addKnowledgeBaseDocuments: jest.fn(),
     createInferenceEndpoint: jest.fn(),
     createKnowledgeBaseEntry: jest.fn(),
+    updateKnowledgeBaseEntry: jest.fn(),
+    deleteKnowledgeBaseEntry: jest.fn(),
     findDocuments: jest.fn(),
     getAssistantTools: jest.fn(),
     getKnowledgeBaseDocumentEntries: jest.fn(),

x-pack/solutions/security/plugins/elastic_assistant/server/__mocks__/request.ts

@@ -14,6 +14,9 @@ import {
 import type {
   DefendInsightsGetRequestQuery,
   DefendInsightsPostRequestBody,
+  KnowledgeBaseEntryDeleteRequestParams,
+  KnowledgeBaseEntryUpdateProps,
+  KnowledgeBaseEntryUpdateRequestParams,
 } from '@kbn/elastic-assistant-common';
 import {
   AttackDiscoveryPostRequestBody,
@@ -31,6 +34,7 @@ import {
   ELASTIC_AI_ASSISTANT_EVALUATE_URL,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BULK_ACTION,
+  ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BY_ID,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_FIND,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_INDICES_URL,
   ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_URL,
@@ -94,6 +98,27 @@ export const getBulkActionKnowledgeBaseEntryRequest = (
     body,
   });
 
+export const getUpdateKnowledgeBaseEntryRequest = ({
+  params,
+  body,
+}: {
+  params: KnowledgeBaseEntryUpdateRequestParams;
+  body: KnowledgeBaseEntryUpdateProps;
+}) =>
+  requestMock.create({
+    method: 'put',
+    path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BY_ID,
+    params,
+    body,
+  });
+
+export const getDeleteKnowledgeBaseEntryRequest = (params: KnowledgeBaseEntryDeleteRequestParams) =>
+  requestMock.create({
+    method: 'delete',
+    path: ELASTIC_AI_ASSISTANT_KNOWLEDGE_BASE_ENTRIES_URL_BY_ID,
+    params,
+  });
+
 export const getGetCapabilitiesRequest = () =>
   requestMock.create({
     method: 'get',

x-pack/solutions/security/plugins/elastic_assistant/server/ai_assistant_data_clients/knowledge_base/index.ts

@@ -22,6 +22,7 @@ import {
   KnowledgeBaseEntryCreateProps,
   KnowledgeBaseEntryResponse,
   Metadata,
+  KnowledgeBaseEntryUpdateProps,
 } from '@kbn/elastic-assistant-common';
 import pRetry from 'p-retry';
 import { StructuredTool } from '@langchain/core/tools';
@@ -30,21 +31,36 @@ import { IndexPatternsFetcher } from '@kbn/data-views-plugin/server';
 import { map } from 'lodash';
 import { AIAssistantDataClient, AIAssistantDataClientParams } from '..';
 import { GetElser } from '../../types';
-import { createKnowledgeBaseEntry, transformToCreateSchema } from './create_knowledge_base_entry';
-import { EsDocumentEntry, EsIndexEntry, EsKnowledgeBaseEntrySchema } from './types';
-import { transformESSearchToKnowledgeBaseEntry } from './transforms';
+import {
+  createKnowledgeBaseEntry,
+  getUpdateScript,
+  transformToCreateSchema,
+  transformToUpdateSchema,
+} from './create_knowledge_base_entry';
+import {
+  EsDocumentEntry,
+  EsIndexEntry,
+  EsKnowledgeBaseEntrySchema,
+  UpdateKnowledgeBaseEntrySchema,
+} from './types';
+import { transformESSearchToKnowledgeBaseEntry, transformESToKnowledgeBase } from './transforms';
 import { SECURITY_LABS_RESOURCE, USER_RESOURCE } from '../../routes/knowledge_base/constants';
 import {
   getKBVectorSearchQuery,
   getStructuredToolForIndexEntry,
   isModelAlreadyExistsError,
 } from './helpers';
-import { getKBUserFilter } from '../../routes/knowledge_base/entries/utils';
+import {
+  getKBUserFilter,
+  validateDocumentsModification,
+} from '../../routes/knowledge_base/entries/utils';
 import {
   loadSecurityLabs,
   getSecurityLabsDocsCount,
 } from '../../lib/langchain/content_loaders/security_labs_loader';
 import { ASSISTANT_ELSER_INFERENCE_ID } from './field_maps_configuration';
+import { BulkOperationError } from '../../lib/data_stream/documents_data_writer';
+import { AUDIT_OUTCOME, KnowledgeBaseAuditAction, knowledgeBaseAuditEvent } from './audit_events';
 
 /**
  * Params for when creating KbDataClient in Request Context Factory. Useful if needing to modify
@@ -633,6 +649,116 @@ export class AIAssistantKnowledgeBaseDataClient extends AIAssistantDataClient {
     });
   };
 
+  /**
+   * Updates a Knowledge Base Entry.
+   *
+   * @param knowledgeBaseEntryId
+   */
+  public updateKnowledgeBaseEntry = async ({
+    auditLogger,
+    knowledgeBaseEntry,
+  }: {
+    auditLogger?: AuditLogger;
+    knowledgeBaseEntry: KnowledgeBaseEntryUpdateProps;
+  }): Promise<{
+    errors: BulkOperationError[];
+    updatedEntry: KnowledgeBaseEntryResponse;
+  }> => {
+    const authenticatedUser = this.options.currentUser;
+
+    if (authenticatedUser == null) {
+      throw new Error(
+        'Authenticated user not found! Ensure kbDataClient was initialized from a request.'
+      );
+    }
+
+    await validateDocumentsModification(this, authenticatedUser, [knowledgeBaseEntry.id], 'update');
+
+    this.options.logger.debug(
+      () => `Updating Knowledge Base Entry:\n ${JSON.stringify(knowledgeBaseEntry, null, 2)}`
+    );
+    this.options.logger.debug(`kbIndex: ${this.indexTemplateAndPattern.alias}`);
+
+    const writer = await this.getWriter();
+    const changedAt = new Date().toISOString();
+    const { errors, docs_updated: docsUpdated } = await writer.bulk({
+      documentsToUpdate: [
+        transformToUpdateSchema({
+          user: authenticatedUser,
+          updatedAt: changedAt,
+          entry: knowledgeBaseEntry,
+          global: knowledgeBaseEntry.users != null && knowledgeBaseEntry.users.length === 0,
+        }),
+      ],
+      getUpdateScript: (entry: UpdateKnowledgeBaseEntrySchema) => getUpdateScript({ entry }),
+      authenticatedUser,
+    });
+
+    // @ts-ignore-next-line TS2322
+    const updatedEntry = transformESToKnowledgeBase(docsUpdated)?.[0];
+
+    if (updatedEntry) {
+      auditLogger?.log(
+        knowledgeBaseAuditEvent({
+          action: KnowledgeBaseAuditAction.UPDATE,
+          id: updatedEntry.id,
+          name: updatedEntry.name,
+          outcome: AUDIT_OUTCOME.SUCCESS,
+        })
+      );
+    }
+
+    return { errors, updatedEntry };
+  };
+
+  /**
+   * Deletes a new Knowledge Base Entry.
+   *
+   * @param knowledgeBaseEntryId
+   */
+  public deleteKnowledgeBaseEntry = async ({
+    auditLogger,
+    knowledgeBaseEntryId,
+  }: {
+    auditLogger?: AuditLogger;
+    knowledgeBaseEntryId: string;
+  }): Promise<{ errors: BulkOperationError[]; docsDeleted: string[] } | null> => {
+    const authenticatedUser = this.options.currentUser;
+
+    if (authenticatedUser == null) {
+      throw new Error(
+        'Authenticated user not found! Ensure kbDataClient was initialized from a request.'
+      );
+    }
+
+    await validateDocumentsModification(this, authenticatedUser, [knowledgeBaseEntryId], 'delete');
+
+    this.options.logger.debug(
+      () => `Deleting Knowledge Base Entry:\n ID: ${JSON.stringify(knowledgeBaseEntryId, null, 2)}`
+    );
+    this.options.logger.debug(`kbIndex: ${this.indexTemplateAndPattern.alias}`);
+
+    const writer = await this.getWriter();
+    const { errors, docs_deleted: docsDeleted } = await writer.bulk({
+      documentsToDelete: [knowledgeBaseEntryId],
+      authenticatedUser,
+    });
+
+    if (docsDeleted.length) {
+      docsDeleted.forEach((docsDeletedId) => {
+        auditLogger?.log(
+          knowledgeBaseAuditEvent({
+            action: KnowledgeBaseAuditAction.DELETE,
+            id: docsDeletedId,
+            outcome: AUDIT_OUTCOME.SUCCESS,
+          })
+        );
+      });
+    }
+
+    return { errors, docsDeleted };
+  };
+
   /**
    * Returns AssistantTools for any 'relevant' KB IndexEntries that exist in the knowledge base.
    *

x-pack/solutions/security/plugins/elastic_assistant/server/routes/knowledge_base/entries/bulk_actions_route.ts

@@ -48,7 +48,7 @@ import {
   transformToCreateSchema,
   transformToUpdateSchema,
 } from '../../../ai_assistant_data_clients/knowledge_base/create_knowledge_base_entry';
-import { getKBUserFilter } from './utils';
+import { validateDocumentsModification } from './utils';
 
 export interface BulkOperationError {
   message: string;
@@ -235,7 +235,6 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
           const kbDataClient = await ctx.elasticAssistant.getAIAssistantKnowledgeBaseDataClient();
           const spaceId = ctx.elasticAssistant.getSpaceId();
           const authenticatedUser = checkResponse.currentUser;
-          const userFilter = getKBUserFilter(authenticatedUser);
           const manageGlobalKnowledgeBaseAIAssistant =
             kbDataClient?.options.manageGlobalKnowledgeBaseAIAssistant;
 
@@ -266,39 +265,15 @@ export const bulkActionKnowledgeBaseEntriesRoute = (router: ElasticAssistantPlug
             }
           }
 
-          const validateDocumentsModification = async (
-            documentIds: string[],
-            operation: 'delete' | 'update'
-          ) => {
-            if (!documentIds.length) {
-              return;
-            }
-            const documentsFilter = documentIds.map((id) => `_id:${id}`).join(' OR ');
-            const entries = await kbDataClient?.findDocuments<EsKnowledgeBaseEntrySchema>({
-              page: 1,
-              perPage: 100,
-              filter: `${documentsFilter} AND ${userFilter}`,
-            });
-            const availableEntries = entries
-              ? transformESSearchToKnowledgeBaseEntry(entries.data)
-              : [];
-            availableEntries.forEach((entry) => {
-              // RBAC validation
-              const isGlobal = entry.users != null && entry.users.length === 0;
-              if (isGlobal && !manageGlobalKnowledgeBaseAIAssistant) {
-                throw new Error(
-                  `User lacks privileges to ${operation} global knowledge base entries`
-                );
-              }
-            });
-            const availableIds = availableEntries.map((doc) => doc.id);
-            const nonAvailableIds = documentIds.filter((id) => !availableIds.includes(id));
-            if (nonAvailableIds.length > 0) {
-              throw new Error(`Could not find documents to ${operation}: ${nonAvailableIds}.`);
-            }
-          };
-          await validateDocumentsModification(body.delete?.ids ?? [], 'delete');
           await validateDocumentsModification(
+            kbDataClient,
+            authenticatedUser,
+            body.delete?.ids ?? [],
+            'delete'
+          );
+          await validateDocumentsModification(
+            kbDataClient,
+            authenticatedUser,
             body.update?.map((entry) => entry.id) ?? [],
             'update'
           );