[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 

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

@@ -137,7 +137,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,10 +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/public/components/message_panel/esql_code_block.tsx

@@ -10,13 +10,48 @@ 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};
+  }
+`;
+
+export function CodeBlock({ children }: { children: React.ReactNode }) {
+  const theme = useEuiTheme();
+
+  return (
+    <EuiPanel
+      hasShadow={false}
+      hasBorder={false}
+      paddingSize="s"
+      className={getCodeBlockClassName(theme)}
+    >
+      <EuiFlexGroup direction="column" gutterSize="xs">
+        <EuiFlexItem grow={false}>
+          <EuiCodeBlock isCopyable fontSize="m">
+            {children}
+          </EuiCodeBlock>
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </EuiPanel>
+  );
+}
+
 export function EsqlCodeBlock({
   value,
   actionsDisabled,
@@ -33,18 +68,7 @@ export function EsqlCodeBlock({
       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};
-        }
-      `}
+      className={getCodeBlockClassName(theme)}
     >
       <EuiFlexGroup direction="column" gutterSize="xs">
         <EuiFlexItem grow={false}>

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

@@ -17,7 +17,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;
@@ -100,6 +100,8 @@ const esqlLanguagePlugin = () => {
 
     if (node.type === 'code' && node.lang === 'esql') {
       node.type = 'esql';
+    } else if (node.type === 'code') {
+      node.type = 'codeBlock';
     }
   };
 
@@ -127,6 +129,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();

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

@@ -18,11 +18,7 @@ import {
   isTokenLimitReachedError,
   StreamingChatResponseEventType,
 } from '../../common';
-import {
-  getAssistantSystemMessage,
-  type ObservabilityAIAssistantChatService,
-  type ObservabilityAIAssistantService,
-} from '..';
+import type { ObservabilityAIAssistantChatService, ObservabilityAIAssistantService } from '..';
 import { useKibana } from './use_kibana';
 import { useOnce } from './use_once';
 import { useUserPreferredLanguage } from './use_user_preferred_language';
@@ -75,13 +71,11 @@ function useChatWithoutContext({
   persist,
 }: UseChatPropsWithoutContext): UseChatResult {
   const [chatState, setChatState] = useState(ChatState.Ready);
-
   const systemMessage = useMemo(() => {
-    return getAssistantSystemMessage({ contexts: chatService.getContexts() });
+    return chatService.getSystemMessage();
   }, [chatService]);
 
   useOnce(initialMessages);
-
   useOnce(initialConversationId);
 
   const [conversationId, setConversationId] = useState(initialConversationId);

x-pack/plugins/observability_solution/observability_ai_assistant/public/index.ts

@@ -56,8 +56,6 @@ export {
   VISUALIZE_ESQL_USER_INTENTIONS,
 } from '../common/functions/visualize_esql';
 
-export { getAssistantSystemMessage } from './service/get_assistant_system_message';
-
 export { isSupportedConnectorType } from '../common';
 export { FunctionVisibility } from '../common';
 

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

@@ -23,7 +23,6 @@ export const mockChatService: ObservabilityAIAssistantChatService = {
   sendAnalyticsEvent: noop,
   chat: (options) => new Observable<StreamingChatResponseEventWithoutError>(),
   complete: (options) => new Observable<StreamingChatResponseEventWithoutError>(),
-  getContexts: () => [],
   getFunctions: () => [buildFunctionElasticsearch(), buildFunctionServiceSummary()],
   renderFunction: (name) => (
     <div>