[Obs AI Assistant] Instructions & Claude improvements

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

@@ -20,7 +20,6 @@ export function registerGetApmDatasetInfoFunction({
   registerFunction(
     {
       name: 'get_apm_dataset_info',
-      contexts: ['core'],
       visibility: FunctionVisibility.AssistantOnly,
       description: `Use this function to get information about APM data.`,
       parameters: {

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

@@ -16,7 +16,6 @@ export function registerGetApmDownstreamDependenciesFunction({
   registerFunction(
     {
       name: 'get_apm_downstream_dependencies',
-      contexts: ['core'],
       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 
@@ -39,7 +38,8 @@ export function registerGetApmDownstreamDependenciesFunction({
           },
           'service.environment': {
             type: 'string',
-            description: 'The environment that the service is running in',
+            description:
+              'The environment that the service is running in. Leave empty to query for all environments.',
           },
           start: {
             type: 'string',

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

@@ -22,7 +22,6 @@ export function registerGetApmServicesListFunction({
   registerFunction(
     {
       name: 'get_apm_services_list',
-      contexts: ['apm'],
       description: `Gets a list of services`,
       descriptionForUser: i18n.translate(
         'xpack.apm.observabilityAiAssistant.functions.registerGetApmServicesList.descriptionForUser',

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

@@ -122,7 +122,6 @@ export function registerGetApmTimeseriesFunction({
 }: FunctionRegistrationParameters) {
   registerFunction(
     {
-      contexts: ['core'],
       name: 'get_apm_timeseries',
       description: `Visualise and analyse different APM metrics, like throughput, failure rate, or latency, for any service or all services, or any or all of its dependencies, both as a timeseries and as a single statistic. A visualisation will be displayed above your reply - DO NOT attempt to display or generate an image yourself, or any other placeholder. Additionally, the function will return any changes, such as spikes, step and trend changes, or dips. You can also use it to compare data by requesting two different time ranges, or for instance two different service versions.`,
       parameters,

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

@@ -48,7 +48,7 @@ export function registerAssistantFunctions({
   ruleDataClient: IRuleDataClient;
   plugins: APMRouteHandlerResources['plugins'];
 }): RegistrationCallback {
-  return async ({ resources, functions: { registerContext, registerFunction } }) => {
+  return async ({ resources, functions: { registerFunction } }) => {
     const apmRouteHandlerResources: MinimalAPMRouteHandlerResources = {
       context: resources.context,
       request: resources.request,

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

@@ -28,11 +28,6 @@ export type CompatibleJSONSchema = {
   description?: string;
 };
 
-export interface ContextDefinition {
-  name: string;
-  description: string;
-}
-
 export type FunctionResponse =
   | {
       content?: any;
@@ -46,10 +41,6 @@ export interface FunctionDefinition<TParameters extends CompatibleJSONSchema = a
   visibility?: FunctionVisibility;
   descriptionForUser?: string;
   parameters?: TParameters;
-  contexts: string[];
 }
 
-export type RegisterContextDefinition = (options: ContextDefinition) => void;
-
-export type ContextRegistry = Map<string, ContextDefinition>;
 export type FunctionRegistry = Map<string, FunctionDefinition>;

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

@@ -94,6 +94,8 @@ export interface UserInstruction {
   text: string;
 }
 
+export type UserInstructionOrPlainText = string | UserInstruction;
+
 export interface ObservabilityAIAssistantScreenContextRequest {
   screenDescription?: string;
   data?: Array<{

x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/emit_with_concatenated_message.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import { concat, last, mergeMap, Observable, shareReplay, withLatestFrom } from 'rxjs';
+import { concat, from, last, mergeMap, Observable, shareReplay, withLatestFrom } from 'rxjs';
 import {
   ChatCompletionChunkEvent,
   MessageAddEvent,
@@ -16,8 +16,32 @@ import {
   ConcatenatedMessage,
 } from './concatenate_chat_completion_chunks';
 
+type ConcatenateMessageCallback = (
+  concatenatedMessage: ConcatenatedMessage
+) => Promise<ConcatenatedMessage>;
+
+function mergeWithEditedMessage(
+  originalMessage: ConcatenatedMessage,
+  chunkEvent: ChatCompletionChunkEvent,
+  callback?: ConcatenateMessageCallback
+): Observable<MessageAddEvent> {
+  return from(
+    (callback ? callback(originalMessage) : Promise.resolve(originalMessage)).then((message) => {
+      const next: MessageAddEvent = {
+        type: StreamingChatResponseEventType.MessageAdd as const,
+        id: chunkEvent.id,
+        message: {
+          '@timestamp': new Date().toISOString(),
+          ...message,
+        },
+      };
+      return next;
+    })
+  );
+}
+
 export function emitWithConcatenatedMessage(
-  callback?: (concatenatedMessage: ConcatenatedMessage) => Promise<ConcatenatedMessage>
+  callback?: ConcatenateMessageCallback
 ): (
   source$: Observable<ChatCompletionChunkEvent>
 ) => Observable<ChatCompletionChunkEvent | MessageAddEvent> {
@@ -30,17 +54,8 @@ export function emitWithConcatenatedMessage(
         concatenateChatCompletionChunks(),
         last(),
         withLatestFrom(source$),
-        mergeMap(async ([message, chunkEvent]) => {
-          const next: MessageAddEvent = {
-            type: StreamingChatResponseEventType.MessageAdd as const,
-            id: chunkEvent.id,
-            message: {
-              '@timestamp': new Date().toISOString(),
-              ...(callback ? await callback(message) : message),
-            },
-          };
-
-          return next;
+        mergeMap(([message, chunkEvent]) => {
+          return mergeWithEditedMessage(message, chunkEvent, callback);
         })
       )
     );

x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/filter_function_definitions.ts

@@ -8,22 +8,18 @@
 import type { FunctionDefinition } from '../functions/types';
 
 export function filterFunctionDefinitions({
-  contexts,
   filter,
   definitions,
 }: {
-  contexts?: string[];
   filter?: string;
   definitions: FunctionDefinition[];
 }) {
-  return contexts || filter
+  return filter
     ? definitions.filter((fn) => {
-        const matchesContext =
-          !contexts || fn.contexts.some((context) => contexts.includes(context));
         const matchesFilter =
           !filter || fn.name.includes(filter) || fn.description.includes(filter);
 
-        return matchesContext && matchesFilter;
+        return matchesFilter;
       })
     : definitions;
 }

x-pack/plugins/observability_solution/observability_ai_assistant/common/utils/throw_serialized_chat_completion_errors.ts

@@ -5,20 +5,21 @@
  * 2.0.
  */
 
-import { filter, Observable, tap } from 'rxjs';
+import { filter, OperatorFunction, tap } from 'rxjs';
 import {
   ChatCompletionError,
   ChatCompletionErrorCode,
   type StreamingChatResponseEvent,
   StreamingChatResponseEventType,
   type ChatCompletionErrorEvent,
+  BufferFlushEvent,
 } from '../conversation_complete';
 
-export function throwSerializedChatCompletionErrors() {
-  return <T extends StreamingChatResponseEvent>(
-    source$: Observable<StreamingChatResponseEvent>
-  ): Observable<Exclude<T, ChatCompletionErrorEvent>> => {
-    return source$.pipe(
+export function throwSerializedChatCompletionErrors<
+  T extends StreamingChatResponseEvent | BufferFlushEvent
+>(): OperatorFunction<T, Exclude<T, ChatCompletionErrorEvent>> {
+  return (source$) =>
+    source$.pipe(
       tap((event) => {
         // de-serialise error
         if (event.type === StreamingChatResponseEventType.ChatCompletionError) {
@@ -33,5 +34,4 @@ export function throwSerializedChatCompletionErrors() {
           event.type !== StreamingChatResponseEventType.ChatCompletionError
       )
     );
-  };
 }

x-pack/plugins/observability_solution/observability_ai_assistant/public/components/message_panel/esql_code_block.tsx

@@ -10,13 +10,55 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
   EuiPanel,
+  UseEuiTheme,
   useEuiTheme,
 } from '@elastic/eui';
 import { css } from '@emotion/css';
 import { i18n } from '@kbn/i18n';
 import React from 'react';
 import { ChatActionClickHandler, ChatActionClickType } from '../chat/types';
 
+const getCodeBlockClassName = (theme: UseEuiTheme) => css`
+  background-color: ${theme.euiTheme.colors.lightestShade};
+  .euiCodeBlock__pre {
+    margin-bottom: 0;
+    padding: ${theme.euiTheme.size.m};
+    min-block-size: 48px;
+  }
+  .euiCodeBlock__controls {
+    inset-block-start: ${theme.euiTheme.size.m};
+    inset-inline-end: ${theme.euiTheme.size.m};
+  }
+`;
+
+function CodeBlockWrapper({ children }: { children: React.ReactNode }) {
+  const theme = useEuiTheme();
+  return (
+    <EuiPanel
+      hasShadow={false}
+      hasBorder={false}
+      paddingSize="s"
+      className={getCodeBlockClassName(theme)}
+    >
+      {children}
+    </EuiPanel>
+  );
+}
+
+export function CodeBlock({ children }: { children: React.ReactNode }) {
+  return (
+    <CodeBlockWrapper>
+      <EuiFlexGroup direction="column" gutterSize="xs">
+        <EuiFlexItem grow={false}>
+          <EuiCodeBlock isCopyable fontSize="m">
+            {children}
+          </EuiCodeBlock>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </CodeBlockWrapper>
+  );
+}
+
 export function EsqlCodeBlock({
   value,
   actionsDisabled,
@@ -26,26 +68,8 @@ export function EsqlCodeBlock({
   actionsDisabled: boolean;
   onActionClick: ChatActionClickHandler;
 }) {
-  const theme = useEuiTheme();
-
   return (
-    <EuiPanel
-      hasShadow={false}
-      hasBorder={false}
-      paddingSize="s"
-      className={css`
-        background-color: ${theme.euiTheme.colors.lightestShade};
-        .euiCodeBlock__pre {
-          margin-bottom: 0;
-          padding: ${theme.euiTheme.size.m};
-          min-block-size: 48px;
-        }
-        .euiCodeBlock__controls {
-          inset-block-start: ${theme.euiTheme.size.m};
-          inset-inline-end: ${theme.euiTheme.size.m};
-        }
-      `}
-    >
+    <CodeBlockWrapper>
       <EuiFlexGroup direction="column" gutterSize="xs">
         <EuiFlexItem grow={false}>
           <EuiCodeBlock isCopyable fontSize="m">
@@ -87,6 +111,6 @@ export function EsqlCodeBlock({
           </EuiFlexGroup>
         </EuiFlexItem>
       </EuiFlexGroup>
-    </EuiPanel>
+    </CodeBlockWrapper>
   );
 }

x-pack/plugins/observability_solution/observability_ai_assistant/public/components/message_panel/message_text.tsx

@@ -21,7 +21,7 @@ import type { Code, InlineCode, Parent, Text } from 'mdast';
 import React, { useMemo, useRef } from 'react';
 import type { Node } from 'unist';
 import { ChatActionClickHandler } from '../chat/types';
-import { EsqlCodeBlock } from './esql_code_block';
+import { CodeBlock, EsqlCodeBlock } from './esql_code_block';
 
 interface Props {
   content: string;
@@ -104,6 +104,9 @@ const esqlLanguagePlugin = () => {
 
     if (node.type === 'code' && node.lang === 'esql') {
       node.type = 'esql';
+    } else if (node.type === 'code') {
+      // switch to type that allows us to control rendering
+      node.type = 'codeBlock';
     }
   };
 
@@ -131,6 +134,14 @@ export function MessageText({ loading, content, onActionClick }: Props) {
     processingPlugins[1][1].components = {
       ...components,
       cursor: Cursor,
+      codeBlock: (props) => {
+        return (
+          <>
+            <CodeBlock>{props.value}</CodeBlock>
+            <EuiSpacer size="m" />
+          </>
+        );
+      },
       esql: (props) => {
         return (
           <>

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

@@ -26,11 +26,17 @@ const mockChatService: MockedChatService = {
   chat: jest.fn(),
   complete: jest.fn(),
   sendAnalyticsEvent: jest.fn(),
-  getContexts: jest.fn().mockReturnValue([{ name: 'core', description: '' }]),
   getFunctions: jest.fn().mockReturnValue([]),
   hasFunction: jest.fn().mockReturnValue(false),
   hasRenderFunction: jest.fn().mockReturnValue(true),
   renderFunction: jest.fn(),
+  getSystemMessage: jest.fn().mockReturnValue({
+    '@timestamp': new Date().toISOString(),
+    message: {
+      content: 'system',
+      role: MessageRole.System,
+    },
+  }),
 };
 
 const addErrorMock = jest.fn();