diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts index 68595e457a355..38b7460c29ae1 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts @@ -79,9 +79,9 @@ export type ConversationUpdateRequest = ConversationRequestBase & { export interface KnowledgeBaseEntry { '@timestamp': string; - id: string; + id: string; // this is a unique ID generated by the client + doc_id: string; // this is the human readable ID generated by the LLM and used by the LLM to update existing entries text: string; - doc_id: string; confidence: 'low' | 'medium' | 'high'; is_correction: boolean; type?: 'user_instruction' | 'contextual'; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/analytics/recall_ranking.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/analytics/recall_ranking.ts index 4c82f79fcba8d..4371310811edf 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/analytics/recall_ranking.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/analytics/recall_ranking.ts @@ -52,9 +52,9 @@ const schema: RootSchema = { }, }; -export const RecallRankingEventType = 'observability_ai_assistant_recall_ranking'; +export const recallRankingEventType = 'observability_ai_assistant_recall_ranking'; export const recallRankingEvent: EventTypeOpts = { - eventType: RecallRankingEventType, + eventType: recallRankingEventType, schema, }; diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts index 8865861d81f45..74f8e5730a38f 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/summarize.ts @@ -63,15 +63,14 @@ export function registerSummarizationFunction({ }, }, ( - { arguments: { id, text, is_correction: isCorrection, confidence, public: isPublic } }, + { arguments: { id: docId, text, is_correction: isCorrection, confidence, public: isPublic } }, signal ) => { return client .addKnowledgeBaseEntry({ entry: { - doc_id: id, + doc_id: docId, role: KnowledgeBaseEntryRole.AssistantSummarization, - id, text, is_correction: isCorrection, type: KnowledgeBaseType.Contextual, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts index fae7077953699..854aeca17833d 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/functions/route.ts @@ -139,7 +139,6 @@ const functionSummariseRoute = createObservabilityAIAssistantServerRoute({ return client.addKnowledgeBaseEntry({ entry: { confidence, - id, doc_id: id, is_correction: isCorrection, type, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts index 6bb024b913cde..c4b8390fbc0f5 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/routes/knowledge_base/route.ts @@ -109,7 +109,6 @@ const saveKnowledgeBaseUserInstruction = createObservabilityAIAssistantServerRou const { id, text, public: isPublic } = resources.params.body; return client.addKnowledgeBaseEntry({ entry: { - id, doc_id: id, text, public: isPublic, @@ -195,7 +194,6 @@ const saveKnowledgeBaseEntry = createObservabilityAIAssistantServerRoute({ return client.addKnowledgeBaseEntry({ entry: { - id, text, doc_id: id, confidence: confidence ?? 'high', diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts index 293d2da9c04b9..86b02e9b9c628 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/client/index.ts @@ -734,7 +734,7 @@ export class ObservabilityAIAssistantClient { addKnowledgeBaseEntry = async ({ entry, }: { - entry: Omit; + entry: Omit; }): Promise => { return this.dependencies.knowledgeBaseService.addEntry({ namespace: this.dependencies.namespace, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts index 679d59a57dbc8..a7b47a4d5a45f 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/service/knowledge_base_service/index.ts @@ -14,6 +14,7 @@ import pRetry from 'p-retry'; import { map, orderBy } from 'lodash'; import { encode } from 'gpt-tokenizer'; import { MlTrainedModelDeploymentNodesStats } from '@elastic/elasticsearch/lib/api/types'; +import { v4 } from 'uuid'; import { INDEX_QUEUED_DOCUMENTS_TASK_ID, INDEX_QUEUED_DOCUMENTS_TASK_TYPE, @@ -38,6 +39,7 @@ interface Dependencies { export interface RecalledEntry { id: string; + doc_id: string; text: string; score: number | null; is_correction?: boolean; @@ -368,13 +370,13 @@ export class KnowledgeBaseService { }; const response = await this.dependencies.esClient.asInternalUser.search< - Pick + Pick >({ index: [resourceNames.aliases.kb], query: esQuery, size: 20, _source: { - includes: ['text', 'is_correction', 'labels'], + includes: ['text', 'is_correction', 'labels', 'doc_id'], }, }); @@ -598,27 +600,68 @@ export class KnowledgeBaseService { return res.hits.hits[0]?._source?.doc_id; }; + getUuidFromHumanReadableId = async ({ + docId, + user, + namespace, + }: { + docId: string; + user?: { name: string; id?: string }; + namespace?: string; + }) => { + const query = { + bool: { + filter: [ + { term: { doc_id: docId } }, + + // exclude user instructions + { bool: { must_not: { term: { type: KnowledgeBaseType.UserInstruction } } } }, + + // restrict access to user's own entries + ...getAccessQuery({ user, namespace }), + ], + }, + }; + + const response = await this.dependencies.esClient.asInternalUser.search({ + size: 1, + index: resourceNames.aliases.kb, + query, + _source: false, + }); + + const id = response.hits.hits[0]?._id ?? v4(); + + return id; + }; + addEntry = async ({ - entry: { id, ...document }, + entry, user, namespace, }: { - entry: Omit; + entry: Omit; user?: { name: string; id?: string }; namespace?: string; }): Promise => { + let id = ''; + // for now we want to limit the number of user instructions to 1 per user - if (document.type === KnowledgeBaseType.UserInstruction) { + if (entry.type === KnowledgeBaseType.UserInstruction) { const existingId = await this.getExistingUserInstructionId({ - isPublic: document.public, + isPublic: entry.public, user, namespace, }); if (existingId) { id = existingId; - document.doc_id = existingId; + entry.doc_id = existingId; } + + // override previous id if it exists + } else { + id = await this.getUuidFromHumanReadableId({ docId: entry.doc_id, user, namespace }); } try { @@ -627,7 +670,7 @@ export class KnowledgeBaseService { id, document: { '@timestamp': new Date().toISOString(), - ...document, + ...entry, user, namespace, }, diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/recall_and_score.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/recall_and_score.ts index 8885ff7e1d7a2..a285ae85890c4 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/recall_and_score.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/recall_and_score.ts @@ -7,13 +7,18 @@ import type { Logger } from '@kbn/logging'; import { AnalyticsServiceStart } from '@kbn/core/server'; +import { scoreSuggestions } from './score_suggestions'; import type { Message } from '../../../common'; import type { ObservabilityAIAssistantClient } from '../../service/client'; import type { FunctionCallChatFunction } from '../../service/types'; -import { retrieveSuggestions } from './retrieve_suggestions'; -import { scoreSuggestions } from './score_suggestions'; -import type { RetrievedSuggestion } from './types'; -import { RecallRanking, RecallRankingEventType } from '../../analytics/recall_ranking'; +import { RecallRanking, recallRankingEventType } from '../../analytics/recall_ranking'; + +export interface RecalledSuggestion { + id: string; + docId: string; + text: string; + score: number | null; +} export async function recallAndScore({ recall, @@ -34,19 +39,24 @@ export async function recallAndScore({ logger: Logger; signal: AbortSignal; }): Promise<{ - relevantDocuments?: RetrievedSuggestion[]; + relevantDocuments?: RecalledSuggestion[]; scores?: Array<{ id: string; score: number }>; - suggestions: RetrievedSuggestion[]; + suggestions: RecalledSuggestion[]; }> { const queries = [ { text: userPrompt, boost: 3 }, { text: context, boost: 1 }, ].filter((query) => query.text.trim()); - const suggestions = await retrieveSuggestions({ - recall, - queries, - }); + const { entries: recalledEntries } = await recall({ queries }); + const suggestions: RecalledSuggestion[] = recalledEntries.map( + ({ id, doc_id: docId, text, score }) => ({ + id, + docId, + text, + score, + }) + ); if (!suggestions.length) { return { @@ -67,7 +77,7 @@ export async function recallAndScore({ chat, }); - analytics.reportEvent(RecallRankingEventType, { + analytics.reportEvent(recallRankingEventType, { prompt: queries.map((query) => query.text).join('\n\n'), scoredDocuments: suggestions.map((suggestion) => { const llmScore = scores.find((score) => score.id === suggestion.id); diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/retrieve_suggestions.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/retrieve_suggestions.ts deleted file mode 100644 index 3c680229cd5d2..0000000000000 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/retrieve_suggestions.ts +++ /dev/null @@ -1,24 +0,0 @@ -/* - * 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 { omit } from 'lodash'; -import { ObservabilityAIAssistantClient } from '../../service/client'; -import { RetrievedSuggestion } from './types'; - -export async function retrieveSuggestions({ - queries, - recall, -}: { - queries: Array<{ text: string; boost?: number }>; - recall: ObservabilityAIAssistantClient['recall']; -}): Promise { - const recallResponse = await recall({ - queries, - }); - - return recallResponse.entries.map((entry) => omit(entry, 'labels', 'is_correction')); -} diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts index 009b91a7a8c2c..1f35986bae8f0 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/score_suggestions.ts @@ -5,16 +5,14 @@ * 2.0. */ import * as t from 'io-ts'; -import { omit } from 'lodash'; import { Logger } from '@kbn/logging'; import dedent from 'dedent'; import { lastValueFrom } from 'rxjs'; import { decodeOrThrow, jsonRt } from '@kbn/io-ts-utils'; import { concatenateChatCompletionChunks, Message, MessageRole } from '../../../common'; import type { FunctionCallChatFunction } from '../../service/types'; -import type { RetrievedSuggestion } from './types'; import { parseSuggestionScores } from './parse_suggestion_scores'; -import { ShortIdTable } from '../../../common/utils/short_id_table'; +import { RecalledSuggestion } from './recall_and_score'; const scoreFunctionRequestRt = t.type({ message: t.type({ @@ -38,7 +36,7 @@ export async function scoreSuggestions({ signal, logger, }: { - suggestions: RetrievedSuggestion[]; + suggestions: RecalledSuggestion[]; messages: Message[]; userPrompt: string; context: string; @@ -46,17 +44,9 @@ export async function scoreSuggestions({ signal: AbortSignal; logger: Logger; }): Promise<{ - relevantDocuments: RetrievedSuggestion[]; + relevantDocuments: RecalledSuggestion[]; scores: Array<{ id: string; score: number }>; }> { - const shortIdTable = new ShortIdTable(); - - const suggestionsWithShortId = suggestions.map((suggestion) => ({ - ...omit(suggestion, 'score', 'id'), // To not bias the LLM - originalId: suggestion.id, - shortId: shortIdTable.take(suggestion.id), - })); - const newUserMessageContent = dedent(`Given the following question, score the documents that are relevant to the question. on a scale from 0 to 7, 0 being completely irrelevant, and 7 being extremely relevant. Information is relevant to the question if it helps in @@ -76,10 +66,7 @@ export async function scoreSuggestions({ Documents: ${JSON.stringify( - suggestionsWithShortId.map((suggestion) => ({ - id: suggestion.shortId, - content: suggestion.text, - })), + suggestions.map(({ id, docId: title, text }) => ({ id, title, text })), null, 2 )}`); @@ -127,15 +114,7 @@ export async function scoreSuggestions({ scoreFunctionRequest.message.function_call.arguments ); - const scores = parseSuggestionScores(scoresAsString).map(({ id, score }) => { - const originalSuggestion = suggestionsWithShortId.find( - (suggestion) => suggestion.shortId === id - ); - return { - originalId: originalSuggestion?.originalId, - score, - }; - }); + const scores = parseSuggestionScores(scoresAsString); if (scores.length === 0) { // seemingly invalid or no scores, return all @@ -145,11 +124,11 @@ export async function scoreSuggestions({ const suggestionIds = suggestions.map((document) => document.id); const relevantDocumentIds = scores - .filter((document) => suggestionIds.includes(document.originalId ?? '')) // Remove hallucinated documents + .filter((document) => suggestionIds.includes(document.id ?? '')) // Remove hallucinated documents .filter((document) => document.score > 4) .sort((a, b) => b.score - a.score) .slice(0, 5) - .map((document) => document.originalId); + .map((document) => document.id); const relevantDocuments = suggestions.filter((suggestion) => relevantDocumentIds.includes(suggestion.id) @@ -159,6 +138,6 @@ export async function scoreSuggestions({ return { relevantDocuments, - scores: scores.map((score) => ({ id: score.originalId!, score: score.score })), + scores: scores.map((score) => ({ id: score.id, score: score.score })), }; } diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/types.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/types.ts deleted file mode 100644 index 3774df64c1ee1..0000000000000 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/utils/recall/types.ts +++ /dev/null @@ -1,10 +0,0 @@ -/* - * 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 { RecalledEntry } from '../../service/knowledge_base_service'; - -export type RetrievedSuggestion = Omit;