diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts index febbc4156d66d..d255397d21c6c 100644 --- a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/context.ts @@ -19,6 +19,7 @@ import { FunctionVisibility, MessageRole, type Message } from '../../common/type import { concatenateChatCompletionChunks } from '../../common/utils/concatenate_chat_completion_chunks'; import type { ObservabilityAIAssistantClient } from '../service/client'; import { createFunctionResponseMessage } from '../service/util/create_function_response_message'; +import { parseSuggestionScores } from './parse_suggestion_scores'; const MAX_TOKEN_COUNT_FOR_DATA_ON_SCREEN = 1000; @@ -121,7 +122,7 @@ export function registerContextFunction({ }; } - const relevantDocuments = await scoreSuggestions({ + const { relevantDocuments, scores } = await scoreSuggestions({ suggestions, queries: queriesOrUserPrompt, messages, @@ -133,16 +134,21 @@ export function registerContextFunction({ return { content: { ...content, learnings: relevantDocuments as unknown as Serializable }, + data: { + scores, + suggestions, + }, }; } return new Observable((subscriber) => { getContext() - .then(({ content }) => { + .then(({ content, data }) => { subscriber.next( createFunctionResponseMessage({ name: 'context', content, + data, }) ); @@ -271,17 +277,16 @@ async function scoreSuggestions({ 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 }; + const scores = parseSuggestionScores(scoresAsString).map(({ index, score }) => { + return { + id: suggestions[index].id, + score, + }; }); if (scores.length === 0) { - return []; + // seemingly invalid or no scores, return all + return { relevantDocuments: suggestions, scores: [] }; } const suggestionIds = suggestions.map((document) => document.id); @@ -299,5 +304,5 @@ async function scoreSuggestions({ logger.debug(`Relevant documents: ${JSON.stringify(relevantDocuments, null, 2)}`); - return relevantDocuments; + return { relevantDocuments, scores }; } diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.test.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.test.ts new file mode 100644 index 0000000000000..7b62cf21af65b --- /dev/null +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.test.ts @@ -0,0 +1,69 @@ +/* + * 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 { parseSuggestionScores } from './parse_suggestion_scores'; + +describe('parseSuggestionScores', () => { + it('parses newlines as separators', () => { + expect( + parseSuggestionScores( + dedent( + `0,1 + 2,7 + 3,10` + ) + ) + ).toEqual([ + { + index: 0, + score: 1, + }, + { + index: 2, + score: 7, + }, + { + index: 3, + score: 10, + }, + ]); + }); + + it('parses semi-colons as separators', () => { + expect(parseSuggestionScores(`0,1;2,7;3,10`)).toEqual([ + { + index: 0, + score: 1, + }, + { + index: 2, + score: 7, + }, + { + index: 3, + score: 10, + }, + ]); + }); + + it('parses spaces as separators', () => { + expect(parseSuggestionScores(`0,1 2,7 3,10`)).toEqual([ + { + index: 0, + score: 1, + }, + { + index: 2, + score: 7, + }, + { + index: 3, + score: 10, + }, + ]); + }); +}); diff --git a/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.ts b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.ts new file mode 100644 index 0000000000000..9fa39bf1233b5 --- /dev/null +++ b/x-pack/plugins/observability_solution/observability_ai_assistant/server/functions/parse_suggestion_scores.ts @@ -0,0 +1,24 @@ +/* + * 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. + */ + +export function parseSuggestionScores(scoresAsString: string) { + // make sure that spaces, semi-colons etc work as separators as well + const scores = scoresAsString + .replace(/[^0-9,]/g, ' ') + .trim() + .split(/\s+/) + .map((pair) => { + const [index, score] = pair.split(',').map((str) => parseInt(str, 10)); + + return { + index, + score, + }; + }); + + return scores; +}