[8.18] [Obs AI Assistant] Make KB retrieval namespace specific (#213505) (#213984)

# Backport This will backport the following commits from `main` to `8.18`: - [[Obs AI Assistant] Make KB retrieval namespace specific (#213505)](https://github.com/elastic/kibana/pull/213505)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Viduni Wickramarachchi <viduni.wickramarachchi@elastic.co>
2025-04-24 17:59:23 -04:00 · 2025-03-12 06:37:53 +11:00 · 2025-03-12 06:37:53 +11:00 · f8e55ed735
commit f8e55ed735
parent d7ff506e03
4 changed files with 173 additions and 18 deletions
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/client/index.ts
@ -771,7 +771,12 @@ export class ObservabilityAIAssistantClient {
    sortBy: string;
    sortDirection: 'asc' | 'desc';
  }) => {
-    return this.dependencies.knowledgeBaseService.getEntries({ query, sortBy, sortDirection });
+    return this.dependencies.knowledgeBaseService.getEntries({
+      query,
+      sortBy,
+      sortDirection,
+      namespace: this.dependencies.namespace,
+    });
  };

  deleteKnowledgeBaseEntry = async (id: string) => {
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/knowledge_base_service/index.ts
@ -19,6 +19,7 @@ import {
 } from '../../../common/types';
 import { getAccessQuery, getUserAccessFilters } from '../util/get_access_query';
 import { getCategoryQuery } from '../util/get_category_query';
+import { getSpaceQuery } from '../util/get_space_query';
 import {
  createInferenceEndpoint,
  deleteInferenceEndpoint,
@ -259,14 +260,17 @@ export class KnowledgeBaseService {
    query,
    sortBy,
    sortDirection,
+    namespace,
  }: {
    query?: string;
    sortBy?: string;
    sortDirection?: 'asc' | 'desc';
+    namespace: string;
  }): Promise<{ entries: KnowledgeBaseEntry[] }> => {
    if (!this.dependencies.config.enableKnowledgeBase) {
      return { entries: [] };
    }
+
    try {
      const response = await this.dependencies.esClient.asInternalUser.search<
        KnowledgeBaseEntry & { doc_id?: string }
@ -281,8 +285,12 @@ export class KnowledgeBaseService {
                : []),
              {
                // exclude user instructions
-                bool: { must_not: { term: { type: KnowledgeBaseType.UserInstruction } } },
+                bool: {
+                  must_not: { term: { type: KnowledgeBaseType.UserInstruction } },
+                },
              },
+              // filter by space
+              ...getSpaceQuery({ namespace }),
            ],
          },
        },
@ -425,6 +433,7 @@ export class KnowledgeBaseService {
        },
        refresh: 'wait_for',
      });
+
      this.dependencies.logger.debug(`Entry added to knowledge base`);
    } catch (error) {
      this.dependencies.logger.debug(`Failed to add entry to knowledge base ${error}`);
--- a/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/util/get_space_query.ts
+++ b/x-pack/platform/plugins/shared/observability_ai_assistant/server/service/util/get_space_query.ts
@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+export function getSpaceQuery({ namespace }: { namespace?: string }) {
+  return [
+    {
+      bool: {
+        should: [
+          { term: { namespace } },
+          { bool: { must_not: { exists: { field: 'namespace' } } } },
+        ],
+      },
+    },
+  ];
+}
--- a/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base.spec.ts
+++ b/x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base.spec.ts
@ -24,6 +24,31 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
  const retry = getService('retry');
  const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');

+  async function getEntries({
+    query = '',
+    sortBy = 'title',
+    sortDirection = 'asc',
+    spaceId,
+    user = 'editor',
+  }: {
+    query?: string;
+    sortBy?: string;
+    sortDirection?: 'asc' | 'desc';
+    spaceId?: string;
+    user?: 'admin' | 'editor' | 'viewer';
+  } = {}): Promise<KnowledgeBaseEntry[]> {
+    const res = await observabilityAIAssistantAPIClient[user]({
+      endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
+      params: {
+        query: { query, sortBy, sortDirection },
+      },
+      spaceId,
+    });
+    expect(res.status).to.be(200);
+
+    return omitCategories(res.body.entries);
+  }
+
  describe('Knowledge base', function () {
    before(async () => {
      await importTinyElserModel(ml);
@ -124,22 +149,6 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
    });

    describe('when managing multiple entries', () => {
-      async function getEntries({
-        query = '',
-        sortBy = 'title',
-        sortDirection = 'asc',
-      }: { query?: string; sortBy?: string; sortDirection?: 'asc' | 'desc' } = {}) {
-        const res = await observabilityAIAssistantAPIClient.editor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
-          params: {
-            query: { query, sortBy, sortDirection },
-          },
-        });
-        expect(res.status).to.be(200);
-
-        return omitCategories(res.body.entries);
-      }
-
      beforeEach(async () => {
        await clearKnowledgeBase(es);

@ -200,6 +209,119 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
      });
    });

+    describe('when managing multiple entries across spaces', () => {
+      const SPACE_A_ID = 'space_a';
+      const SPACE_B_ID = 'space_b';
+
+      before(async () => {
+        await clearKnowledgeBase(es);
+
+        const { status: importStatusForSpaceA } = await observabilityAIAssistantAPIClient.admin({
+          endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import',
+          params: {
+            body: {
+              entries: [
+                {
+                  id: 'my-doc-1',
+                  title: `Entry in Space A by Admin 1`,
+                  text: `This is a public entry in Space A created by Admin`,
+                  public: true,
+                },
+                {
+                  id: 'my-doc-2',
+                  title: `Entry in Space A by Admin 2`,
+                  text: `This is a private entry in Space A created by Admin`,
+                  public: false,
+                },
+              ],
+            },
+          },
+          spaceId: SPACE_A_ID,
+        });
+        expect(importStatusForSpaceA).to.be(200);
+
+        const { status: importStatusForSpaceB } = await observabilityAIAssistantAPIClient.admin({
+          endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import',
+          params: {
+            body: {
+              entries: [
+                {
+                  id: 'my-doc-3',
+                  title: `Entry in Space B by Admin 3`,
+                  text: `This is a public entry in Space B created by Admin`,
+                  public: true,
+                },
+                {
+                  id: 'my-doc-4',
+                  title: `Entry in Space B by Admin 4`,
+                  text: `This is a private entry in Space B created by Admin`,
+                  public: false,
+                },
+              ],
+            },
+          },
+          spaceId: SPACE_B_ID,
+        });
+        expect(importStatusForSpaceB).to.be(200);
+      });
+
+      after(async () => {
+        await clearKnowledgeBase(es);
+      });
+
+      it('ensures users can only access entries relevant to their namespace', async () => {
+        // User (admin) in space A should only see entries in space A
+        const spaceAEntries = await getEntries({
+          user: 'admin',
+          spaceId: SPACE_A_ID,
+        });
+
+        expect(spaceAEntries.length).to.be(2);
+
+        expect(spaceAEntries[0].id).to.equal('my-doc-1');
+        expect(spaceAEntries[0].public).to.be(true);
+        expect(spaceAEntries[0].title).to.equal('Entry in Space A by Admin 1');
+
+        expect(spaceAEntries[1].id).to.equal('my-doc-2');
+        expect(spaceAEntries[1].public).to.be(false);
+        expect(spaceAEntries[1].title).to.equal('Entry in Space A by Admin 2');
+
+        // User (admin) in space B should only see entries in space B
+        const spaceBEntries = await getEntries({
+          user: 'admin',
+          spaceId: SPACE_B_ID,
+        });
+
+        expect(spaceBEntries.length).to.be(2);
+
+        expect(spaceBEntries[0].id).to.equal('my-doc-3');
+        expect(spaceBEntries[0].public).to.be(true);
+        expect(spaceBEntries[0].title).to.equal('Entry in Space B by Admin 3');
+
+        expect(spaceBEntries[1].id).to.equal('my-doc-4');
+        expect(spaceBEntries[1].public).to.be(false);
+        expect(spaceBEntries[1].title).to.equal('Entry in Space B by Admin 4');
+      });
+
+      it('should allow a user who is not the owner of the entries to access entries relevant to their namespace', async () => {
+        // User (editor) in space B should only see entries in space B
+        const spaceBEntries = await getEntries({
+          user: 'editor',
+          spaceId: SPACE_B_ID,
+        });
+
+        expect(spaceBEntries.length).to.be(2);
+
+        expect(spaceBEntries[0].id).to.equal('my-doc-3');
+        expect(spaceBEntries[0].public).to.be(true);
+        expect(spaceBEntries[0].title).to.equal('Entry in Space B by Admin 3');
+
+        expect(spaceBEntries[1].id).to.equal('my-doc-4');
+        expect(spaceBEntries[1].public).to.be(false);
+        expect(spaceBEntries[1].title).to.equal('Entry in Space B by Admin 4');
+      });
+    });
+
    describe('security roles and access privileges', () => {
      describe('should deny access for users without the ai_assistant privilege', () => {
        it('POST /internal/observability_ai_assistant/kb/entries/save', async () => {