[Obs AI Assistant] Add support for dense vector embeddings in API tests

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base.spec.ts

@@ -7,10 +7,12 @@
 
 import expect from '@kbn/expect';
 import { type KnowledgeBaseEntry } from '@kbn/observability-ai-assistant-plugin/common';
+import { orderBy, size, toPairs } from 'lodash';
 import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import {
   clearKnowledgeBase,
   deleteKnowledgeBaseModel,
+  getKnowledgeBaseEntries,
   setupKnowledgeBase,
 } from '../utils/knowledge_base';
 
@@ -40,7 +42,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
     expect(res.status).to.be(200);
 
-    return omitCategories(res.body.entries);
+    return res.body.entries;
   }
 
   describe('Knowledge base', function () {
@@ -59,48 +61,42 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         title: 'My title',
         text: 'My content',
       };
-      it('returns 200 on create', async () => {
+
+      before(async () => {
         const { status } = await observabilityAIAssistantAPIClient.editor({
           endpoint: 'POST /internal/observability_ai_assistant/kb/entries/save',
           params: { body: knowledgeBaseEntry },
         });
         expect(status).to.be(200);
-        const res = await observabilityAIAssistantAPIClient.editor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
-          params: {
-            query: {
-              query: '',
-              sortBy: 'title',
-              sortDirection: 'asc',
-            },
-          },
-        });
-        const entry = res.body.entries[0];
-        expect(entry.id).to.equal(knowledgeBaseEntry.id);
-        expect(entry.title).to.equal(knowledgeBaseEntry.title);
-        expect(entry.text).to.equal(knowledgeBaseEntry.text);
       });
 
-      it('returns 200 on get entries and entry exists', async () => {
-        const res = await observabilityAIAssistantAPIClient.editor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
-          params: {
-            query: {
-              query: '',
-              sortBy: 'title',
-              sortDirection: 'asc',
-            },
-          },
-        });
-
-        expect(res.status).to.be(200);
-        const entry = res.body.entries[0];
+      it('can retrieve the entry', async () => {
+        const entries = await getEntries();
+        const entry = entries[0];
         expect(entry.id).to.equal(knowledgeBaseEntry.id);
         expect(entry.title).to.equal(knowledgeBaseEntry.title);
         expect(entry.text).to.equal(knowledgeBaseEntry.text);
       });
 
-      it('returns 200 on delete', async () => {
+      it('generates sparse embeddings', async () => {
+        const hits = await getKnowledgeBaseEntries(es);
+        const embeddings =
+          hits[0]._source?._inference_fields?.semantic_text?.inference.chunks.semantic_text[0]
+            .embeddings;
+
+        const sorted = orderBy(toPairs(embeddings), [1], ['desc']).slice(0, 5);
+
+        expect(size(embeddings)).to.be.greaterThan(10);
+        expect(sorted).to.eql([
+          ['temperature', 0.07421875],
+          ['used', 0.068359375],
+          ['definition', 0.03955078],
+          ['only', 0.038208008],
+          ['what', 0.028930664],
+        ]);
+      });
+
+      it('can delete the entry', async () => {
         const entryId = 'my-doc-id-1';
         const { status } = await observabilityAIAssistantAPIClient.editor({
           endpoint: 'DELETE /internal/observability_ai_assistant/kb/entries/{entryId}',
@@ -110,21 +106,8 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         });
         expect(status).to.be(200);
 
-        const res = await observabilityAIAssistantAPIClient.editor({
-          endpoint: 'GET /internal/observability_ai_assistant/kb/entries',
-          params: {
-            query: {
-              query: '',
-              sortBy: 'title',
-              sortDirection: 'asc',
-            },
-          },
-        });
-
-        expect(res.status).to.be(200);
-        expect(res.body.entries.filter((entry) => entry.id.startsWith('my-doc-id')).length).to.eql(
-          0
-        );
+        const entries = await getEntries();
+        expect(entries.length).to.eql(0);
       });
 
       it('returns 500 on delete not found', async () => {
@@ -174,9 +157,9 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
         await clearKnowledgeBase(es);
       });
 
-      it('returns 200 on create', async () => {
+      it('creates multiple entries', async () => {
         const entries = await getEntries();
-        expect(omitCategories(entries).length).to.eql(3);
+        expect(entries.length).to.eql(3);
       });
 
       describe('when sorting ', () => {
@@ -352,7 +335,3 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
     });
   });
 }
-
-function omitCategories(entries: KnowledgeBaseEntry[]) {
-  return entries.filter((entry) => entry.labels?.category === undefined);
-}

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_reindex_and_populate_missing_semantic_text_fields.spec.ts

@@ -8,37 +8,15 @@
 import { orderBy } from 'lodash';
 import expect from '@kbn/expect';
 import { AI_ASSISTANT_KB_INFERENCE_ID } from '@kbn/observability-ai-assistant-plugin/server/service/inference_endpoint';
-import { SearchResponse } from '@elastic/elasticsearch/lib/api/types';
-import { KnowledgeBaseEntry } from '@kbn/observability-ai-assistant-plugin/common';
 import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import {
   deleteKnowledgeBaseModel,
   clearKnowledgeBase,
   setupKnowledgeBase,
+  getKnowledgeBaseEntries,
 } from '../utils/knowledge_base';
 import { restoreIndexAssets } from '../utils/index_assets';
 
-interface InferenceChunk {
-  text: string;
-  embeddings: any;
-}
-
-interface InferenceData {
-  inference_id: string;
-  chunks: {
-    semantic_text: InferenceChunk[];
-  };
-}
-
-interface SemanticTextField {
-  semantic_text: string;
-  _inference_fields?: {
-    semantic_text?: {
-      inference: InferenceData;
-    };
-  };
-}
-
 export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
   const esArchiver = getService('esArchiver');
@@ -48,19 +26,6 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
   const archive =
     'x-pack/test/functional/es_archives/observability/ai_assistant/knowledge_base_8_15';
 
-  async function getKnowledgeBaseEntries() {
-    const res = (await es.search({
-      index: '.kibana-observability-ai-assistant-kb*',
-      // Add fields parameter to include inference metadata
-      fields: ['_inference_fields'],
-      query: {
-        match_all: {},
-      },
-    })) as SearchResponse<KnowledgeBaseEntry & SemanticTextField>;
-
-    return res.hits.hits;
-  }
-
   describe('when the knowledge base index was created before 8.15', function () {
     // Intentionally skipped in all serverless environnments (local and MKI)
     // because the migration scenario being tested is not relevant to MKI and Serverless.
@@ -81,7 +46,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
     describe('before migrating', () => {
       it('the docs do not have semantic_text embeddings', async () => {
-        const hits = await getKnowledgeBaseEntries();
+        const hits = await getKnowledgeBaseEntries(es);
         const hasSemanticTextEmbeddings = hits.some((hit) => hit._source?.semantic_text);
         expect(hasSemanticTextEmbeddings).to.be(false);
       });
@@ -98,7 +63,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
       it('the docs have semantic_text embeddings', async () => {
         await retry.try(async () => {
-          const hits = await getKnowledgeBaseEntries();
+          const hits = await getKnowledgeBaseEntries(es);
           const hasSemanticTextEmbeddings = hits.every((hit) => hit._source?.semantic_text);
           expect(hasSemanticTextEmbeddings).to.be(true);
 

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_setup.spec.ts

@@ -7,7 +7,7 @@
 
 import expect from '@kbn/expect';
 import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
-import { TINY_ELSER, deleteKnowledgeBaseModel, setupKnowledgeBase } from '../utils/knowledge_base';
+import { TINY_MODELS, deleteKnowledgeBaseModel, setupKnowledgeBase } from '../utils/knowledge_base';
 import { restoreIndexAssets } from '../utils/index_assets';
 
 export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
@@ -52,7 +52,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
           endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
           params: {
             query: {
-              model_id: TINY_ELSER.id,
+              model_id: TINY_MODELS.ELSER,
             },
           },
         });

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/knowledge_base/knowledge_base_status.spec.ts

@@ -10,9 +10,9 @@ import { KnowledgeBaseState } from '@kbn/observability-ai-assistant-plugin/commo
 import type { DeploymentAgnosticFtrProviderContext } from '../../../../ftr_provider_context';
 import {
   deleteKnowledgeBaseModel,
-  TINY_ELSER,
   deleteInferenceEndpoint,
   setupKnowledgeBase,
+  TINY_MODELS,
 } from '../utils/knowledge_base';
 
 export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderContext) {
@@ -40,7 +40,7 @@ export default function ApiTest({ getService }: DeploymentAgnosticFtrProviderCon
 
       expect(res.body.kbState).to.be(KnowledgeBaseState.READY);
       expect(res.body.enabled).to.be(true);
-      expect(res.body.endpoint?.service_settings?.model_id).to.eql(TINY_ELSER.id);
+      expect(res.body.endpoint?.service_settings?.model_id).to.eql(TINY_MODELS.ELSER);
     });
 
     it('returns correct status after model is deleted', async () => {

x-pack/test/api_integration/deployment_agnostic/apis/observability/ai_assistant/utils/knowledge_base.ts

@@ -12,6 +12,7 @@ import { ToolingLog } from '@kbn/tooling-log';
 import { RetryService } from '@kbn/ftr-common-functional-services';
 import {
   Instruction,
+  KnowledgeBaseEntry,
   KnowledgeBaseState,
 } from '@kbn/observability-ai-assistant-plugin/common/types';
 import { resourceNames } from '@kbn/observability-ai-assistant-plugin/server/service';
@@ -21,29 +22,29 @@ import { MachineLearningProvider } from '../../../../../services/ml';
 import { SUPPORTED_TRAINED_MODELS } from '../../../../../../functional/services/ml/api';
 import { setAdvancedSettings } from './advanced_settings';
 
-export const TINY_ELSER = {
-  ...SUPPORTED_TRAINED_MODELS.TINY_ELSER,
-  id: SUPPORTED_TRAINED_MODELS.TINY_ELSER.name,
-};
+export const TINY_MODELS = {
+  ELSER: SUPPORTED_TRAINED_MODELS.TINY_ELSER.name,
+  TEXT_EMBEDDING: SUPPORTED_TRAINED_MODELS.TINY_TEXT_EMBEDDING.name,
+} as const;
 
-export async function importTinyElserModel(ml: ReturnType<typeof MachineLearningProvider>) {
-  const config = {
-    ...ml.api.getTrainedModelConfig(TINY_ELSER.name),
-    input: {
-      field_names: ['text_field'],
-    },
-  };
+type TinyModelID = (typeof TINY_MODELS)[keyof typeof TINY_MODELS];
+
+async function importModel(ml: ReturnType<typeof MachineLearningProvider>, modelId: TinyModelID) {
   // necessary for MKI, check indices before importing model.  compatible with stateful
   await ml.api.assureMlStatsIndexExists();
-  await ml.api.importTrainedModel(TINY_ELSER.name, TINY_ELSER.id, config);
+
+  const config = ml.api.getTrainedModelConfig(modelId);
+  await ml.api.importTrainedModel(modelId, modelId, config);
 }
 
 export async function setupKnowledgeBase(
   getService: DeploymentAgnosticFtrProviderContext['getService'],
   {
     deployModel: deployModel = true,
+    modelId = TINY_MODELS.ELSER,
   }: {
     deployModel?: boolean;
+    modelId?: TinyModelID;
   } = {}
 ) {
   const log = getService('log');
@@ -52,14 +53,15 @@ export async function setupKnowledgeBase(
   const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantApi');
 
   if (deployModel) {
-    await importTinyElserModel(ml);
+    await importModel(ml, modelId);
+    await ml.api.startTrainedModelDeploymentES(modelId);
   }
 
   const { status, body } = await observabilityAIAssistantAPIClient.admin({
     endpoint: 'POST /internal/observability_ai_assistant/kb/setup',
     params: {
       query: {
-        model_id: TINY_ELSER.id,
+        model_id: modelId,
       },
     },
   });
@@ -94,17 +96,19 @@ export async function deleteKnowledgeBaseModel(
   getService: DeploymentAgnosticFtrProviderContext['getService'],
   {
     shouldDeleteInferenceEndpoint = true,
+    modelId = TINY_MODELS.ELSER,
   }: {
     shouldDeleteInferenceEndpoint?: boolean;
+    modelId?: TinyModelID;
   } = {}
 ) {
   const log = getService('log');
   const ml = getService('ml');
   const es = getService('es');
 
   try {
-    await ml.api.stopTrainedModelDeploymentES(TINY_ELSER.id, true);
-    await ml.api.deleteTrainedModelES(TINY_ELSER.id);
+    await ml.api.stopTrainedModelDeploymentES(modelId, true);
+    await ml.api.deleteTrainedModelES(modelId);
     await ml.testResources.cleanMLSavedObjects();
 
     if (shouldDeleteInferenceEndpoint) {
@@ -233,3 +237,34 @@ export async function hasIndexWriteBlock(es: Client, index: string) {
   const writeBlockSetting = Object.values(response)[0]?.settings?.index?.blocks?.write;
   return writeBlockSetting === 'true' || writeBlockSetting === true;
 }
+
+interface SemanticTextField {
+  semantic_text: string;
+  _inference_fields?: {
+    semantic_text?: {
+      inference: {
+        inference_id: string;
+        chunks: {
+          semantic_text: Array<{
+            embeddings:
+              | Record<string, number> // sparse embedding
+              | number[]; // dense embedding;
+          }>;
+        };
+      };
+    };
+  };
+}
+
+export async function getKnowledgeBaseEntries(es: Client) {
+  const res = await es.search<KnowledgeBaseEntry & SemanticTextField>({
+    index: resourceNames.aliases.kb,
+    // Add fields parameter to include inference metadata
+    fields: ['_inference_fields'],
+    query: {
+      match_all: {},
+    },
+  });
+
+  return res.hits.hits;
+}

x-pack/test/functional/services/ml/api.ts

@@ -1591,6 +1591,16 @@ export function MachineLearningAPIProvider({ getService }: FtrProviderContext) {
       }
     },
 
+    async startTrainedModelDeploymentES(modelId: string) {
+      log.debug(`Starting trained model deployment with id "${modelId}"`);
+      const url = `/_ml/trained_models/${modelId}/deployment/_start`;
+
+      const { body, status } = await esSupertest.post(url);
+      this.assertResponseStatusCode(200, status, body);
+
+      log.debug('> Trained model deployment started');
+    },
+
     async stopTrainedModelDeploymentES(deploymentId: string, force: boolean = false) {
       log.debug(`Stopping trained model deployment with id "${deploymentId}"`);
       const url = `/_ml/trained_models/${deploymentId}/deployment/_stop${