From 006e1d62385f3a879d7e207178cfdf16a580aea3 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 3 Aug 2023 13:34:25 +0200 Subject: [PATCH] [Observability AI Assistant]: Summarise & recall (#22) --- .../src/to_boolean_rt/index.ts | 2 +- .../common/types.ts | 10 ++ .../components/chat/chat_body.stories.tsx | 4 +- .../components/chat/chat_flyout.stories.tsx | 4 +- .../public/components/insight/insight.tsx | 1 + .../public/functions/elasticsearch.ts | 55 ++++++++ .../public/functions/index.ts | 34 +++++ .../public/functions/recall.ts | 49 +++++++ .../public/functions/setup_kb.ts | 38 +++++ .../public/functions/summarise.ts | 85 +++++++++++ .../public/hooks/use_timeline.ts | 57 +++++--- .../public/plugin.tsx | 53 +------ .../public/service/get_system_message.ts | 41 ++++++ .../public/utils/builders.ts | 3 +- .../get_timeline_items_from_conversation.ts | 24 ++-- .../server/index.ts | 9 +- .../server/routes/functions/route.ts | 84 +++++++++++ .../server/routes/runtime_types.ts | 2 + .../server/service/client/index.ts | 133 +++++++++++++++++- .../conversation_component_template.ts | 3 + .../server/service/index.ts | 76 +++++++++- .../server/service/kb_component_template.ts | 56 ++++++++ .../server/service/types.ts | 11 ++ 23 files changed, 740 insertions(+), 94 deletions(-) create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/index.ts create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/recall.ts create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/setup_kb.ts create mode 100644 x-pack/plugins/observability_ai_assistant/public/functions/summarise.ts create mode 100644 x-pack/plugins/observability_ai_assistant/public/service/get_system_message.ts create mode 100644 x-pack/plugins/observability_ai_assistant/server/service/kb_component_template.ts diff --git a/packages/kbn-io-ts-utils/src/to_boolean_rt/index.ts b/packages/kbn-io-ts-utils/src/to_boolean_rt/index.ts index eb8277adb281e..bf2ca188e24d6 100644 --- a/packages/kbn-io-ts-utils/src/to_boolean_rt/index.ts +++ b/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( +export const toBooleanRt = new t.Type( 'ToBoolean', t.boolean.is, (input) => { diff --git a/x-pack/plugins/observability_ai_assistant/common/types.ts b/x-pack/plugins/observability_ai_assistant/common/types.ts index 78b90d8bd551f..b98e255555904 100644 --- a/x-pack/plugins/observability_ai_assistant/common/types.ts +++ b/x-pack/plugins/observability_ai_assistant/common/types.ts @@ -49,6 +49,7 @@ export interface Conversation { labels: Record; numeric_labels: Record; namespace: string; + public: boolean; } export type ConversationRequestBase = Omit & { @@ -58,6 +59,15 @@ export type ConversationRequestBase = Omit; export interface ContextDefinition { diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx index 720c20d0f7b90..8eae1fb7df8ab 100644 --- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.stories.tsx +++ b/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: [ diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx index d13bdeb354409..dcd9121c532c1 100644 --- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx +++ b/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: () => {}, }; diff --git a/x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx b/x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx index 1fdbe758c4749..6b629da3fd5b0 100644 --- a/x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx +++ b/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]); diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts b/x-pack/plugins/observability_ai_assistant/public/functions/elasticsearch.ts new file mode 100644 index 0000000000000..154f1ac40c20a --- /dev/null +++ b/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 })); + } + ); +} diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/index.ts b/x-pack/plugins/observability_ai_assistant/public/functions/index.ts new file mode 100644 index 0000000000000..450793554f19e --- /dev/null +++ b/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 }); +} diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/public/functions/recall.ts new file mode 100644 index 0000000000000..576eba182c659 --- /dev/null +++ b/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 })); + } + ); +} diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/setup_kb.ts b/x-pack/plugins/observability_ai_assistant/public/functions/setup_kb.ts new file mode 100644 index 0000000000000..9cb498e1a7793 --- /dev/null +++ b/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 })); + } + ); +} diff --git a/x-pack/plugins/observability_ai_assistant/public/functions/summarise.ts b/x-pack/plugins/observability_ai_assistant/public/functions/summarise.ts new file mode 100644 index 0000000000000..723839fd6da6f --- /dev/null +++ b/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`, + }, + })); + } + ); +} diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts index 4fc2cb8093289..71ac00c98183e 100644 --- a/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts +++ b/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts @@ -15,16 +15,18 @@ import type { ChatTimelineProps } from '../components/chat/chat_timeline'; import type { ObservabilityAIAssistantService, PendingMessage } from '../types'; import { getTimelineItemsfromConversation } from '../utils/get_timeline_items_from_conversation'; import type { UseGenAIConnectorsResult } from './use_genai_connectors'; +import { getSystemMessage } from '../service/get_system_message'; export function createNewConversation(): ConversationCreateRequest { return { '@timestamp': new Date().toISOString(), - messages: [], + messages: [getSystemMessage()], conversation: { title: '', }, labels: {}, numeric_labels: {}, + public: false, }; } @@ -120,26 +122,39 @@ export function useTimeline({ if (nextMessage?.message.function_call?.name) { const name = nextMessage.message.function_call.name; - const message = await service.executeFunction( - name, - nextMessage.message.function_call.arguments, - controller.signal - ); - - await chat( - nextMessages.concat({ - '@timestamp': new Date().toISOString(), - message: { - role: MessageRole.System, - name, - content: `The following data was returned by the function: - \`\`\` - ${JSON.stringify(message.content, null, 4)} - \`\`\``, - data: JSON.stringify(message.data), - }, - }) - ); + try { + const message = await service.executeFunction( + name, + nextMessage.message.function_call.arguments, + controller.signal + ); + + await chat( + nextMessages.concat({ + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + name, + content: JSON.stringify(message.content), + data: JSON.stringify(message.data), + }, + }) + ); + } catch (error) { + await chat( + nextMessages.concat({ + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.User, + name, + content: JSON.stringify({ + message: error.toString(), + ...error.body, + }), + }, + }) + ); + } } }) .catch((err) => {}); diff --git a/x-pack/plugins/observability_ai_assistant/public/plugin.tsx b/x-pack/plugins/observability_ai_assistant/public/plugin.tsx index 949178a4b3649..3f4ac40998e85 100644 --- a/x-pack/plugins/observability_ai_assistant/public/plugin.tsx +++ b/x-pack/plugins/observability_ai_assistant/public/plugin.tsx @@ -4,19 +4,17 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import { EuiCodeBlock } from '@elastic/eui'; import { AppNavLinkStatus, - type CoreStart, DEFAULT_APP_CATEGORIES, type AppMountParameters, type CoreSetup, + type CoreStart, type Plugin, type PluginInitializerContext, } from '@kbn/core/public'; import { i18n } from '@kbn/i18n'; import type { Logger } from '@kbn/logging'; -import type { Serializable } from '@kbn/utility-types'; import React from 'react'; import ReactDOM from 'react-dom'; import type { @@ -25,6 +23,7 @@ import type { RegisterContextDefinition, RegisterFunctionDefinition, } from '../common/types'; +import { registerFunctions } from './functions'; import { createService } from './service/create_service'; import type { ConfigSchema, @@ -120,52 +119,12 @@ export class ObservabilityAIAssistantPlugin functionRegistry.set(def.name, { options: def, respond, render }); }; - registerContext({ - name: 'core', - description: - 'Core functions, like calling Elasticsearch APIs, storing embeddables for instructions or creating base visualisations.', + registerFunctions({ + registerContext, + registerFunction, + service, }); - 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 })); - }, - ({ response: { content } }) => { - return {JSON.stringify(content, null, 2)}; - } - ); - return { ...service, registerContext, diff --git a/x-pack/plugins/observability_ai_assistant/public/service/get_system_message.ts b/x-pack/plugins/observability_ai_assistant/public/service/get_system_message.ts new file mode 100644 index 0000000000000..816a6d68333b8 --- /dev/null +++ b/x-pack/plugins/observability_ai_assistant/public/service/get_system_message.ts @@ -0,0 +1,41 @@ +/* + * 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 dedent from 'dedent'; +import { MessageRole } from '../../common'; + +export function getSystemMessage() { + return { + '@timestamp': new Date().toISOString(), + message: { + role: MessageRole.System as const, + content: + dedent(`You are a helpful assistant for Elastic Observability. Your goal is to help the Elastic Observability users to quickly assess what is happening in their observed systems. You can help them visualise and analyze data, investigate their systems, perform root cause analysis or identify optimisation opportunities. + + You can use the "summarise" functions to store new information you have learned in a knowledge database. Once you have established that you did not know the answer to a question, and the user gave you this information, it's important that you create a summarisation of what you have learned and store it in the knowledge database. When you create this summarisation, make sure you craft it in a way that can be recalled with a semantic search later. + + Additionally, you can use the "recall" function to retrieve relevant information from the knowledge database, using semantic search. Use this when: + + - a user asks a question which you would normally not be able to answer, but you would if you had access to an knowledge database or history of past conversations. + - you are unsure about what the user specifically means. + + In these cases, use the "recall" function with a semantic search query to extract relevant information from the knowledge database. + + It is absolutely critical to use the "recall" function, even when the user has not asked for it. Using the "recall" function will allow you to create a much more tailored and improved user experience. Assume the user has told you before. + + You do have access to private data that the user has shared with you in previous conversations. + + Given this ability to store and recall data, it is very important to not make assumptions about the user's data or preferences, but rather, first query the knowledge database, and if nothing relevant comes up, ask the user for clarification. This is very important! They could be storing their data in any field, and in any data stream or index. + + You should autonomously execute these functions - do not wait on the user's permission, but be proactive. + + Note that any visualisations will be displayed ABOVE your textual response, not below. + + Feel free to use Markdown in your replies.`), + }, + }; +} diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts index 04f0b3efd19bb..597d28499f9d2 100644 --- a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts +++ b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts @@ -8,6 +8,7 @@ import { uniqueId } from 'lodash'; import { MessageRole, Conversation, FunctionDefinition } from '../../common/types'; import { ChatTimelineItem } from '../components/chat/chat_timeline'; +import { getSystemMessage } from '../service/get_system_message'; type ChatItemBuildProps = Partial & Pick; @@ -91,7 +92,7 @@ export function buildConversation(params?: Partial) { title: '', last_updated: '', }, - messages: [], + messages: [getSystemMessage()], labels: {}, numeric_labels: {}, namespace: '', diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.ts b/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.ts index 276a7f23dadee..8fc8261ede0a0 100644 --- a/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.ts +++ b/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.ts @@ -7,6 +7,7 @@ import { v4 } from 'uuid'; import { i18n } from '@kbn/i18n'; import type { AuthenticatedUser } from '@kbn/security-plugin/common'; +import dedent from 'dedent'; import { MessageRole } from '../../common'; import type { ConversationCreateRequest } from '../../common/types'; import type { ChatTimelineItem } from '../components/chat/chat_timeline'; @@ -38,16 +39,27 @@ export function getTimelineItemsfromConversation({ const isSystemPrompt = message.message.role === MessageRole.System; let title: string; + let content: string | undefined; + if (hasFunction) { title = i18n.translate('xpack.observabilityAiAssistant.suggestedFunctionEvent', { defaultMessage: 'suggested a function', }); + content = dedent(`I have requested your system performs the function _${ + message.message.function_call?.name + }_ with the payload + \`\`\` + ${JSON.stringify(JSON.parse(message.message.function_call?.arguments || ''), null, 4)} + \`\`\` + and return its results for me to look at.`); } else if (isSystemPrompt) { title = i18n.translate('xpack.observabilityAiAssistant.addedSystemPromptEvent', { - defaultMessage: 'returned data', + defaultMessage: 'added a prompt', }); + content = ''; } else { title = ''; + content = message.message.content; } const props = { @@ -58,15 +70,7 @@ export function getTimelineItemsfromConversation({ canGiveFeedback: message.message.role === MessageRole.Assistant, loading: false, title, - content: hasFunction - ? `I have requested your system performs the function _${ - message.message.function_call?.name - }_ with the payload - \`\`\` - ${JSON.stringify(JSON.parse(message.message.function_call?.arguments || ''), null, 4)} - \`\`\` - and return its results for me to look at.` - : message.message.content, + content, currentUser, }; diff --git a/x-pack/plugins/observability_ai_assistant/server/index.ts b/x-pack/plugins/observability_ai_assistant/server/index.ts index f471407678309..38a2caffca3e4 100644 --- a/x-pack/plugins/observability_ai_assistant/server/index.ts +++ b/x-pack/plugins/observability_ai_assistant/server/index.ts @@ -5,11 +5,18 @@ * 2.0. */ -import type { PluginInitializerContext } from '@kbn/core/server'; +import type { PluginConfigDescriptor, PluginInitializerContext } from '@kbn/core/server'; import type { ObservabilityAIAssistantConfig } from './config'; import { ObservabilityAIAssistantPlugin } from './plugin'; export type { ObservabilityAIAssistantServerRouteRepository } from './routes/get_global_observability_ai_assistant_route_repository'; +import { config as configSchema } from './config'; + +export const config: PluginConfigDescriptor = { + exposeToBrowser: {}, + schema: configSchema, +}; + export const plugin = (ctx: PluginInitializerContext) => new ObservabilityAIAssistantPlugin(ctx); diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts index b2469fe7622e9..e9b4426171f63 100644 --- a/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts +++ b/x-pack/plugins/observability_ai_assistant/server/routes/functions/route.ts @@ -5,7 +5,10 @@ * 2.0. */ import * as t from 'io-ts'; +import { nonEmptyStringRt, toBooleanRt } from '@kbn/io-ts-utils'; +import { notImplemented } from '@hapi/boom'; import { createObservabilityAIAssistantServerRoute } from '../create_observability_ai_assistant_server_route'; +import { KnowledgeBaseEntry } from '../../../common/types'; const functionElasticsearchRoute = createObservabilityAIAssistantServerRoute({ endpoint: 'POST /internal/observability_ai_assistant/functions/elasticsearch', @@ -44,6 +47,87 @@ const functionElasticsearchRoute = createObservabilityAIAssistantServerRoute({ }, }); +const functionRecallRoute = createObservabilityAIAssistantServerRoute({ + endpoint: 'POST /internal/observability_ai_assistant/functions/recall', + params: t.type({ + body: t.type({ + query: nonEmptyStringRt, + }), + }), + options: { + tags: ['access:ai_assistant'], + }, + handler: async (resources): Promise<{ entries: KnowledgeBaseEntry[] }> => { + const client = await resources.service.getClient({ request: resources.request }); + + if (!client) { + throw notImplemented(); + } + + return client.recall(resources.params.body.query); + }, +}); + +const functionSummariseRoute = createObservabilityAIAssistantServerRoute({ + endpoint: 'POST /internal/observability_ai_assistant/functions/summarise', + params: t.type({ + body: t.type({ + id: t.string, + text: nonEmptyStringRt, + confidence: t.union([t.literal('low'), t.literal('medium'), t.literal('high')]), + is_correction: toBooleanRt, + public: toBooleanRt, + }), + }), + options: { + tags: ['access:ai_assistant'], + }, + handler: async (resources): Promise => { + const client = await resources.service.getClient({ request: resources.request }); + + if (!client) { + throw notImplemented(); + } + + const { + confidence, + id, + is_correction: isCorrection, + text, + public: isPublic, + } = resources.params.body; + + return client.summarise({ + entry: { + confidence, + id, + is_correction: isCorrection, + text, + public: isPublic, + }, + }); + }, +}); + +const setupKnowledgeBaseRoute = createObservabilityAIAssistantServerRoute({ + endpoint: 'POST /internal/observability_ai_assistant/functions/setup_kb', + options: { + tags: ['access:ai_assistant'], + }, + handler: async (resources): Promise => { + const client = await resources.service.getClient({ request: resources.request }); + + if (!client) { + throw notImplemented(); + } + + await client.setupKnowledgeBase(); + }, +}); + export const functionRoutes = { ...functionElasticsearchRoute, + ...functionRecallRoute, + ...functionSummariseRoute, + ...setupKnowledgeBaseRoute, }; diff --git a/x-pack/plugins/observability_ai_assistant/server/routes/runtime_types.ts b/x-pack/plugins/observability_ai_assistant/server/routes/runtime_types.ts index 12b66ff988ac5..41d0d9d19492a 100644 --- a/x-pack/plugins/observability_ai_assistant/server/routes/runtime_types.ts +++ b/x-pack/plugins/observability_ai_assistant/server/routes/runtime_types.ts @@ -5,6 +5,7 @@ * 2.0. */ import * as t from 'io-ts'; +import { toBooleanRt } from '@kbn/io-ts-utils'; import { Conversation, ConversationCreateRequest, @@ -58,6 +59,7 @@ export const baseConversationRt: t.Type = t.type({ messages: t.array(messageRt), labels: t.record(t.string, t.string), numeric_labels: t.record(t.string, t.number), + public: toBooleanRt, }); export const conversationCreateRt: t.Type = t.intersection([ diff --git a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts index 50eef7dbf2377..bd582c18ad615 100644 --- a/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/plugins/observability_ai_assistant/server/service/client/index.ts @@ -4,8 +4,9 @@ * 2.0; you may not use this file except in compliance with the Elastic License * 2.0. */ -import type { SearchHit } from '@elastic/elasticsearch/lib/api/types'; -import { internal, notFound } from '@hapi/boom'; +import { errors } from '@elastic/elasticsearch'; +import type { SearchHit, QueryDslTextExpansionQuery } from '@elastic/elasticsearch/lib/api/types'; +import { internal, notFound, serverUnavailable } from '@hapi/boom'; import type { ActionsClient } from '@kbn/actions-plugin/server/actions_client'; import type { ElasticsearchClient } from '@kbn/core/server'; import type { Logger } from '@kbn/logging'; @@ -20,11 +21,12 @@ import type { } from 'openai'; import { v4 } from 'uuid'; import { - type FunctionDefinition, + KnowledgeBaseEntry, MessageRole, type Conversation, type ConversationCreateRequest, type ConversationUpdateRequest, + type FunctionDefinition, type Message, } from '../../../common/types'; import type { @@ -32,6 +34,12 @@ import type { ObservabilityAIAssistantResourceNames, } from '../types'; +const ELSER_MODEL_ID = '.elser_model_1'; + +function throwKnowledgeBaseNotReady(body: any) { + throw serverUnavailable(`Knowledge base is not ready yet`, body); +} + export class ObservabilityAIAssistantClient implements IObservabilityAIAssistantClient { constructor( private readonly dependencies: { @@ -53,8 +61,20 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant bool: { filter: [ { - term: { - 'user.name': this.dependencies.user.name, + bool: { + should: [ + { + term: { + 'user.name': this.dependencies.user.name, + }, + }, + { + term: { + public: true, + }, + }, + ], + minimum_should_match: 1, }, }, { @@ -152,7 +172,7 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant messages: messagesForOpenAI, stream: true, functions: functionsForOpenAI, - temperature: 0.1, + temperature: 0, }; const executeResult = await this.dependencies.actionsClient.execute({ @@ -235,4 +255,105 @@ export class ObservabilityAIAssistantClient implements IObservabilityAIAssistant return createdConversation; }; + + recall = async (query: string): Promise<{ entries: KnowledgeBaseEntry[] }> => { + try { + const response = await this.dependencies.esClient.search({ + index: this.dependencies.resources.aliases.kb, + query: { + bool: { + should: [ + { + text_expansion: { + 'ml.tokens': { + model_text: query, + model_id: '.elser_model_1', + }, + } as unknown as QueryDslTextExpansionQuery, + }, + ], + filter: [...this.getAccessQuery()], + }, + }, + _source: { + excludes: ['ml.tokens'], + }, + }); + + return { entries: response.hits.hits.map((hit) => hit._source!) }; + } catch (error) { + if ( + (error instanceof errors.ResponseError && + error.body.error.type === 'resource_not_found_exception') || + error.body.error.type === 'status_exception' + ) { + throwKnowledgeBaseNotReady(error.body); + } + throw error; + } + }; + + summarise = async ({ + entry: { id, ...document }, + }: { + entry: Omit; + }): Promise => { + try { + await this.dependencies.esClient.index({ + index: this.dependencies.resources.aliases.kb, + id, + document: { + '@timestamp': new Date().toISOString(), + ...document, + user: this.dependencies.user, + namespace: this.dependencies.namespace, + }, + pipeline: this.dependencies.resources.pipelines.kb, + }); + } catch (error) { + if (error instanceof errors.ResponseError && error.body.error.type === 'status_exception') { + throwKnowledgeBaseNotReady(error.body); + } + throw error; + } + }; + + setupKnowledgeBase = async () => { + // if this fails, it's fine to propagate the error to the user + await this.dependencies.esClient.ml.putTrainedModel({ + model_id: ELSER_MODEL_ID, + input: { + field_names: ['text_field'], + }, + }); + + try { + await this.dependencies.esClient.ml.startTrainedModelDeployment({ + model_id: ELSER_MODEL_ID, + }); + + const modelStats = await this.dependencies.esClient.ml.getTrainedModelsStats({ + model_id: ELSER_MODEL_ID, + }); + + const elserModelStats = modelStats.trained_model_stats[0]; + + if (elserModelStats?.deployment_stats?.state !== 'started') { + throwKnowledgeBaseNotReady({ + message: `Deployment has not started`, + deployment_stats: elserModelStats.deployment_stats, + }); + } + return; + } catch (error) { + if ( + (error instanceof errors.ResponseError && + error.body.error.type === 'resource_not_found_exception') || + error.body.error.type === 'status_exception' + ) { + throwKnowledgeBaseNotReady(error.body); + } + throw error; + } + }; } diff --git a/x-pack/plugins/observability_ai_assistant/server/service/conversation_component_template.ts b/x-pack/plugins/observability_ai_assistant/server/service/conversation_component_template.ts index 2ce8180b0fdc9..c00e2c5e3a1fb 100644 --- a/x-pack/plugins/observability_ai_assistant/server/service/conversation_component_template.ts +++ b/x-pack/plugins/observability_ai_assistant/server/service/conversation_component_template.ts @@ -87,6 +87,9 @@ export const conversationComponentTemplate: ClusterComponentTemplate['component_ }, }, }, + public: { + type: 'boolean', + }, }, }, }; diff --git a/x-pack/plugins/observability_ai_assistant/server/service/index.ts b/x-pack/plugins/observability_ai_assistant/server/service/index.ts index 5896a56bd634c..5a76bd9125797 100644 --- a/x-pack/plugins/observability_ai_assistant/server/service/index.ts +++ b/x-pack/plugins/observability_ai_assistant/server/service/index.ts @@ -14,6 +14,7 @@ import { getSpaceIdFromPath } from '@kbn/spaces-plugin/common'; import { once } from 'lodash'; import { ObservabilityAIAssistantClient } from './client'; import { conversationComponentTemplate } from './conversation_component_template'; +import { kbComponentTemplate } from './kb_component_template'; import type { IObservabilityAIAssistantClient, IObservabilityAIAssistantService, @@ -31,19 +32,26 @@ export class ObservabilityAIAssistantService implements IObservabilityAIAssistan private readonly resourceNames: ObservabilityAIAssistantResourceNames = { componentTemplate: { conversations: getResourceName('component-template-conversations'), + kb: getResourceName('component-template-kb'), }, aliases: { conversations: getResourceName('conversations'), + kb: getResourceName('kb'), }, indexPatterns: { conversations: getResourceName('conversations*'), + kb: getResourceName('kb*'), }, indexTemplate: { conversations: getResourceName('index-template-conversations'), + kb: getResourceName('index-template-kb'), }, ilmPolicy: { conversations: getResourceName('ilm-policy-conversations'), }, + pipelines: { + kb: getResourceName('kb-ingest-pipeline'), + }, }; constructor({ logger, core }: { logger: Logger; core: CoreSetup }) { @@ -94,20 +102,78 @@ export class ObservabilityAIAssistantService implements IObservabilityAIAssistan }, }); - const aliasName = this.resourceNames.aliases.conversations; + const conversationAliasName = this.resourceNames.aliases.conversations; await createConcreteWriteIndex({ esClient, logger: this.logger, totalFieldsLimit: 10000, indexPatterns: { - alias: aliasName, - pattern: `${aliasName}*`, - basePattern: `${aliasName}*`, - name: `${aliasName}-000001`, + alias: conversationAliasName, + pattern: `${conversationAliasName}*`, + basePattern: `${conversationAliasName}*`, + name: `${conversationAliasName}-000001`, template: this.resourceNames.indexTemplate.conversations, }, }); + + await esClient.cluster.putComponentTemplate({ + create: false, + name: this.resourceNames.componentTemplate.kb, + template: kbComponentTemplate, + }); + + await esClient.ingest.putPipeline({ + id: this.resourceNames.pipelines.kb, + processors: [ + { + inference: { + model_id: '.elser_model_1', + target_field: 'ml', + field_map: { + text: 'text_field', + }, + inference_config: { + // @ts-expect-error + text_expansion: { + results_field: 'tokens', + }, + }, + }, + }, + ], + }); + + await esClient.indices.putIndexTemplate({ + name: this.resourceNames.indexTemplate.kb, + composed_of: [this.resourceNames.componentTemplate.kb], + create: false, + index_patterns: [this.resourceNames.indexPatterns.kb], + template: { + settings: { + number_of_shards: 1, + auto_expand_replicas: '0-1', + refresh_interval: '1s', + }, + }, + }); + + const kbAliasName = this.resourceNames.aliases.kb; + + await createConcreteWriteIndex({ + esClient, + logger: this.logger, + totalFieldsLimit: 10000, + indexPatterns: { + alias: kbAliasName, + pattern: `${kbAliasName}*`, + basePattern: `${kbAliasName}*`, + name: `${kbAliasName}-000001`, + template: this.resourceNames.indexTemplate.kb, + }, + }); + + this.logger.info('Successfully set up index assets'); } catch (error) { this.logger.error(`Failed to initialize service: ${error.message}`); this.logger.debug(error); diff --git a/x-pack/plugins/observability_ai_assistant/server/service/kb_component_template.ts b/x-pack/plugins/observability_ai_assistant/server/service/kb_component_template.ts new file mode 100644 index 0000000000000..55d6bbd15519c --- /dev/null +++ b/x-pack/plugins/observability_ai_assistant/server/service/kb_component_template.ts @@ -0,0 +1,56 @@ +/* + * 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 { ClusterComponentTemplate } from '@elastic/elasticsearch/lib/api/typesWithBodyKey'; + +const keyword = { + type: 'keyword' as const, + ignore_above: 1024, +}; + +const text = { + type: 'text' as const, +}; + +const date = { + type: 'date' as const, +}; + +export const kbComponentTemplate: ClusterComponentTemplate['component_template']['template'] = { + mappings: { + dynamic: false, + properties: { + '@timestamp': date, + id: keyword, + user: { + properties: { + id: keyword, + name: keyword, + }, + }, + conversation: { + properties: { + id: keyword, + title: text, + last_updated: date, + }, + }, + namespace: keyword, + text, + 'ml.tokens': { + type: 'rank_features', + }, + confidence: keyword, + is_correction: { + type: 'boolean', + }, + public: { + type: 'boolean', + }, + }, + }, +}; diff --git a/x-pack/plugins/observability_ai_assistant/server/service/types.ts b/x-pack/plugins/observability_ai_assistant/server/service/types.ts index 56824d5506c18..27bdfb21de5c4 100644 --- a/x-pack/plugins/observability_ai_assistant/server/service/types.ts +++ b/x-pack/plugins/observability_ai_assistant/server/service/types.ts @@ -12,6 +12,7 @@ import type { ConversationCreateRequest, ConversationUpdateRequest, FunctionDefinition, + KnowledgeBaseEntry, Message, } from '../../common/types'; @@ -26,6 +27,9 @@ export interface IObservabilityAIAssistantClient { create: (conversation: ConversationCreateRequest) => Promise; update: (conversation: ConversationUpdateRequest) => Promise; delete: (conversationId: string) => Promise; + recall: (query: string) => Promise<{ entries: KnowledgeBaseEntry[] }>; + summarise: (options: { entry: Omit }) => Promise; + setupKnowledgeBase: () => Promise; } export interface IObservabilityAIAssistantService { @@ -37,17 +41,24 @@ export interface IObservabilityAIAssistantService { export interface ObservabilityAIAssistantResourceNames { componentTemplate: { conversations: string; + kb: string; }; indexTemplate: { conversations: string; + kb: string; }; ilmPolicy: { conversations: string; }; aliases: { conversations: string; + kb: string; }; indexPatterns: { conversations: string; + kb: string; + }; + pipelines: { + kb: string; }; }