diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts index 951b1fc716730..7d20cd927b42f 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts @@ -68,13 +68,14 @@ export function registerSummarizationFunction({ signal ) => { // The LLM should be able to update an existing entry by providing the same doc_id - // if no id is provided, we generate a new one + // if no existing entry is found, we generate a uuid const id = await client.getUuidFromDocId(docId); return client .addKnowledgeBaseEntry({ entry: { id: id ?? v4(), + title: docId, // use doc_id as title for now doc_id: docId, role: KnowledgeBaseEntryRole.AssistantSummarization, text, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts index 3ed206d68451c..e3959ed35250f 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts @@ -12,7 +12,6 @@ import type { import { notImplemented } from '@hapi/boom'; import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils'; import * as t from 'io-ts'; -import { v4 } from 'uuid'; import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route'; import { Instruction, @@ -240,10 +239,15 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({ params: t.type({ body: t.type({ entries: t.array( - t.type({ - doc_id: t.string, - text: nonEmptyStringRt, - }) + t.intersection([ + t.type({ + id: t.string, + text: nonEmptyStringRt, + }), + t.partial({ + title: t.string, + }), + ]) ), }), }), @@ -257,18 +261,48 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({ throw notImplemented(); } - const entries = resources.params.body.entries.map((entry) => ({ - id: v4(), + const formattedEntries = resources.params.body.entries.map((entry) => ({ + id: entry.id, + title: entry.title, + text: entry.text, confidence: 'high' as KnowledgeBaseEntry['confidence'], is_correction: false, type: 'contextual' as const, public: true, labels: {}, role: KnowledgeBaseEntryRole.UserEntry, - ...entry, })); - return await client.importKnowledgeBaseEntries({ entries }); + return await client.importKnowledgeBaseEntries({ entries: formattedEntries }); + }, +}); + +const importKnowledgeBaseCategoryEntries = createObservabilityAIAssistantServerRoute({ + endpoint: 'POST /internal/observability_ai_assistant/kb/entries/category/import', + params: t.type({ + body: t.type({ + category: t.string, + entries: t.array( + t.type({ + id: t.string, + texts: t.array(t.string), + }) + ), + }), + }), + options: { + tags: ['access:ai_assistant'], + }, + handler: async (resources): Promise => { + const client = await resources.service.getClient({ request: resources.request }); + + if (!client) { + throw notImplemented(); + } + + const { entries, category } = resources.params.body; + + return resources.service.addCategoryToKnowledgeBase(category, entries); }, }); @@ -279,6 +313,7 @@ export const knowledgeBaseRoutes = { ...saveKnowledgeBaseUserInstruction, ...getKnowledgeBaseUserInstructions, ...importKnowledgeBaseEntries, + ...importKnowledgeBaseCategoryEntries, ...saveKnowledgeBaseEntry, ...deleteKnowledgeBaseEntry, }; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts index 002847e9349e0..21126748c3372 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/index.ts @@ -326,35 +326,34 @@ export class ObservabilityAIAssistantService { addToKnowledgeBaseQueue(entries: KnowledgeBaseEntryRequest[]): void { this.init() .then(() => { - this.kbService!.queue( - entries.flatMap((entry) => { - const entryWithSystemProperties = { - ...entry, - doc_id: entry.id, - '@timestamp': new Date().toISOString(), - public: true, - confidence: 'high' as const, - type: 'contextual' as const, - is_correction: false, - labels: { - ...entry.labels, - }, - role: KnowledgeBaseEntryRole.Elastic, - }; - - const operations = - 'texts' in entryWithSystemProperties - ? splitKbText(entryWithSystemProperties) - : [ - { - type: KnowledgeBaseEntryOperationType.Index, - document: entryWithSystemProperties, - }, - ]; - - return operations; - }) - ); + const operations = entries.flatMap((entry) => { + const entryWithSystemProperties = { + ...entry, + doc_id: entry.id, + '@timestamp': new Date().toISOString(), + public: true, + confidence: 'high' as const, + type: 'contextual' as const, + is_correction: false, + labels: { + ...entry.labels, + }, + role: KnowledgeBaseEntryRole.Elastic, + }; + + return 'texts' in entryWithSystemProperties + ? splitKbText(entryWithSystemProperties) + : [ + { + type: KnowledgeBaseEntryOperationType.Index, + document: entryWithSystemProperties, + }, + ]; + }); + + this.logger.debug(`Queuing ${operations.length} operations (${entries.length} entries)`); + + this.kbService!.queue(operations); }) .catch((error) => { this.logger.error( diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/kb_component_template.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/kb_component_template.ts index 5467556a0e3ab..0c654aef5db39 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/kb_component_template.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/kb_component_template.ts @@ -32,7 +32,15 @@ export const kbComponentTemplate: ClusterComponentTemplate['component_template'] '@timestamp': date, id: keyword, doc_id: { type: 'text', fielddata: true }, - title: { type: 'text', fielddata: true }, + title: { + type: 'text', + fields: { + keyword: { + type: 'keyword', + ignore_above: 256, + }, + }, + }, user: { properties: { id: keyword, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts index c1ade1d4f6be4..d5f55307cb727 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts @@ -48,8 +48,8 @@ export interface RecalledEntry { function isModelMissingOrUnavailableError(error: Error) { return ( error instanceof errors.ResponseError && - (error.body.error.type === 'resource_not_found_exception' || - error.body.error.type === 'status_exception') + (error.body.error?.type === 'resource_not_found_exception' || + error.body.error?.type === 'status_exception') ); } function isCreateModelValidationError(error: Error) { @@ -70,7 +70,7 @@ export enum KnowledgeBaseEntryOperationType { interface KnowledgeBaseDeleteOperation { type: KnowledgeBaseEntryOperationType.Delete; - doc_id?: string; + groupId?: string; labels?: Record; } @@ -84,7 +84,7 @@ export type KnowledgeBaseEntryOperation = | KnowledgeBaseIndexOperation; export class KnowledgeBaseService { - private hasSetup: boolean = false; + private isModelReady: boolean = false; private _queue: KnowledgeBaseEntryOperation[] = []; @@ -93,6 +93,7 @@ export class KnowledgeBaseService { } setup = async () => { + this.dependencies.logger.debug('Setting up knowledge base'); const elserModelId = await this.dependencies.getModelId(); const retryOptions = { factor: 1, minTimeout: 10000, retries: 12 }; @@ -189,7 +190,7 @@ export class KnowledgeBaseService { ); if (isReady) { - return Promise.resolve(); + return; } this.dependencies.logger.debug('Model is not allocated yet'); @@ -234,7 +235,7 @@ export class KnowledgeBaseService { query: { bool: { filter: [ - ...(operation.doc_id ? [{ term: { _id: operation.doc_id } }] : []), + ...(operation.groupId ? [{ term: { doc_id: operation.groupId } }] : []), ...(operation.labels ? map(operation.labels, (value, key) => { return { term: { [key]: value } }; @@ -247,7 +248,7 @@ export class KnowledgeBaseService { return; } catch (error) { this.dependencies.logger.error( - `Failed to delete document "${operation?.doc_id}" due to ${error.message}` + `Failed to delete document "${operation?.groupId}" due to ${error.message}` ); this.dependencies.logger.debug(() => JSON.stringify(operation)); throw error; @@ -275,7 +276,7 @@ export class KnowledgeBaseService { this.dependencies.logger.debug(`Processing queue`); - this.hasSetup = true; + this.isModelReady = true; this.dependencies.logger.info(`Processing ${this._queue.length} queue operations`); @@ -292,7 +293,7 @@ export class KnowledgeBaseService { ) ); - this.dependencies.logger.info('Processed all queued operations'); + this.dependencies.logger.info(`Finished processing ${operations.length} queued operations`); } queue(operations: KnowledgeBaseEntryOperation[]): void { @@ -300,8 +301,15 @@ export class KnowledgeBaseService { return; } - if (!this.hasSetup) { - this._queue.push(...operations); + this.dependencies.logger.debug( + `Adding ${operations.length} operations to queue. Queue size now: ${this._queue.length})` + ); + this._queue.push(...operations); + + if (!this.isModelReady) { + this.dependencies.logger.debug( + `Delay processing ${operations.length} operations until knowledge base is ready` + ); return; } @@ -311,13 +319,18 @@ export class KnowledgeBaseService { limiter(() => this.processOperation(operation)) ); - Promise.all(limitedFunctions).catch((err) => { - this.dependencies.logger.error(`Failed to process all queued operations`); - this.dependencies.logger.error(err); - }); + Promise.all(limitedFunctions) + .then(() => { + this.dependencies.logger.debug(`Processed all queued operations`); + }) + .catch((err) => { + this.dependencies.logger.error(`Failed to process all queued operations`); + this.dependencies.logger.error(err); + }); } status = async () => { + this.dependencies.logger.debug('Checking model status'); const elserModelId = await this.dependencies.getModelId(); try { @@ -327,14 +340,23 @@ export class KnowledgeBaseService { const elserModelStats = modelStats.trained_model_stats[0]; const deploymentState = elserModelStats.deployment_stats?.state; const allocationState = elserModelStats.deployment_stats?.allocation_status.state; + const ready = deploymentState === 'started' && allocationState === 'fully_allocated'; + + this.dependencies.logger.debug( + `Model deployment state: ${deploymentState}, allocation state: ${allocationState}, ready: ${ready}` + ); return { - ready: deploymentState === 'started' && allocationState === 'fully_allocated', + ready, deployment_state: deploymentState, allocation_state: allocationState, model_name: elserModelId, }; } catch (error) { + this.dependencies.logger.debug( + `Failed to get status for model "${elserModelId}" due to ${error.message}` + ); + return { error: error instanceof errors.ResponseError ? error.body.error : String(error), ready: false, @@ -533,8 +555,10 @@ export class KnowledgeBaseService { query: { bool: { filter: [ - // filter title by query - ...(query ? [{ wildcard: { doc_id: { value: `${query}*` } } }] : []), + // filter by search query + ...(query + ? [{ query_string: { query: `${query}*`, fields: ['doc_id', 'title'] } }] + : []), { // exclude user instructions bool: { must_not: { term: { type: KnowledgeBaseType.UserInstruction } } }, @@ -542,13 +566,13 @@ export class KnowledgeBaseService { ], }, }, - sort: [ - { - [String(sortBy)]: { - order: sortDirection, - }, - }, - ], + sort: + sortBy === 'title' + ? [ + { ['title.keyword']: { order: sortDirection } }, + { doc_id: { order: sortDirection } }, // sort by doc_id for backwards compatibility + ] + : [{ [String(sortBy)]: { order: sortDirection } }], size: 500, _source: { includes: [ @@ -661,6 +685,7 @@ export class KnowledgeBaseService { document: { '@timestamp': new Date().toISOString(), ...doc, + title: doc.title ?? doc.doc_id, user, namespace, }, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/util/split_kb_text.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/util/split_kb_text.ts index 9a2f047b60f9b..4b3dd559df3eb 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/util/split_kb_text.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/util/split_kb_text.ts @@ -20,14 +20,14 @@ export function splitKbText({ return [ { type: KnowledgeBaseEntryOperationType.Delete, - doc_id: id, + groupId: id, // delete all entries with the same groupId labels: {}, }, ...texts.map((text, index) => ({ type: KnowledgeBaseEntryOperationType.Index, document: merge({}, rest, { id: [id, index].join('_'), - doc_id: id, + doc_id: id, // group_id is used to group entries together labels: {}, text, }), diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts index 1f35986bae8f0..c71676c60a6c4 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts @@ -66,7 +66,7 @@ export async function scoreSuggestions({ Documents: ${JSON.stringify( - suggestions.map(({ id, docId: title, text }) => ({ id, title, text })), + suggestions.map(({ id, doc_id: title, text }) => ({ id, title, text })), null, 2 )}`); diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant_management/public/routes/components/knowledge_base_tab.test.tsx b/x-pack/plugins/observability_solution/observability_ai_assistant_management/public/routes/components/knowledge_base_tab.test.tsx index d8e2897c6878c..94ae564ac7f96 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant_management/public/routes/components/knowledge_base_tab.test.tsx +++ b/x-pack/plugins/observability_solution/observability_ai_assistant_management/public/routes/components/knowledge_base_tab.test.tsx @@ -81,7 +81,9 @@ describe('KnowledgeBaseTab', () => { getByTestId('knowledgeBaseEditManualEntryFlyoutSaveButton').click(); - expect(createMock).toHaveBeenCalledWith({ entry: { id: 'foo', public: false, text: 'bar' } }); + expect(createMock).toHaveBeenCalledWith({ + entry: { id: expect.any(String), title: 'foo', public: false, text: 'bar' }, + }); }); it('should require an id', () => { diff --git a/x-pack/test/observability_ai_assistant_api_integration/common/config.ts b/x-pack/test/observability_ai_assistant_api_integration/common/config.ts index 0a2a4c796eb5f..50d660ae1a0dd 100644 --- a/x-pack/test/observability_ai_assistant_api_integration/common/config.ts +++ b/x-pack/test/observability_ai_assistant_api_integration/common/config.ts @@ -23,6 +23,10 @@ export type CreateTestConfig = ReturnType; export type CreateTest = ReturnType; +export type ObservabilityAIAssistantApiClients = Awaited< + ReturnType +>; + export type ObservabilityAIAssistantAPIClient = Awaited< ReturnType >; diff --git a/x-pack/test/observability_ai_assistant_api_integration/tests/complete/functions/summarize.spec.ts b/x-pack/test/observability_ai_assistant_api_integration/tests/complete/functions/summarize.spec.ts index 238be31220aa9..ea929c34526b8 100644 --- a/x-pack/test/observability_ai_assistant_api_integration/tests/complete/functions/summarize.spec.ts +++ b/x-pack/test/observability_ai_assistant_api_integration/tests/complete/functions/summarize.spec.ts @@ -64,7 +64,7 @@ export default function ApiTest({ getService }: FtrProviderContext) { params: { query: { query: '', - sortBy: 'doc_id', + sortBy: 'title', sortDirection: 'asc', }, }, diff --git a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base.spec.ts b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base.spec.ts index b9983e5e5c58c..a2ebee43b54dc 100644 --- a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base.spec.ts +++ b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base.spec.ts @@ -6,31 +6,35 @@ */ import expect from '@kbn/expect'; +import { type KnowledgeBaseEntry } from '@kbn/observability-ai-assistant-plugin/common'; +import pRetry from 'p-retry'; +import { ToolingLog } from '@kbn/tooling-log'; +import { uniq } from 'lodash'; import { FtrProviderContext } from '../../common/ftr_provider_context'; import { clearKnowledgeBase, createKnowledgeBaseModel, deleteKnowledgeBaseModel } from './helpers'; +import { ObservabilityAIAssistantApiClients } from '../../common/config'; +import { ObservabilityAIAssistantApiClient } from '../../common/observability_ai_assistant_api_client'; export default function ApiTest({ getService }: FtrProviderContext) { const ml = getService('ml'); const es = getService('es'); - + const log = getService('log'); const observabilityAIAssistantAPIClient = getService('observabilityAIAssistantAPIClient'); describe('Knowledge base', () => { before(async () => { await createKnowledgeBaseModel(ml); - }); - after(async () => { - await deleteKnowledgeBaseModel(ml); - }); - - it('returns 200 on knowledge base setup', async () => { - const res = await observabilityAIAssistantAPIClient + await observabilityAIAssistantAPIClient .editorUser({ endpoint: 'POST /internal/observability_ai_assistant/kb/setup', }) .expect(200); - expect(res.body).to.eql({}); + }); + + after(async () => { + await deleteKnowledgeBaseModel(ml); + await clearKnowledgeBase(es); }); describe('when managing a single entry', () => { @@ -51,7 +55,7 @@ export default function ApiTest({ getService }: FtrProviderContext) { params: { query: { query: '', - sortBy: 'doc_id', + sortBy: 'title', sortDirection: 'asc', }, }, @@ -69,7 +73,7 @@ export default function ApiTest({ getService }: FtrProviderContext) { params: { query: { query: '', - sortBy: 'doc_id', + sortBy: 'title', sortDirection: 'asc', }, }, @@ -98,7 +102,7 @@ export default function ApiTest({ getService }: FtrProviderContext) { params: { query: { query: '', - sortBy: 'doc_id', + sortBy: 'title', sortDirection: 'asc', }, }, @@ -123,195 +127,343 @@ export default function ApiTest({ getService }: FtrProviderContext) { }); describe('when managing multiple entries', () => { - before(async () => { - await clearKnowledgeBase(es); - }); - afterEach(async () => { - await clearKnowledgeBase(es); - }); - const knowledgeBaseEntries = [ - { - doc_id: 'my_doc_a', - text: 'My content a', - }, - { - doc_id: 'my_doc_b', - text: 'My content b', - }, - { - doc_id: 'my_doc_c', - text: 'My content c', - }, - ]; - it('returns 200 on create', async () => { - await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import', - params: { body: { entries: knowledgeBaseEntries } }, - }) - .expect(200); - + async function getEntries({ + query = '', + sortBy = 'title', + sortDirection = 'asc', + }: { query?: string; sortBy?: string; sortDirection?: 'asc' | 'desc' } = {}) { const res = await observabilityAIAssistantAPIClient .editorUser({ endpoint: 'GET /internal/observability_ai_assistant/kb/entries', params: { - query: { - query: '', - sortBy: 'doc_id', - sortDirection: 'asc', - }, + query: { query, sortBy, sortDirection }, }, }) .expect(200); - expect(res.body.entries.filter((entry) => entry.id.startsWith('my_doc')).length).to.eql(3); - }); - it('allows sorting', async () => { + return omitCategories(res.body.entries); + } + + beforeEach(async () => { + await clearKnowledgeBase(es); + await observabilityAIAssistantAPIClient .editorUser({ endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import', - params: { body: { entries: knowledgeBaseEntries } }, - }) - .expect(200); - - const res = await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'GET /internal/observability_ai_assistant/kb/entries', params: { - query: { - query: '', - sortBy: 'doc_id', - sortDirection: 'desc', + body: { + entries: [ + { + id: 'my_doc_a', + title: 'My title a', + text: 'My content a', + }, + { + id: 'my_doc_b', + title: 'My title b', + text: 'My content b', + }, + { + id: 'my_doc_c', + title: 'My title c', + text: 'My content c', + }, + ], }, }, }) .expect(200); + }); - const entries = res.body.entries.filter((entry) => entry.id.startsWith('my_doc')); - expect(entries[0].id).to.eql('my_doc_c'); - expect(entries[1].id).to.eql('my_doc_b'); - expect(entries[2].id).to.eql('my_doc_a'); - - // asc - const resAsc = await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'GET /internal/observability_ai_assistant/kb/entries', - params: { - query: { - query: '', - sortBy: 'doc_id', - sortDirection: 'asc', - }, - }, - }) - .expect(200); + afterEach(async () => { + await clearKnowledgeBase(es); + }); - const entriesAsc = resAsc.body.entries.filter((entry) => entry.id.startsWith('my_doc')); - expect(entriesAsc[0].id).to.eql('my_doc_a'); - expect(entriesAsc[1].id).to.eql('my_doc_b'); - expect(entriesAsc[2].id).to.eql('my_doc_c'); + it('returns 200 on create', async () => { + const entries = await getEntries(); + expect(omitCategories(entries).length).to.eql(3); }); - it('allows searching', async () => { - await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import', - params: { body: { entries: knowledgeBaseEntries } }, - }) - .expect(200); - const res = await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'GET /internal/observability_ai_assistant/kb/entries', - params: { - query: { - query: 'my_doc_a', - sortBy: 'doc_id', - sortDirection: 'asc', - }, - }, - }) - .expect(200); + describe('when sorting ', () => { + const ascendingOrder = ['my_doc_a', 'my_doc_b', 'my_doc_c']; + + it('allows sorting ascending', async () => { + const entries = await getEntries({ sortBy: 'title', sortDirection: 'asc' }); + expect(entries.map(({ id }) => id)).to.eql(ascendingOrder); + }); - expect(res.body.entries.length).to.eql(1); - expect(res.body.entries[0].id).to.eql('my_doc_a'); + it('allows sorting descending', async () => { + const entries = await getEntries({ sortBy: 'title', sortDirection: 'desc' }); + expect(entries.map(({ id }) => id)).to.eql([...ascendingOrder].reverse()); + }); + }); + + it('allows searching by title', async () => { + const entries = await getEntries({ query: 'b' }); + expect(entries.length).to.eql(1); + expect(entries[0].title).to.eql('My title b'); }); }); - describe('When the LLM creates entries', () => { - before(async () => { + describe('when importing categories', () => { + beforeEach(async () => { await clearKnowledgeBase(es); }); + afterEach(async () => { await clearKnowledgeBase(es); }); - it('can replace an existing entry using the `doc_id`', async () => { - await observabilityAIAssistantAPIClient + const importCategories = () => + observabilityAIAssistantAPIClient .editorUser({ - endpoint: 'POST /internal/observability_ai_assistant/functions/summarize', + endpoint: 'POST /internal/observability_ai_assistant/kb/entries/category/import', params: { body: { - doc_id: 'my_doc_id', - text: 'My content', - confidence: 'high', - is_correction: false, - public: false, - labels: {}, + category: 'my_new_category', + entries: [ + { + id: 'my_new_category_a', + texts: [ + 'My first category content a', + 'My second category content a', + 'my third category content a', + ], + }, + { + id: 'my_new_category_b', + texts: [ + 'My first category content b', + 'My second category content b', + 'my third category content b', + ], + }, + { + id: 'my_new_category_c', + texts: [ + 'My first category content c', + 'My second category content c', + 'my third category content c', + ], + }, + ], }, }, }) .expect(200); + it('overwrites existing entries on subsequent import', async () => { + await waitForModelReady(observabilityAIAssistantAPIClient, log); + await importCategories(); + await importCategories(); + + await pRetry( + async () => { + const res = await observabilityAIAssistantAPIClient + .editorUser({ + endpoint: 'GET /internal/observability_ai_assistant/kb/entries', + params: { + query: { + query: '', + sortBy: 'title', + sortDirection: 'asc', + }, + }, + }) + .expect(200); + + const categoryEntries = res.body.entries.filter( + (entry) => entry.labels?.category === 'my_new_category' + ); + + const entryGroups = uniq(categoryEntries.map((entry) => entry.doc_id)); + + log.debug( + `Waiting for entries to be created. Found ${categoryEntries.length} entries and ${entryGroups.length} groups` + ); + + if (categoryEntries.length !== 9 || entryGroups.length !== 3) { + throw new Error( + `Expected 9 entries, found ${categoryEntries.length} and ${entryGroups.length} groups` + ); + } + + expect(categoryEntries.length).to.eql(9); + expect(entryGroups.length).to.eql(3); + }, + { + retries: 100, + factor: 1, + } + ); + }); + }); + + describe('When the LLM creates entries', () => { + async function addEntryWithDocId({ + apiClient, + docId, + text, + }: { + apiClient: ObservabilityAIAssistantApiClient; + docId: string; + text: string; + }) { + return apiClient({ + endpoint: 'POST /internal/observability_ai_assistant/functions/summarize', + params: { + body: { + doc_id: docId, + text, + confidence: 'high', + is_correction: false, + public: false, + labels: {}, + }, + }, + }).expect(200); + } + + async function getEntriesWithDocId(docId: string) { const res = await observabilityAIAssistantAPIClient .editorUser({ endpoint: 'GET /internal/observability_ai_assistant/kb/entries', params: { query: { query: '', - sortBy: 'doc_id', + sortBy: 'title', sortDirection: 'asc', }, }, }) .expect(200); - const id = res.body.entries[0].id; - expect(res.body.entries.length).to.eql(1); - expect(res.body.entries[0].text).to.eql('My content'); + return res.body.entries.filter((entry) => entry.doc_id === docId); + } - await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'POST /internal/observability_ai_assistant/functions/summarize', - params: { - body: { - doc_id: 'my_doc_id', - text: 'My content_2', - confidence: 'high', - is_correction: false, - public: false, - labels: {}, - }, - }, - }) - .expect(200); + describe('when the LLM uses the same doc_id for two entries created by the same user', () => { + let entries1: KnowledgeBaseEntry[]; + let entries2: KnowledgeBaseEntry[]; - const res2 = await observabilityAIAssistantAPIClient - .editorUser({ - endpoint: 'GET /internal/observability_ai_assistant/kb/entries', - params: { - query: { - query: '', - sortBy: 'doc_id', - sortDirection: 'asc', + before(async () => { + const docId = 'my_favourite_color'; + + await addEntryWithDocId({ + apiClient: observabilityAIAssistantAPIClient.editorUser, + docId, + text: 'My favourite color is blue', + }); + entries1 = await getEntriesWithDocId(docId); + + await addEntryWithDocId({ + apiClient: observabilityAIAssistantAPIClient.editorUser, + docId, + text: 'My favourite color is green', + }); + entries2 = await getEntriesWithDocId(docId); + }); + + after(async () => { + await clearKnowledgeBase(es); + }); + + it('overwrites the first entry so there is only one', async () => { + expect(entries1.length).to.eql(1); + expect(entries2.length).to.eql(1); + }); + + it('replaces the text content of the first entry with the new text content', async () => { + expect(entries1[0].text).to.eql('My favourite color is blue'); + expect(entries2[0].text).to.eql('My favourite color is green'); + }); + + it('updates the timestamp', async () => { + const getAsMs = (timestamp: string) => new Date(timestamp).getTime(); + expect(getAsMs(entries1[0]['@timestamp'])).to.be.lessThan( + getAsMs(entries2[0]['@timestamp']) + ); + }); + + it('does not change the _id', () => { + expect(entries1[0].id).to.eql(entries2[0].id); + }); + }); + + describe('when the LLM uses same doc_id for two entries created by different users', () => { + let entries: KnowledgeBaseEntry[]; + + before(async () => { + await addEntryWithDocId({ + apiClient: observabilityAIAssistantAPIClient.editorUser, + docId: 'users_favorite_animal', + text: "The user's favourite animal is a dog", + }); + await addEntryWithDocId({ + apiClient: observabilityAIAssistantAPIClient.secondaryEditorUser, + docId: 'users_favorite_animal', + text: "The user's favourite animal is a cat", + }); + + const res = await observabilityAIAssistantAPIClient + .editorUser({ + endpoint: 'GET /internal/observability_ai_assistant/kb/entries', + params: { + query: { + query: '', + sortBy: 'title', + sortDirection: 'asc', + }, }, - }, - }) - .expect(200); + }) + .expect(200); + + entries = omitCategories(res.body.entries); + }); + + after(async () => { + await clearKnowledgeBase(es); + }); + + it('creates two separate entries with the same doc_id', async () => { + expect(entries.map(({ doc_id: docId }) => docId)).to.eql([ + 'users_favorite_animal', + 'users_favorite_animal', + ]); + }); + + it('creates two entries with different text content', async () => { + expect(entries.map(({ text }) => text)).to.eql([ + "The user's favourite animal is a cat", + "The user's favourite animal is a dog", + ]); + }); - expect(res2.body.entries.length).to.eql(1); - expect(res2.body.entries[0].text).to.eql('My content_2'); - expect(res2.body.entries[0].id).to.eql(id); + it('creates two entries by different users', async () => { + expect(entries.map(({ user }) => user?.name)).to.eql(['secondary_editor', 'editor']); + }); }); }); }); } + +function omitCategories(entries: KnowledgeBaseEntry[]) { + return entries.filter((entry) => entry.labels?.category === undefined); +} + +async function waitForModelReady( + observabilityAIAssistantAPIClient: ObservabilityAIAssistantApiClients, + log: ToolingLog +) { + return pRetry(async () => { + const res = await observabilityAIAssistantAPIClient + .editorUser({ endpoint: 'GET /internal/observability_ai_assistant/kb/status' }) + .expect(200); + + const isModelReady = res.body.ready; + log.debug(`Model status: ${isModelReady ? 'ready' : 'not ready'}`); + + if (!isModelReady) { + throw new Error('Model not ready'); + } + }); +} diff --git a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts index c1ce378be79c4..3083374c309a0 100644 --- a/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts +++ b/x-pack/test/observability_ai_assistant_api_integration/tests/knowledge_base/knowledge_base_user_instructions.spec.ts @@ -88,9 +88,9 @@ export default function ApiTest({ getService }: FtrProviderContext) { }); const instructions = res.body.userInstructions; - const sortByDocId = (data: any) => sortBy(data, 'doc_id'); - expect(sortByDocId(instructions)).to.eql( - sortByDocId([ + const sortById = (data: any) => sortBy(data, 'id'); + expect(sortById(instructions)).to.eql( + sortById([ { id: 'private-doc-from-editorUser', public: false, @@ -116,9 +116,9 @@ export default function ApiTest({ getService }: FtrProviderContext) { }); const instructions = res.body.userInstructions; - const sortByDocId = (data: any) => sortBy(data, 'doc_id'); - expect(sortByDocId(instructions)).to.eql( - sortByDocId([ + const sortById = (data: any) => sortBy(data, 'id'); + expect(sortById(instructions)).to.eql( + sortById([ { id: 'public-doc-from-editorUser', public: true,