[Obs AI Assistant] Extract get_data_on_screen function (#214362)

Minor refactor to extract get registration of `get_data_on_screen`. This
way the function co-located with the other LLM functions making it
easier to find and test.

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/get_data_on_screen.ts

@@ -0,0 +1,61 @@
+/*
+ * 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 { compact } from 'lodash';
+import dedent from 'dedent';
+import { ObservabilityAIAssistantScreenContextRequest } from '../../common/types';
+import { FunctionVisibility } from '../../common';
+import { ChatFunctionClient } from '../service/chat_function_client';
+
+export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
+
+export function registerGetDataOnScreenFunction(
+  functions: ChatFunctionClient,
+  screenContexts: ObservabilityAIAssistantScreenContextRequest[]
+) {
+  const allData = compact(screenContexts.flatMap((context) => context.data));
+
+  if (!allData.length) {
+    return;
+  }
+
+  functions.registerFunction(
+    {
+      name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
+      description: `Retrieve the structured data of content currently visible on the user's screen. Use this tool to understand what the user is viewing at this moment to provide more accurate and context-aware responses to their questions.`,
+      visibility: FunctionVisibility.AssistantOnly,
+      parameters: {
+        type: 'object',
+        properties: {
+          data: {
+            type: 'array',
+            description:
+              'The pieces of data you want to look at it. You can request one, or multiple',
+            items: {
+              type: 'string',
+              enum: allData.map((data) => data.name),
+            },
+          },
+        },
+        required: ['data' as const],
+      },
+    },
+    async ({ arguments: { data: dataNames } }) => {
+      return {
+        content: allData.filter((data) => dataNames.includes(data.name)),
+      };
+    }
+  );
+
+  functions.registerInstruction(({ availableFunctionNames }) =>
+    availableFunctionNames.includes(GET_DATA_ON_SCREEN_FUNCTION_NAME)
+      ? `The ${GET_DATA_ON_SCREEN_FUNCTION_NAME} function will retrieve specific content from the user's screen by specifying a data key. Use this tool to provide context-aware responses. Available data: ${dedent(
+          allData.map((data) => `${data.name}: ${data.description}`).join('\n')
+        )}`
+      : undefined
+  );
+}

x-pack/platform/plugins/shared/observability_ai_assistant/server/functions/index.ts

@@ -13,7 +13,7 @@ import { registerElasticsearchFunction } from './elasticsearch';
 import { GET_DATASET_INFO_FUNCTION_NAME, registerGetDatasetInfoFunction } from './get_dataset_info';
 import { registerKibanaFunction } from './kibana';
 import { registerExecuteConnectorFunction } from './execute_connector';
-import { GET_DATA_ON_SCREEN_FUNCTION_NAME } from '../service/chat_function_client';
+import { GET_DATA_ON_SCREEN_FUNCTION_NAME } from './get_data_on_screen';
 
 // cannot be imported from x-pack/solutions/observability/plugins/observability_ai_assistant_app/server/functions/query/index.ts due to circular dependency
 export const QUERY_FUNCTION_NAME = 'query';

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/chat_function_client/index.test.ts

@@ -5,10 +5,11 @@
  * 2.0.
  */
 import dedent from 'dedent';
-import { ChatFunctionClient, GET_DATA_ON_SCREEN_FUNCTION_NAME } from '.';
+import { ChatFunctionClient } from '.';
 import { FunctionVisibility } from '../../../common/functions/types';
 import { Logger } from '@kbn/logging';
 import { RegisterInstructionCallback } from '../types';
+import { GET_DATA_ON_SCREEN_FUNCTION_NAME } from '../../functions/get_data_on_screen';
 
 describe('chatFunctionClient', () => {
   describe('when executing a function with invalid arguments', () => {

x-pack/platform/plugins/shared/observability_ai_assistant/server/service/chat_function_client/index.ts

@@ -7,10 +7,9 @@
 /* eslint-disable max-classes-per-file*/
 
 import Ajv, { type ErrorObject, type ValidateFunction } from 'ajv';
-import dedent from 'dedent';
 import { compact, keyBy } from 'lodash';
 import { Logger } from '@kbn/logging';
-import { FunctionVisibility, type FunctionResponse } from '../../../common/functions/types';
+import { type FunctionResponse } from '../../../common/functions/types';
 import type { Message, ObservabilityAIAssistantScreenContextRequest } from '../../../common/types';
 import { filterFunctionDefinitions } from '../../../common/utils/filter_function_definitions';
 import type {
@@ -21,6 +20,7 @@ import type {
   RegisterFunction,
   RegisterInstruction,
 } from '../types';
+import { registerGetDataOnScreenFunction } from '../../functions/get_data_on_screen';
 
 export class FunctionArgsValidationError extends Error {
   constructor(public readonly errors: ErrorObject[]) {
@@ -32,8 +32,6 @@ const ajv = new Ajv({
   strict: false,
 });
 
-export const GET_DATA_ON_SCREEN_FUNCTION_NAME = 'get_data_on_screen';
-
 export class ChatFunctionClient {
   private readonly instructions: InstructionOrCallback[] = [];
 
@@ -43,47 +41,9 @@ export class ChatFunctionClient {
   private readonly actions: Required<ObservabilityAIAssistantScreenContextRequest>['actions'];
 
   constructor(private readonly screenContexts: ObservabilityAIAssistantScreenContextRequest[]) {
-    const allData = compact(screenContexts.flatMap((context) => context.data));
-
     this.actions = compact(screenContexts.flatMap((context) => context.actions));
 
-    if (allData.length) {
-      this.registerFunction(
-        {
-          name: GET_DATA_ON_SCREEN_FUNCTION_NAME,
-          description: `Retrieve the structured data of content currently visible on the user's screen. Use this tool to understand what the user is viewing at this moment to provide more accurate and context-aware responses to their questions.`,
-          visibility: FunctionVisibility.AssistantOnly,
-          parameters: {
-            type: 'object',
-            properties: {
-              data: {
-                type: 'array',
-                description:
-                  'The pieces of data you want to look at it. You can request one, or multiple',
-                items: {
-                  type: 'string',
-                  enum: allData.map((data) => data.name),
-                },
-              },
-            },
-            required: ['data' as const],
-          },
-        },
-        async ({ arguments: { data: dataNames } }) => {
-          return {
-            content: allData.filter((data) => dataNames.includes(data.name)),
-          };
-        }
-      );
-
-      this.registerInstruction(({ availableFunctionNames }) =>
-        availableFunctionNames.includes(GET_DATA_ON_SCREEN_FUNCTION_NAME)
-          ? `The ${GET_DATA_ON_SCREEN_FUNCTION_NAME} function will retrieve specific content from the user's screen by specifying a data key. Use this tool to provide context-aware responses. Available data: ${dedent(
-              allData.map((data) => `${data.name}: ${data.description}`).join('\n')
-            )}`
-          : undefined
-      );
-    }
+    registerGetDataOnScreenFunction(this, screenContexts);
 
     this.actions.forEach((action) => {
       if (action.parameters) {