From 260ba9bf4a94e37393ab68aa4b240c1a927d284e Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Wed, 7 Feb 2024 15:34:46 +0100 Subject: [PATCH 1/9] Log response times --- .../server/service/client/index.ts | 44 ++++++++++++------- 1 file changed, 29 insertions(+), 15 deletions(-) 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 76749e75daed1..27f866a19607d 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 @@ -449,6 +449,8 @@ export class ObservabilityAIAssistantClient { ): Promise> => { const span = apm.startSpan(`chat ${name}`); + const spanId = (span?.ids['span.id'] || '').substring(0, 6); + const messagesForOpenAI: Array< Omit & { role: MessageRole; @@ -484,6 +486,8 @@ export class ObservabilityAIAssistantClient { this.dependencies.logger.debug(`Sending conversation to connector`); this.dependencies.logger.trace(JSON.stringify(request, null, 2)); + let now = performance.now(); + const executeResult = await this.dependencies.actionsClient.execute({ actionId: connectorId, params: { @@ -495,7 +499,13 @@ export class ObservabilityAIAssistantClient { }, }); - this.dependencies.logger.debug(`Received action client response: ${executeResult.status}`); + this.dependencies.logger.debug( + `Received action client response: ${executeResult.status} (took: ${Math.round( + performance.now() - now + )}ms)${spanId ? ` (${spanId})` : ''}` + ); + + now = performance.now(); if (executeResult.status === 'error' && executeResult?.serviceMessage) { const tokenLimitRegex = @@ -518,20 +528,24 @@ export class ObservabilityAIAssistantClient { const observable = streamIntoObservable(response).pipe(processOpenAiStream(), shareReplay()); - if (span) { - lastValueFrom(observable) - .then( - () => { - span.setOutcome('success'); - }, - () => { - span.setOutcome('failure'); - } - ) - .finally(() => { - span.end(); - }); - } + lastValueFrom(observable) + .then( + () => { + span?.setOutcome('success'); + }, + () => { + span?.setOutcome('failure'); + } + ) + .finally(() => { + this.dependencies.logger.debug( + `Completed response in ${Math.round(performance.now() - now)}ms${ + spanId ? ` (${spanId})` : '' + }` + ); + + span?.end(); + }); return observable; }; From 749f22bcd209c99099c4190a642ddcc6a13c27d4 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Wed, 7 Feb 2024 16:59:05 +0100 Subject: [PATCH 2/9] Count tokens when sending over request --- .../server/service/client/index.ts | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) 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 27f866a19607d..aa51c8571a3e7 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 @@ -14,7 +14,15 @@ import apm from 'elastic-apm-node'; import { decode, encode } from 'gpt-tokenizer'; import { compact, isEmpty, last, merge, noop, omit, pick, take } from 'lodash'; import type OpenAI from 'openai'; -import { filter, isObservable, lastValueFrom, Observable, shareReplay, toArray } from 'rxjs'; +import { + filter, + firstValueFrom, + isObservable, + lastValueFrom, + Observable, + shareReplay, + toArray, +} from 'rxjs'; import { Readable } from 'stream'; import { v4 } from 'uuid'; import { @@ -483,10 +491,12 @@ export class ObservabilityAIAssistantClient { function_call: functionCall ? { name: functionCall } : undefined, }; - this.dependencies.logger.debug(`Sending conversation to connector`); + this.dependencies.logger.debug( + `Sending conversation to connector (${encode(JSON.stringify(request)).length} tokens)` + ); this.dependencies.logger.trace(JSON.stringify(request, null, 2)); - let now = performance.now(); + const now = performance.now(); const executeResult = await this.dependencies.actionsClient.execute({ actionId: connectorId, @@ -505,8 +515,6 @@ export class ObservabilityAIAssistantClient { )}ms)${spanId ? ` (${spanId})` : ''}` ); - now = performance.now(); - if (executeResult.status === 'error' && executeResult?.serviceMessage) { const tokenLimitRegex = /This model's maximum context length is (\d+) tokens\. However, your messages resulted in (\d+) tokens/g; @@ -528,6 +536,19 @@ export class ObservabilityAIAssistantClient { const observable = streamIntoObservable(response).pipe(processOpenAiStream(), shareReplay()); + firstValueFrom(observable) + .then( + () => {}, + () => {} + ) + .finally(() => { + this.dependencies.logger.debug( + `Received first value after ${Math.round(performance.now() - now)}ms${ + spanId ? ` (${spanId})` : '' + }` + ); + }); + lastValueFrom(observable) .then( () => { From 46d82b52c68d60c532510080e58a57327452cf46 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Wed, 7 Feb 2024 18:06:19 +0100 Subject: [PATCH 3/9] Use CSV output format for recall --- .../server/functions/recall.ts | 28 ++++++++++++------- .../server/service/client/index.ts | 4 +-- 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts index 7e966fa0e5508..1b5785ea1cf73 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts +++ b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts @@ -9,7 +9,7 @@ import { decodeOrThrow, jsonRt } from '@kbn/io-ts-utils'; import type { Serializable } from '@kbn/utility-types'; import dedent from 'dedent'; import * as t from 'io-ts'; -import { last, omit } from 'lodash'; +import { chunk, last, omit } from 'lodash'; import { lastValueFrom } from 'rxjs'; import { FunctionRegistrationParameters } from '.'; import { MessageRole, type Message } from '../../common/types'; @@ -105,15 +105,23 @@ export function registerRecallFunction({ }; } - const relevantDocuments = await scoreSuggestions({ - suggestions, - systemMessage, - userMessage, - queries, - client, - connectorId, - signal, - }); + const chunks = chunk(suggestions, 1); + + const relevantDocuments = ( + await Promise.all( + chunks.map((suggestionsInChunk) => + scoreSuggestions({ + suggestions: suggestionsInChunk, + systemMessage, + userMessage, + queries, + client, + connectorId, + signal, + }) + ) + ) + ).flat(); resources.logger.debug(`Received ${relevantDocuments.length} relevant documents`); resources.logger.debug(JSON.stringify(relevantDocuments, null, 2)); 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 aa51c8571a3e7..da29c472520b9 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 @@ -491,9 +491,7 @@ export class ObservabilityAIAssistantClient { function_call: functionCall ? { name: functionCall } : undefined, }; - this.dependencies.logger.debug( - `Sending conversation to connector (${encode(JSON.stringify(request)).length} tokens)` - ); + this.dependencies.logger.debug(`Sending conversation to connector`); this.dependencies.logger.trace(JSON.stringify(request, null, 2)); const now = performance.now(); From 80877b7006e112cef158587fdcce4a16a7b4bf16 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 08:35:32 +0100 Subject: [PATCH 4/9] Don't log (incorrect) token count --- .../observability_ai_assistant/server/service/client/index.ts | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 aa51c8571a3e7..da29c472520b9 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 @@ -491,9 +491,7 @@ export class ObservabilityAIAssistantClient { function_call: functionCall ? { name: functionCall } : undefined, }; - this.dependencies.logger.debug( - `Sending conversation to connector (${encode(JSON.stringify(request)).length} tokens)` - ); + this.dependencies.logger.debug(`Sending conversation to connector`); this.dependencies.logger.trace(JSON.stringify(request, null, 2)); const now = performance.now(); From 8db2c428ed199b9d25886ba5e8a64fc848025a98 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 08:36:43 +0100 Subject: [PATCH 5/9] Use CSV --- .../server/functions/recall.ts | 106 ++++++++---------- 1 file changed, 47 insertions(+), 59 deletions(-) diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts index 7e966fa0e5508..331e63d3c49bc 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts +++ b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts @@ -9,7 +9,7 @@ import { decodeOrThrow, jsonRt } from '@kbn/io-ts-utils'; import type { Serializable } from '@kbn/utility-types'; import dedent from 'dedent'; import * as t from 'io-ts'; -import { last, omit } from 'lodash'; +import { compact, last, omit } from 'lodash'; import { lastValueFrom } from 'rxjs'; import { FunctionRegistrationParameters } from '.'; import { MessageRole, type Message } from '../../common/types'; @@ -87,12 +87,17 @@ export function registerRecallFunction({ messages.filter((message) => message.message.role === MessageRole.User) ); + const nonEmptyQueries = queries.filter(Boolean); + + const queriesOrUserPrompt = nonEmptyQueries.length + ? nonEmptyQueries + : compact([userMessage?.message.content]); + const suggestions = await retrieveSuggestions({ userMessage, client, - signal, categories, - queries, + queries: queriesOrUserPrompt, }); resources.logger.debug(`Received ${suggestions.length} suggestions`); @@ -107,9 +112,8 @@ export function registerRecallFunction({ const relevantDocuments = await scoreSuggestions({ suggestions, - systemMessage, - userMessage, - queries, + queries: queriesOrUserPrompt, + messages, client, connectorId, signal, @@ -126,25 +130,17 @@ export function registerRecallFunction({ } async function retrieveSuggestions({ - userMessage, queries, client, categories, - signal, }: { userMessage?: Message; queries: string[]; client: ObservabilityAIAssistantClient; categories: Array<'apm' | 'lens'>; - signal: AbortSignal; }) { - const queriesWithUserPrompt = - userMessage && userMessage.message.content - ? [userMessage.message.content, ...queries] - : queries; - const recallResponse = await client.recall({ - queries: queriesWithUserPrompt, + queries, categories, }); @@ -161,50 +157,42 @@ const scoreFunctionRequestRt = t.type({ }); const scoreFunctionArgumentsRt = t.type({ - scores: t.array( - t.type({ - id: t.string, - score: t.number, - }) - ), + scores: t.string, }); async function scoreSuggestions({ suggestions, - systemMessage, - userMessage, + messages, queries, client, connectorId, signal, }: { suggestions: Awaited>; - systemMessage: Message; - userMessage?: Message; + messages: Message[]; queries: string[]; client: ObservabilityAIAssistantClient; connectorId: string; signal: AbortSignal; }) { - const systemMessageExtension = - dedent(`You have the function called score available to help you inform the user about how relevant you think a given document is to the conversation. - Please give a score between 1 and 7, fractions are allowed. - A higher score means it is more relevant.`); - const extendedSystemMessage = { - ...systemMessage, - message: { - ...systemMessage.message, - content: `${systemMessage.message.content}\n\n${systemMessageExtension}`, - }, - }; - - const userMessageOrQueries = - userMessage && userMessage.message.content ? userMessage.message.content : queries.join(','); + const indexedSuggestions = suggestions.map((suggestion, index) => ({ ...suggestion, id: index })); const newUserMessageContent = - dedent(`Given the question "${userMessageOrQueries}", can you give me a score for how relevant the following documents are? + dedent(`Given the following question, score the documents that are relevant to the question. on a scale from 0 to 7, + 0 being completely relevant, and 10 being extremely relevant. Information is relevant to the question if it helps in + answering the question. Judge it according to the following criteria: + + - The document is relevant to the question, and the rest of the conversation + - The document has information relevant to the question that is not mentioned, + or more detailed than what is available in the conversation + - The document has a high amount of information relevant to the question compared to other documents + - The document contains new information not mentioned before in the conversation - ${JSON.stringify(suggestions, null, 2)}`); + Question: + ${queries.join('\n')} + + Documents: + ${JSON.stringify(indexedSuggestions, null, 2)}`); const newUserMessage: Message = { '@timestamp': new Date().toISOString(), @@ -223,22 +211,13 @@ async function scoreSuggestions({ additionalProperties: false, properties: { scores: { - description: 'The document IDs and their scores', - type: 'array', - items: { - type: 'object', - additionalProperties: false, - properties: { - id: { - description: 'The ID of the document', - type: 'string', - }, - score: { - description: 'The score for the document', - type: 'number', - }, - }, - }, + description: `The document IDs and their scores, as CSV. Example: + + my_id,7 + my_other_id,3 + my_third_id,4 + `, + type: 'string', }, }, required: ['score'], @@ -250,7 +229,7 @@ async function scoreSuggestions({ ( await client.chat('score_suggestions', { connectorId, - messages: [extendedSystemMessage, newUserMessage], + messages: [...messages.slice(-1), newUserMessage], functions: [scoreFunction], functionCall: 'score', signal, @@ -258,10 +237,19 @@ async function scoreSuggestions({ ).pipe(concatenateChatCompletionChunks()) ); const scoreFunctionRequest = decodeOrThrow(scoreFunctionRequestRt)(response); - const { scores } = decodeOrThrow(jsonRt.pipe(scoreFunctionArgumentsRt))( + const { scores: scoresAsString } = decodeOrThrow(jsonRt.pipe(scoreFunctionArgumentsRt))( scoreFunctionRequest.message.function_call.arguments ); + const scores = scoresAsString.split('\n').map((line) => { + const [index, score] = line + .split(',') + .map((value) => value.trim()) + .map(Number); + + return { id: suggestions[index].id, score }; + }); + if (scores.length === 0) { return []; } From 4ee18470c928411d5aec262cb45d8e9b3b7b724e Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 14:01:10 +0100 Subject: [PATCH 6/9] Update x-pack/plugins/observability_ai_assistant/server/functions/recall.ts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Søren Louv-Jansen --- .../observability_ai_assistant/server/functions/recall.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts index 331e63d3c49bc..8b58c60f95a21 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts +++ b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts @@ -87,7 +87,7 @@ export function registerRecallFunction({ messages.filter((message) => message.message.role === MessageRole.User) ); - const nonEmptyQueries = queries.filter(Boolean); + const nonEmptyQueries = compact(queries); const queriesOrUserPrompt = nonEmptyQueries.length ? nonEmptyQueries From 394e3ce736f8c189578f03093957a3aafb1c9181 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 14:22:24 +0100 Subject: [PATCH 7/9] Review feedback --- .../observability_ai_assistant/server/functions/recall.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts index 331e63d3c49bc..e6c8b0cb71a5f 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts +++ b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts @@ -179,7 +179,7 @@ async function scoreSuggestions({ 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 relevant, and 10 being extremely relevant. Information is relevant to the question if it helps in + 0 being completely relevant, and 7 being extremely relevant. Information is relevant to the question if it helps in answering the question. Judge it according to the following criteria: - The document is relevant to the question, and the rest of the conversation @@ -229,7 +229,7 @@ async function scoreSuggestions({ ( await client.chat('score_suggestions', { connectorId, - messages: [...messages.slice(-1), newUserMessage], + messages: [...messages.slice(0, -1), newUserMessage], functions: [scoreFunction], functionCall: 'score', signal, From 9bfbc32e0cb2581fddacf2eb63bdc31219c07699 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 14:31:38 +0100 Subject: [PATCH 8/9] Review feedback --- .../server/service/client/index.ts | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) 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 f28d68ed0757b..afd34aa8ea966 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 @@ -541,10 +541,7 @@ export class ObservabilityAIAssistantClient { const observable = streamIntoObservable(response).pipe(processOpenAiStream(), shareReplay()); firstValueFrom(observable) - .then( - () => {}, - () => {} - ) + .catch(noop) .finally(() => { this.dependencies.logger.debug( `Received first value after ${Math.round(performance.now() - now)}ms${ From 1b36de7a5433c4b8d30404e1f05fdebf20fe27d2 Mon Sep 17 00:00:00 2001 From: Dario Gieselaar Date: Thu, 8 Feb 2024 15:18:22 +0100 Subject: [PATCH 9/9] Resolve merge conflict --- .../observability_ai_assistant/server/functions/recall.ts | 4 ---- 1 file changed, 4 deletions(-) diff --git a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts index 465db7f0b30cb..909a823286cc6 100644 --- a/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts +++ b/x-pack/plugins/observability_ai_assistant/server/functions/recall.ts @@ -238,7 +238,6 @@ async function scoreSuggestions({ scoreFunctionRequest.message.function_call.arguments ); -<<<<<<< HEAD const scores = scoresAsString.split('\n').map((line) => { const [index, score] = line .split(',') @@ -247,9 +246,6 @@ async function scoreSuggestions({ return { id: suggestions[index].id, score }; }); -======= - resources.logger.debug(`Scores: ${JSON.stringify(scores, null, 2)}`); ->>>>>>> 1c3fa24be396176454df3dd3a67c7acdfcea46d4 if (scores.length === 0) { return [];