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

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

@@ -79,9 +79,10 @@ export type ConversationUpdateRequest = ConversationRequestBase & {
 
 export interface KnowledgeBaseEntry {
   '@timestamp': string;
-  id: string;
+  id: string; // unique ID
+  doc_id?: string; // human readable ID generated by the LLM and used by the LLM to lookup and update existing entries. TODO: rename `doc_id` to `lookup_id`
+  title?: string;
   text: string;
-  doc_id: string;
   confidence: 'low' | 'medium' | 'high';
   is_correction: boolean;
   type?: 'user_instruction' | 'contextual';
@@ -94,12 +95,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/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,6 +5,7 @@
  * 2.0.
  */
 
+import { v4 } from 'uuid';
 import { KnowledgeBaseType } from '../../common/types';
 import type { FunctionRegistrationParameters } from '.';
 import { KnowledgeBaseEntryRole } from '../../common';
@@ -31,7 +32,7 @@ export function registerSummarizationFunction({
           id: {
             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 lookup id for the document. This should be a short human-readable keyword field with only alphabetic characters and underscores, that allow you to find and update it later.',
           },
           text: {
             type: 'string',
@@ -62,16 +63,21 @@ export function registerSummarizationFunction({
         ],
       },
     },
-    (
-      { arguments: { id, text, is_correction: isCorrection, confidence, public: isPublic } },
+    async (
+      { arguments: { id: docId, text, is_correction: isCorrection, confidence, public: isPublic } },
       signal
     ) => {
+      // The LLM should be able to update an existing entry by providing the same doc_id
+      // if no existing entry is found, we generate a uuid
+      const id = await client.getUuidFromDocId(docId);
+
       return client
         .addKnowledgeBaseEntry({
           entry: {
-            doc_id: id,
+            id: id ?? v4(),
+            title: docId, // use doc_id as title for now
+            doc_id: docId,
             role: KnowledgeBaseEntryRole.AssistantSummarization,
-            id,
             text,
             is_correction: isCorrection,
             type: KnowledgeBaseType.Contextual,

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,8 +7,9 @@
 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 { KnowledgeBaseEntryRole, KnowledgeBaseType } from '../../../common/types';
 import type { RecalledEntry } from '../../service/knowledge_base_service';
 import { getSystemMessageFromInstructions } from '../../service/util/get_system_message_from_instructions';
 import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route';
@@ -107,11 +108,10 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
   endpoint: 'POST /internal/observability_ai_assistant/functions/summarize',
   params: t.type({
     body: t.type({
-      id: t.string,
+      doc_id: 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),
     }),
@@ -127,22 +127,22 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({
     }
 
     const {
+      doc_id: docId,
       confidence,
-      id,
       is_correction: isCorrection,
-      type,
       text,
       public: isPublic,
       labels,
     } = resources.params.body;
 
+    const id = await client.getUuidFromDocId(docId);
     return client.addKnowledgeBaseEntry({
       entry: {
         confidence,
-        id,
-        doc_id: id,
+        id: id ?? v4(),
+        doc_id: docId,
         is_correction: isCorrection,
-        type,
+        type: KnowledgeBaseType.Contextual,
         text,
         public: isPublic,
         labels,

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

@@ -107,10 +107,9 @@ const saveKnowledgeBaseUserInstruction = createObservabilityAIAssistantServerRou
     }
 
     const { id, text, public: isPublic } = resources.params.body;
-    return client.addKnowledgeBaseEntry({
+    return client.addUserInstruction({
       entry: {
         id,
-        doc_id: id,
         text,
         public: isPublic,
         confidence: 'high',
@@ -158,9 +157,11 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
     body: t.intersection([
       t.type({
         id: t.string,
+        title: t.string,
         text: nonEmptyStringRt,
       }),
       t.partial({
+        doc_id: t.string,
         confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]),
         is_correction: toBooleanRt,
         public: toBooleanRt,
@@ -185,6 +186,8 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
 
     const {
       id,
+      doc_id: docId,
+      title,
       text,
       public: isPublic,
       confidence,
@@ -195,9 +198,10 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
 
     return client.addKnowledgeBaseEntry({
       entry: {
-        id,
         text,
-        doc_id: id,
+        id,
+        doc_id: docId,
+        title,
         confidence: confidence ?? 'high',
         is_correction: isCorrection ?? false,
         type: 'contextual',
@@ -235,10 +239,15 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
   params: t.type({
     body: t.type({
       entries: t.array(
-        t.type({
-          id: t.string,
-          text: nonEmptyStringRt,
-        })
+        t.intersection([
+          t.type({
+            id: t.string,
+            text: nonEmptyStringRt,
+          }),
+          t.partial({
+            title: t.string,
+          }),
+        ])
       ),
     }),
   }),
@@ -252,18 +261,48 @@ const importKnowledgeBaseEntries = createObservabilityAIAssistantServerRoute({
       throw notImplemented();
     }
 
-    const entries = resources.params.body.entries.map((entry) => ({
-      doc_id: entry.id,
+    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<void> => {
+    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);
   },
 });
 
@@ -274,6 +313,7 @@ export const knowledgeBaseRoutes = {
   ...saveKnowledgeBaseUserInstruction,
   ...getKnowledgeBaseUserInstructions,
   ...importKnowledgeBaseEntries,
+  ...importKnowledgeBaseCategoryEntries,
   ...saveKnowledgeBaseEntry,
   ...deleteKnowledgeBaseEntry,
 };

x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts

@@ -731,6 +731,40 @@ export class ObservabilityAIAssistantClient {
     return this.dependencies.knowledgeBaseService.setup();
   };
 
+  getUuidFromDocId = async (docId: string) => {
+    return this.dependencies.knowledgeBaseService.getUuidFromDocId({
+      namespace: this.dependencies.namespace,
+      user: this.dependencies.user,
+      docId,
+    });
+  };
+
+  addUserInstruction = async ({
+    entry,
+  }: {
+    entry: Omit<KnowledgeBaseEntry, '@timestamp'>;
+  }): Promise<void> => {
+    // for now we want to limit the number of user instructions to 1 per user
+    // if a user instruction already exists for the user, we get the id and update it
+    this.dependencies.logger.debug('Adding user instruction entry');
+    const existingId = await this.dependencies.knowledgeBaseService.getPersonalUserInstructionId({
+      isPublic: entry.public,
+      namespace: this.dependencies.namespace,
+      user: this.dependencies.user,
+    });
+
+    if (existingId) {
+      entry.id = existingId;
+      this.dependencies.logger.debug(`Updating user instruction with id "${existingId}"`);
+    }
+
+    return this.dependencies.knowledgeBaseService.addEntry({
+      namespace: this.dependencies.namespace,
+      user: this.dependencies.user,
+      entry,
+    });
+  };
+
   addKnowledgeBaseEntry = async ({
     entry,
   }: {
@@ -753,7 +787,7 @@ export class ObservabilityAIAssistantClient {
       document: { ...entry, '@timestamp': new Date().toISOString() },
     }));
 
-    await this.dependencies.knowledgeBaseService.addEntries({ operations });
+    await this.dependencies.knowledgeBaseService.importEntries({ operations });
   };
 
   getKnowledgeBaseEntries = async ({

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,
-              '@timestamp': new Date().toISOString(),
-              doc_id: entry.id,
-              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(