[Observability AI Assistant]: Summarise & recall (elastic#22)

packages/kbn-io-ts-utils/src/to_boolean_rt/index.ts

@@ -8,7 +8,7 @@
 
 import * as t from 'io-ts';
 
-export const toBooleanRt = new t.Type<boolean, unknown, unknown>(
+export const toBooleanRt = new t.Type<boolean, boolean, unknown>(
   'ToBoolean',
   t.boolean.is,
   (input) => {

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

@@ -49,6 +49,7 @@ export interface Conversation {
   labels: Record<string, string>;
   numeric_labels: Record<string, number>;
   namespace: string;
+  public: boolean;
 }
 
 export type ConversationRequestBase = Omit<Conversation, 'user' | 'conversation' | 'namespace'> & {
@@ -58,6 +59,15 @@ export type ConversationRequestBase = Omit<Conversation, 'user' | 'conversation'
 export type ConversationCreateRequest = ConversationRequestBase;
 export type ConversationUpdateRequest = ConversationRequestBase & { conversation: { id: string } };
 
+export interface KnowledgeBaseEntry {
+  '@timestamp': string;
+  id: string;
+  text: string;
+  confidence: 'low' | 'medium' | 'high';
+  is_correction: boolean;
+  public: boolean;
+}
+
 type CompatibleJSONSchema = Exclude<JSONSchema, boolean>;
 
 export interface ContextDefinition {

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx

@@ -8,6 +8,7 @@
 import { ComponentStory } from '@storybook/react';
 import React from 'react';
 import { Observable } from 'rxjs';
+import { getSystemMessage } from '../../service/get_system_message';
 import { ObservabilityAIAssistantService } from '../../types';
 import { ChatBody as Component } from './chat_body';
 
@@ -34,7 +35,8 @@ const defaultProps: ChatBodyProps = {
     },
     labels: {},
     numeric_labels: {},
-    messages: [],
+    messages: [getSystemMessage()],
+    public: false,
   },
   connectors: {
     connectors: [

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx

@@ -7,6 +7,7 @@
 
 import { ComponentStory } from '@storybook/react';
 import React from 'react';
+import { getSystemMessage } from '../../service/get_system_message';
 import { KibanaReactStorybookDecorator } from '../../utils/storybook_decorator';
 import { ChatFlyout as Component } from './chat_flyout';
 
@@ -33,9 +34,10 @@ const defaultProps: ChatFlyoutProps = {
     conversation: {
       title: 'How is this working',
     },
-    messages: [],
+    messages: [getSystemMessage()],
     labels: {},
     numeric_labels: {},
+    public: false,
   },
   onClose: () => {},
 };

x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx

@@ -70,6 +70,7 @@ function ChatContent({ messages, connectorId }: { messages: Message[]; connector
       },
       labels: {},
       numeric_labels: {},
+      public: false,
     };
   }, [pendingMessage, messages]);
 

x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts

@@ -0,0 +1,55 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Serializable } from '@kbn/utility-types';
+import type { RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+
+export function registerElasticsearchFunction({
+  service,
+  registerFunction,
+}: {
+  service: ObservabilityAIAssistantService;
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'elasticsearch',
+      contexts: ['core'],
+      description: 'Call Elasticsearch APIs on behalf of the user',
+      parameters: {
+        type: 'object',
+        properties: {
+          method: {
+            type: 'string',
+            description: 'The HTTP method of the Elasticsearch endpoint',
+            enum: ['GET', 'PUT', 'POST', 'DELETE', 'PATCH'] as const,
+          },
+          path: {
+            type: 'string',
+            description: 'The path of the Elasticsearch endpoint, including query parameters',
+          },
+        },
+        required: ['method' as const, 'path' as const],
+      },
+    },
+    ({ arguments: { method, path, body } }, signal) => {
+      return service
+        .callApi(`POST /internal/observability_ai_assistant/functions/elasticsearch`, {
+          signal,
+          params: {
+            body: {
+              method,
+              path,
+              body,
+            },
+          },
+        })
+        .then((response) => ({ content: response as Serializable }));
+    }
+  );
+}

x-pack/plugins/observability_ai_assistant/public/functions/index.ts

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { RegisterContextDefinition, RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+import { registerElasticsearchFunction } from './elasticsearch';
+import { registerRecallFunction } from './recall';
+import { registerSetupKbFunction } from './setup_kb';
+import { registerSummarisationFunction } from './summarise';
+
+export function registerFunctions({
+  registerFunction,
+  registerContext,
+  service,
+}: {
+  registerFunction: RegisterFunctionDefinition;
+  registerContext: RegisterContextDefinition;
+  service: ObservabilityAIAssistantService;
+}) {
+  registerContext({
+    name: 'core',
+    description:
+      'Core functions, like calling Elasticsearch APIs, storing embeddables for instructions or creating base visualisations.',
+  });
+
+  registerElasticsearchFunction({ service, registerFunction });
+  registerSummarisationFunction({ service, registerFunction });
+  registerRecallFunction({ service, registerFunction });
+  registerSetupKbFunction({ service, registerFunction });
+}

x-pack/plugins/observability_ai_assistant/public/functions/recall.ts

@@ -0,0 +1,49 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Serializable } from '@kbn/utility-types';
+import type { RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+
+export function registerRecallFunction({
+  service,
+  registerFunction,
+}: {
+  service: ObservabilityAIAssistantService;
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'recall',
+      contexts: ['core'],
+      description:
+        'Use this function to recall earlier learnings. Anything you will summarise can be retrieved again later via this function.',
+      parameters: {
+        type: 'object',
+        properties: {
+          query: {
+            type: 'string',
+            description: 'The query for the semantic search',
+          },
+        },
+        required: ['query' as const],
+      },
+    },
+    ({ arguments: { query } }, signal) => {
+      return service
+        .callApi('POST /internal/observability_ai_assistant/functions/recall', {
+          params: {
+            body: {
+              query,
+            },
+          },
+          signal,
+        })
+        .then((response) => ({ content: response as unknown as Serializable }));
+    }
+  );
+}

x-pack/plugins/observability_ai_assistant/public/functions/setup_kb.ts

@@ -0,0 +1,38 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { Serializable } from '@kbn/utility-types';
+import type { RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+
+export function registerSetupKbFunction({
+  service,
+  registerFunction,
+}: {
+  service: ObservabilityAIAssistantService;
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'setup_kb',
+      contexts: ['core'],
+      description:
+        'Use this function to set up the knowledge base. ONLY use this if you got an error from the recall or summarise function, or if the user has explicitly requested it. Note that it might take a while (e.g. ten minutes) until the knowledge base is available. Assume it will not be ready for the rest of the current conversation.',
+      parameters: {
+        type: 'object',
+        properties: {},
+      },
+    },
+    ({}, signal) => {
+      return service
+        .callApi('POST /internal/observability_ai_assistant/functions/setup_kb', {
+          signal,
+        })
+        .then((response) => ({ content: response as unknown as Serializable }));
+    }
+  );
+}

x-pack/plugins/observability_ai_assistant/public/functions/summarise.ts

@@ -0,0 +1,85 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import type { RegisterFunctionDefinition } from '../../common/types';
+import type { ObservabilityAIAssistantService } from '../types';
+
+export function registerSummarisationFunction({
+  service,
+  registerFunction,
+}: {
+  service: ObservabilityAIAssistantService;
+  registerFunction: RegisterFunctionDefinition;
+}) {
+  registerFunction(
+    {
+      name: 'summarise',
+      contexts: ['core'],
+      description:
+        'Use this function to summarise things learned from the conversation. You can score the learnings with a confidence metric, whether it is a correction on a previous learning. An embedding will be created that you can recall later with a semantic search. There is no need to ask the user for permission to store something you have learned, unless you do not feel confident.',
+      parameters: {
+        type: 'object',
+        properties: {
+          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.',
+          },
+          text: {
+            type: 'string',
+            description:
+              'A human-readable summary of what you have learned, described in such a way that you can recall it later with semantic search.',
+          },
+          is_correction: {
+            type: 'boolean',
+            description: 'Whether this is a correction for a previous learning.',
+          },
+          confidence: {
+            type: 'string',
+            description: 'How confident you are about this being a correct and useful learning',
+            enum: ['low' as const, 'medium' as const, 'high' as const],
+          },
+          public: {
+            type: 'boolean',
+            description:
+              'Whether this information is specific to the user, or generally applicable to any user of the product',
+          },
+        },
+        required: [
+          'id' as const,
+          'text' as const,
+          'is_correction' as const,
+          'confidence' as const,
+          'public' as const,
+        ],
+      },
+    },
+    (
+      { arguments: { id, text, is_correction: isCorrection, confidence, public: isPublic } },
+      signal
+    ) => {
+      return service
+        .callApi('POST /internal/observability_ai_assistant/functions/summarise', {
+          params: {
+            body: {
+              id,
+              text,
+              is_correction: isCorrection,
+              confidence,
+              public: isPublic,
+            },
+          },
+          signal,
+        })
+        .then(() => ({
+          content: {
+            message: `The document has been stored`,
+          },
+        }));
+    }
+  );
+}