[AI Assistant] Add scopes to AI assistant (elastic#192611)

## Summary

This adds scopes to the Observability AI Assistant framework so we can
start consuming it in different places with different instructions and
functions available.

Instructions and functions become responsible for registering which
'scope' they apply to, and the assistant should be called with a scope
when chatting so that it can decide which functions are available and
which registered instructions to send to the LLM.

x-pack/plugins/observability_solution/apm/public/components/app/error_group_details/error_sampler/error_sample_contextual_insight.tsx

@@ -38,19 +38,19 @@ export function ErrorSampleContextualInsight({
       instructions: `I'm an SRE. I am looking at an exception and trying to understand what it means.
 
       Your task is to describe what the error means and what it could be caused by.
-      
+
       The error occurred on a service called ${serviceName}, which is a ${runtimeName} service written in ${languageName}. The
       runtime version is ${runtimeVersion}.
-      
+
       The request it occurred for is called ${transactionName}.
-      
+
       ${
         logStacktrace
           ? `The log stacktrace:
       ${logStacktrace}`
           : ''
       }
-      
+
       ${
         exceptionStacktrace
           ? `The exception stacktrace:

x-pack/plugins/observability_solution/apm/server/assistant_functions/get_apm_dataset_info.ts

@@ -164,6 +164,7 @@ export function registerGetApmDatasetInfoFunction({
           `,
         },
       };
-    }
+    },
+    ['observability']
   );
 }

x-pack/plugins/observability_solution/apm/server/assistant_functions/get_apm_downstream_dependencies.ts

@@ -22,16 +22,16 @@ export function registerGetApmDownstreamDependenciesFunction({
   registerFunction(
     {
       name: 'get_apm_downstream_dependencies',
-      description: `Get the downstream dependencies (services or uninstrumented backends) for a 
-      service. This allows you to map the downstream dependency name to a service, by 
-      returning both span.destination.service.resource and service.name. Use this to 
+      description: `Get the downstream dependencies (services or uninstrumented backends) for a
+      service. This allows you to map the downstream dependency name to a service, by
+      returning both span.destination.service.resource and service.name. Use this to
       drilldown further if needed.`,
       descriptionForUser: i18n.translate(
         'xpack.apm.observabilityAiAssistant.functions.registerGetApmDownstreamDependencies.descriptionForUser',
         {
-          defaultMessage: `Get the downstream dependencies (services or uninstrumented backends) for a 
-      service. This allows you to map the dowstream dependency name to a service, by 
-      returning both span.destination.service.resource and service.name. Use this to 
+          defaultMessage: `Get the downstream dependencies (services or uninstrumented backends) for a
+      service. This allows you to map the dowstream dependency name to a service, by
+      returning both span.destination.service.resource and service.name. Use this to
       drilldown further if needed.`,
         }
       ),
@@ -67,6 +67,7 @@ export function registerGetApmDownstreamDependenciesFunction({
           randomSampler,
         }),
       };
-    }
+    },
+    ['observability']
   );
 }

x-pack/plugins/observability_solution/apm/server/assistant_functions/get_apm_services_list.ts

@@ -84,6 +84,7 @@ export function registerGetApmServicesListFunction({
           arguments: args,
         }),
       };
-    }
+    },
+    ['observability']
   );
 }

x-pack/plugins/observability_solution/apm/server/assistant_functions/get_apm_timeseries.ts

@@ -138,7 +138,8 @@ export function registerGetApmTimeseriesFunction({
         content: timeseries.map((series): Omit<ApmTimeseries, 'data'> => omit(series, 'data')),
         data: timeseries,
       };
-    }
+    },
+    ['observability']
   );
 }
 

x-pack/plugins/observability_solution/infra/public/components/asset_details/tabs/processes/process_row.tsx

@@ -48,7 +48,7 @@ export const ContextualInsightProcessRow = ({ command }: { command: string }) =>
       with the arguments to the process you should then explain its arguments and how they influence the behaviour
       of the process. If I do not provide any arguments then explain the behaviour of the process when no arguments are
       provided.
-      
+
       Here is an example with arguments.
       Process: metricbeat -c /etc/metricbeat.yml -d autodiscover,kafka -e -system.hostfs=/hostfs
       Explanation: Metricbeat is part of the Elastic Stack. It is a lightweight shipper that you can install on your

x-pack/plugins/observability_solution/observability_ai_assistant/common/functions/types.ts

@@ -8,6 +8,7 @@ import type { JSONSchema7TypeName } from 'json-schema';
 import type { Observable } from 'rxjs';
 import { ChatCompletionChunkEvent, MessageAddEvent } from '../conversation_complete';
 import { FunctionVisibility } from './function_visibility';
+import { AssistantScope } from '../types';
 export { FunctionVisibility };
 
 type JSONSchemaOrPrimitive = CompatibleJSONSchema | string | number | boolean;
@@ -41,6 +42,7 @@ export interface FunctionDefinition<TParameters extends CompatibleJSONSchema = a
   visibility?: FunctionVisibility;
   descriptionForUser?: string;
   parameters?: TParameters;
+  scopes?: AssistantScope[];
 }
 
 export type FunctionRegistry = Map<string, FunctionDefinition>;

x-pack/plugins/observability_solution/observability_ai_assistant/common/types.ts

@@ -157,3 +157,5 @@ export interface ObservabilityAIAssistantScreenContext {
   actions?: Array<ScreenContextActionDefinition<any>>;
   starterPrompts?: StarterPrompt[];
 }
+
+export type AssistantScope = 'observability' | 'search' | 'all';

x-pack/plugins/observability_solution/observability_ai_assistant/public/components/insight/insight.tsx

@@ -56,6 +56,7 @@ function ChatContent({
 }) {
   const service = useObservabilityAIAssistant();
   const chatService = useObservabilityAIAssistantChatService();
+  const { scope } = service;
 
   const initialMessagesRef = useRef(initialMessages);
 
@@ -68,6 +69,7 @@ function ChatContent({
     initialMessages,
     persist: false,
     disableFunctions: true,
+    scope,
   });
 
   const lastAssistantResponse = getLastMessageOfType(

x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.test.ts

@@ -80,6 +80,7 @@ describe('useChat', () => {
           service: {
             getScreenContexts: () => [],
           } as unknown as ObservabilityAIAssistantService,
+          scope: 'observability',
         } as UseChatProps,
       });
     });
@@ -109,6 +110,7 @@ describe('useChat', () => {
           service: {
             getScreenContexts: () => [],
           } as unknown as ObservabilityAIAssistantService,
+          scope: 'observability',
         } as UseChatProps,
       });
 

x-pack/plugins/observability_solution/observability_ai_assistant/public/hooks/use_chat.ts

@@ -10,6 +10,7 @@ import { merge } from 'lodash';
 import { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
 import type { NotificationsStart } from '@kbn/core/public';
+import { AssistantScope } from '../../common/types';
 import {
   MessageRole,
   type Message,
@@ -55,6 +56,7 @@ interface UseChatPropsWithoutContext {
   disableFunctions?: boolean;
   onConversationUpdate?: (event: ConversationCreateEvent | ConversationUpdateEvent) => void;
   onChatComplete?: (messages: Message[]) => void;
+  scope: AssistantScope;
 }
 
 export type UseChatProps = Omit<UseChatPropsWithoutContext, 'notifications'>;
@@ -70,6 +72,7 @@ function useChatWithoutContext({
   onChatComplete,
   persist,
   disableFunctions,
+  scope,
 }: UseChatPropsWithoutContext): UseChatResult {
   const [chatState, setChatState] = useState(ChatState.Ready);
   const systemMessage = useMemo(() => {
@@ -161,6 +164,7 @@ function useChatWithoutContext({
         disableFunctions: disableFunctions ?? false,
         signal: abortControllerRef.current.signal,
         conversationId,
+        scope,
       });
 
       function getPendingMessages() {
@@ -259,6 +263,7 @@ function useChatWithoutContext({
       disableFunctions,
       service,
       systemMessage,
+      scope,
     ]
   );
 

x-pack/plugins/observability_solution/observability_ai_assistant/public/mock.tsx

@@ -60,6 +60,7 @@ export const mockService: ObservabilityAIAssistantService = {
     predefinedConversation$: new Observable(),
   },
   navigate: async () => of(),
+  scope: 'all',
 };
 
 function createSetupContract(): ObservabilityAIAssistantPublicSetup {

x-pack/plugins/observability_solution/observability_ai_assistant/public/plugin.tsx

@@ -65,6 +65,7 @@ export class ObservabilityAIAssistantPlugin
         coreStart.application.capabilities.observabilityAIAssistant[
           aiAssistantCapabilities.show
         ] === true,
+      scope: 'observability',
     }));
 
     const withProviders = <P extends {}, R = {}>(

x-pack/plugins/observability_solution/observability_ai_assistant/public/service/complete.test.ts

@@ -102,6 +102,7 @@ describe('complete', () => {
         disableFunctions: false,
         signal: new AbortController().signal,
         ...params,
+        scope: 'all',
       },
       requestCallback
     );

x-pack/plugins/observability_solution/observability_ai_assistant/public/service/complete.ts

@@ -43,6 +43,7 @@ export function complete(
     disableFunctions,
     signal,
     instructions,
+    scope,
   }: {
     client: Pick<ObservabilityAIAssistantChatService, 'chat' | 'complete'>;
     getScreenContexts: () => ObservabilityAIAssistantScreenContext[];
@@ -65,6 +66,7 @@ export function complete(
           screenContexts,
           conversationId,
           instructions,
+          scope,
         },
       },
     }).pipe(shareReplay());
@@ -131,6 +133,7 @@ export function complete(
               persist,
               disableFunctions,
               instructions,
+              scope,
             },
             requestCallback
           ).subscribe(subscriber);

x-pack/plugins/observability_solution/observability_ai_assistant/public/service/create_chat_service.test.ts

@@ -70,6 +70,7 @@ describe('createChatService', () => {
       apiClient: clientSpy,
       registrations: [],
       signal: new AbortController().signal,
+      scope: 'observability',
     });
   });
 
@@ -79,7 +80,12 @@ describe('createChatService', () => {
 
   describe('chat', () => {
     function chat({ signal }: { signal: AbortSignal } = { signal: new AbortController().signal }) {
-      return service.chat('my_test', { signal, messages: [], connectorId: '' });
+      return service.chat('my_test', {
+        signal,
+        messages: [],
+        connectorId: '',
+        scope: 'observability',
+      });
     }
 
     it('correctly parses a stream of JSON lines', async () => {

x-pack/plugins/observability_solution/observability_ai_assistant/public/service/create_chat_service.ts

@@ -23,6 +23,7 @@ import {
   throwError,
   timestamp,
 } from 'rxjs';
+import { AssistantScope } from '../../common/types';
 import { ChatCompletionChunkEvent, Message, MessageRole } from '../../common';
 import {
   StreamingChatResponseEventType,
@@ -137,19 +138,26 @@ export async function createChatService({
   signal: setupAbortSignal,
   registrations,
   apiClient,
+  scope,
 }: {
   analytics: AnalyticsServiceStart;
   signal: AbortSignal;
   registrations: ChatRegistrationRenderFunction[];
   apiClient: ObservabilityAIAssistantAPIClient;
+  scope: AssistantScope;
 }): Promise<ObservabilityAIAssistantChatService> {
   const functionRegistry: FunctionRegistry = new Map();
 
   const renderFunctionRegistry: Map<string, RenderFunction<unknown, FunctionResponse>> = new Map();
 
   const [{ functionDefinitions, systemMessage }] = await Promise.all([
-    apiClient('GET /internal/observability_ai_assistant/functions', {
+    apiClient('GET /internal/observability_ai_assistant/{scope}/functions', {
       signal: setupAbortSignal,
+      params: {
+        path: {
+          scope,
+        },
+      },
     }),
     ...registrations.map((registration) => {
       return registration({
@@ -196,6 +204,7 @@ export async function createChatService({
             connectorId,
             functionCall,
             functions: functions ?? [],
+            scope,
           },
         },
         signal,
@@ -228,6 +237,7 @@ export async function createChatService({
           signal,
           client,
           instructions,
+          scope,
         },
         ({ params }) => {
           return callStreamingApi('POST /internal/observability_ai_assistant/chat/complete', {

x-pack/plugins/observability_solution/observability_ai_assistant/public/service/create_service.ts

@@ -8,7 +8,11 @@
 import type { AnalyticsServiceStart, CoreStart } from '@kbn/core/public';
 import { compact, without } from 'lodash';
 import { BehaviorSubject, debounceTime, filter, lastValueFrom, of, Subject, take } from 'rxjs';
-import type { Message, ObservabilityAIAssistantScreenContext } from '../../common/types';
+import type {
+  AssistantScope,
+  Message,
+  ObservabilityAIAssistantScreenContext,
+} from '../../common/types';
 import { createFunctionRequestMessage } from '../../common/utils/create_function_request_message';
 import { createFunctionResponseMessage } from '../../common/utils/create_function_response_message';
 import { createCallObservabilityAIAssistantAPI } from '../api';
@@ -19,10 +23,12 @@ export function createService({
   analytics,
   coreStart,
   enabled,
+  scope,
 }: {
   analytics: AnalyticsServiceStart;
   coreStart: CoreStart;
   enabled: boolean;
+  scope: AssistantScope;
 }): ObservabilityAIAssistantService {
   const apiClient = createCallObservabilityAIAssistantAPI(coreStart);
 
@@ -42,7 +48,7 @@ export function createService({
     },
     start: async ({ signal }) => {
       const mod = await import('./create_chat_service');
-      return await mod.createChatService({ analytics, apiClient, signal, registrations });
+      return await mod.createChatService({ analytics, apiClient, signal, registrations, scope });
     },
     callApi: apiClient,
     getScreenContexts() {
@@ -89,5 +95,6 @@ export function createService({
       },
       predefinedConversation$: predefinedConversation$.asObservable(),
     },
+    scope,
   };
 }

x-pack/plugins/observability_solution/observability_ai_assistant/public/storybook_mock.tsx

@@ -52,4 +52,5 @@ export const createStorybookService = (): ObservabilityAIAssistantService => ({
     predefinedConversation$: new Observable(),
   },
   navigate: async () => of(),
+  scope: 'observability',
 });