[Obs AI Assistant] Add uuid to knowledge base entries to avoid overwriting accidentally

x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts

@@ -82,8 +82,8 @@ export type ConversationUpdateRequest = ConversationRequestBase & {
 export interface KnowledgeBaseEntry {
   '@timestamp': string;
   id: string;
+  title?: string;
   text: string;
-  doc_id: string;
   confidence: 'low' | 'medium' | 'high';
   is_correction: boolean;
   type?: 'user_instruction' | 'contextual';
@@ -96,12 +96,12 @@ export interface KnowledgeBaseEntry {
 }
 
 export interface Instruction {
-  doc_id: string;
+  id: string;
   text: string;
 }
 
 export interface AdHocInstruction {
-  doc_id?: string;
+  id?: string;
   text: string;
   instruction_type: 'user_instruction' | 'application_instruction';
 }

x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/short_id_table.test.ts

@@ -7,6 +7,16 @@
 import { ShortIdTable } from './short_id_table';
 
 describe('shortIdTable', () => {
+  it('generates a short id from a uuid', () => {
+    const table = new ShortIdTable();
+
+    const uuid = 'd877f65c-4036-42c4-b105-19e2f1a1c045';
+    const shortId = table.take(uuid);
+
+    expect(shortId.length).toBe(4);
+    expect(table.lookup(shortId)).toBe(uuid);
+  });
+
   it('generates at least 10k unique ids consistently', () => {
     const ids = new Set();
 

x-pack/plugins/observability_solution/observability_ai_assistant/server/analytics/recall_ranking.ts

@@ -52,9 +52,9 @@ const schema: RootSchema<RecallRanking> = {
   },
 };
 
-export const RecallRankingEventType = 'observability_ai_assistant_recall_ranking';
+export const recallRankingEventType = 'observability_ai_assistant_recall_ranking';
 
 export const recallRankingEvent: EventTypeOpts<RecallRanking> = {
-  eventType: RecallRankingEventType,
+  eventType: recallRankingEventType,
   schema,
 };

x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { KnowledgeBaseType } from '../../common/types';
+import { v4 } from 'uuid';
 import type { FunctionRegistrationParameters } from '.';
 import { KnowledgeBaseEntryRole } from '../../common';
 
@@ -14,6 +14,7 @@ export const SUMMARIZE_FUNCTION_NAME = 'summarize';
 export function registerSummarizationFunction({
   client,
   functions,
+  resources,
 }: FunctionRegistrationParameters) {
   functions.registerFunction(
     {
@@ -28,10 +29,10 @@ export function registerSummarizationFunction({
       parameters: {
         type: 'object',
         properties: {
-          id: {
+          title: {
             type: 'string',
             description:
-              'An id for the document. This should be a short human-readable keyword field with only alphabetic characters and underscores, that allow you to update it later.',
+              'A human readable title that can be used to identify the document later. This should be no longer than 255 characters',
           },
           text: {
             type: 'string',
@@ -54,29 +55,31 @@ export function registerSummarizationFunction({
           },
         },
         required: [
-          'id' as const,
+          'title' as const,
           'text' as const,
           'is_correction' as const,
           'confidence' as const,
           'public' as const,
         ],
       },
     },
-    (
-      { arguments: { id, text, is_correction: isCorrection, confidence, public: isPublic } },
+    async (
+      { arguments: { title, text, is_correction: isCorrection, confidence, public: isPublic } },
       signal
     ) => {
+      const id = v4();
+      resources.logger.debug(`Creating new knowledge base entry with id: ${id}`);
+
       return client
         .addKnowledgeBaseEntry({
           entry: {
-            doc_id: id,
-            role: KnowledgeBaseEntryRole.AssistantSummarization,
             id,
+            title,
             text,
-            is_correction: isCorrection,
-            type: KnowledgeBaseType.Contextual,
-            confidence,
             public: isPublic,
+            role: KnowledgeBaseEntryRole.AssistantSummarization,
+            confidence,
+            is_correction: isCorrection,
             labels: {},
           },
           // signal,

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/chat/route.ts

@@ -42,7 +42,7 @@ const chatCompleteBaseRt = t.type({
       ]),
       instructions: t.array(
         t.intersection([
-          t.partial({ doc_id: t.string }),
+          t.partial({ id: t.string }),
           t.type({
             text: t.string,
             instruction_type: t.union([

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts

@@ -7,6 +7,7 @@
 import { notImplemented } from '@hapi/boom';
 import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils';
 import * as t from 'io-ts';
+import { v4 } from 'uuid';
 import { FunctionDefinition } from '../../../common/functions/types';
 import { KnowledgeBaseEntryRole } from '../../../common/types';
 import type { RecalledEntry } from '../../service/knowledge_base_service';
@@ -114,19 +115,19 @@ const functionRecallRoute = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    return client.recall({ queries, categories });
+    const entries = await client.recall({ queries, categories });
+    return { entries };
   },
 });
 
 const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/functions/summarize',
   params: t.type({
     body: t.type({
-      id: t.string,
+      title: t.string,
       text: nonEmptyStringRt,
       confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
       is_correction: toBooleanRt,
-      type: t.union([t.literal('user_instruction'), t.literal('contextual')]),
       public: toBooleanRt,
       labels: t.record(t.string, t.string),
     }),
@@ -142,22 +143,20 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
     }
 
     const {
+      title,
       confidence,
-      id,
       is_correction: isCorrection,
-      type,
       text,
       public: isPublic,
       labels,
     } = resources.params.body;
 
     return client.addKnowledgeBaseEntry({
       entry: {
+        title,
         confidence,
-        id,
-        doc_id: id,
+        id: v4(),
         is_correction: isCorrection,
-        type,
         text,
         public: isPublic,
         labels,

x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts

@@ -9,16 +9,12 @@ import type {
   MlDeploymentAllocationState,
   MlDeploymentState,
 } from '@elastic/elasticsearch/lib/api/types';
+import pLimit from 'p-limit';
 import { notImplemented } from '@hapi/boom';
 import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils';
 import * as t from 'io-ts';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
-import {
-  Instruction,
-  KnowledgeBaseEntry,
-  KnowledgeBaseEntryRole,
-  KnowledgeBaseType,
-} from '../../../common/types';
+import { Instruction, KnowledgeBaseEntry, KnowledgeBaseEntryRole } from '../../../common/types';
 
 const getKnowledgeBaseStatus = createObservabilityAIAssistantServerRoute({
   endpoint: 'GET /internal/observability_ai_assistant/kb/status',
@@ -108,18 +104,8 @@ const saveKnowledgeBaseUserInstruction = createObservabilityAIAssistantServerRou
     }
 
     const { id, text, public: isPublic } = resources.params.body;
-    return client.addKnowledgeBaseEntry({
-      entry: {
-        id,
-        doc_id: id,
-        text,
-        public: isPublic,
-        confidence: 'high',
-        is_correction: false,
-        type: KnowledgeBaseType.UserInstruction,
-        labels: {},
-        role: KnowledgeBaseEntryRole.UserEntry,
-      },
+    return client.addUserInstruction({
+      entry: { id, text, public: isPublic },
     });
   },
 });
@@ -153,26 +139,29 @@ const getKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
   },
 });
 
+const knowledgeBaseEntryRt = t.intersection([
+  t.type({
+    id: t.string,
+    title: t.string,
+    text: nonEmptyStringRt,
+  }),
+  t.partial({
+    confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
+    is_correction: toBooleanRt,
+    public: toBooleanRt,
+    labels: t.record(t.string, t.string),
+    role: t.union([
+      t.literal(KnowledgeBaseEntryRole.AssistantSummarization),
+      t.literal(KnowledgeBaseEntryRole.UserEntry),
+      t.literal(KnowledgeBaseEntryRole.Elastic),
+    ]),
+  }),
+]);
+
 const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/entries/save',
   params: t.type({
-    body: t.intersection([
-      t.type({
-        id: t.string,
-        text: nonEmptyStringRt,
-      }),
-      t.partial({
-        confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
-        is_correction: toBooleanRt,
-        public: toBooleanRt,
-        labels: t.record(t.string, t.string),
-        role: t.union([
-          t.literal('assistant_summarization'),
-          t.literal('user_entry'),
-          t.literal('elastic'),
-        ]),
-      }),
-    ]),
+    body: knowledgeBaseEntryRt,
   }),
   options: {
     tags: ['access:ai_assistant'],
@@ -184,27 +173,15 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    const {
-      id,
-      text,
-      public: isPublic,
-      confidence,
-      is_correction: isCorrection,
-      labels,
-      role,
-    } = resources.params.body;
-
+    const entry = resources.params.body;
     return client.addKnowledgeBaseEntry({
       entry: {
-        id,
-        text,
-        doc_id: id,
-        confidence: confidence ?? 'high',
-        is_correction: isCorrection ?? false,
-        type: 'contextual',
-        public: isPublic ?? true,
-        labels: labels ?? {},
-        role: (role as KnowledgeBaseEntryRole) ?? KnowledgeBaseEntryRole.UserEntry,
+        confidence: 'high',
+        is_correction: false,
+        public: true,
+        labels: {},
+        role: KnowledgeBaseEntryRole.UserEntry,
+        ...entry,
       },
     });
   },
@@ -235,12 +212,7 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/kb/entries/import',
   params: t.type({
     body: t.type({
-      entries: t.array(
-        t.type({
-          id: t.string,
-          text: nonEmptyStringRt,
-        })
-      ),
+      entries: t.array(knowledgeBaseEntryRt),
     }),
   }),
   options: {
@@ -253,18 +225,29 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    const entries = resources.params.body.entries.map((entry) => ({
-      doc_id: entry.id,
-      confidence: 'high' as KnowledgeBaseEntry['confidence'],
-      is_correction: false,
-      type: 'contextual' as const,
-      public: true,
-      labels: {},
-      role: KnowledgeBaseEntryRole.UserEntry,
-      ...entry,
-    }));
-
-    return await client.importKnowledgeBaseEntries({ entries });
+    const status = await client.getKnowledgeBaseStatus();
+    if (!status.ready) {
+      throw new Error('Knowledge base is not ready');
+    }
+
+    const limiter = pLimit(5);
+
+    const promises = resources.params.body.entries.map(async (entry) => {
+      return limiter(async () => {
+        return client.addKnowledgeBaseEntry({
+          entry: {
+            confidence: 'high',
+            is_correction: false,
+            public: true,
+            labels: {},
+            role: KnowledgeBaseEntryRole.UserEntry,
+            ...entry,
+          },
+        });
+      });
+    });
+
+    await Promise.all(promises);
   },
 });
 
@@ -273,8 +256,8 @@ export const knowledgeBaseRoutes = {
   ...getKnowledgeBaseStatus,
   ...getKnowledgeBaseEntries,
   ...saveKnowledgeBaseUserInstruction,
-  ...getKnowledgeBaseUserInstructions,
   ...importKnowledgeBaseEntries,
+  ...getKnowledgeBaseUserInstructions,
   ...saveKnowledgeBaseEntry,
   ...deleteKnowledgeBaseEntry,
 };