[8.13] [Obs AI Assistant] More lenient parsing of suggestion scores (#177898)

x-pack/plugins/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<MessageAddEvent>((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 };
 }

x-pack/plugins/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,
+      },
+    ]);
+  });
+});

x-pack/plugins/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;
+}