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

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

@@ -81,9 +81,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';
@@ -96,12 +97,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(
     {
@@ -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,21 +63,31 @@ 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);
+
+      resources.logger.debug(
+        id
+          ? `Updating knowledge base entry with id: ${id}, doc_id: ${docId}`
+          : `Creating new knowledge base entry with doc_id: ${docId}`
+      );
+
       return client
         .addKnowledgeBaseEntry({
           entry: {
-            doc_id: id,
-            role: KnowledgeBaseEntryRole.AssistantSummarization,
-            id,
+            id: id ?? v4(),
+            doc_id: docId,
+            title: docId, // use doc_id as title for now
             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,
+      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),
     }),
@@ -142,22 +143,21 @@ 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,
         text,
         public: isPublic,
         labels,

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

@@ -13,12 +13,7 @@ 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 +103,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 },
     });
   },
 });
@@ -159,17 +144,19 @@ 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,
         labels: t.record(t.string, t.string),
         role: t.union([
-          t.literal('assistant_summarization'),
-          t.literal('user_entry'),
-          t.literal('elastic'),
+          t.literal(KnowledgeBaseEntryRole.AssistantSummarization),
+          t.literal(KnowledgeBaseEntryRole.UserEntry),
+          t.literal(KnowledgeBaseEntryRole.Elastic),
         ]),
       }),
     ]),
@@ -186,6 +173,8 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({
 
     const {
       id,
+      doc_id: docId,
+      title,
       text,
       public: isPublic,
       confidence,
@@ -196,15 +185,15 @@ 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',
         public: isPublic ?? true,
         labels: labels ?? {},
-        role: (role as KnowledgeBaseEntryRole) ?? KnowledgeBaseEntryRole.UserEntry,
+        role: role ?? KnowledgeBaseEntryRole.UserEntry,
       },
     });
   },
@@ -236,10 +225,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,
+          }),
+        ])
       ),
     }),
   }),
@@ -253,18 +247,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);
   },
 });
 
@@ -275,6 +299,7 @@ export const knowledgeBaseRoutes = {
   ...saveKnowledgeBaseUserInstruction,
   ...getKnowledgeBaseUserInstructions,
   ...importKnowledgeBaseEntries,
+  ...importKnowledgeBaseCategoryEntries,
   ...saveKnowledgeBaseEntry,
   ...deleteKnowledgeBaseEntry,
 };