[O11y AI Ass] Consolidate system message & recall for insights (#168007)

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

@@ -40,14 +40,6 @@ export function ErrorSampleContextualInsight({
     const transactionName = transaction?.transaction.name ?? '';
 
     return [
-      {
-        '@timestamp': now,
-        message: {
-          role: MessageRole.System,
-          content: `You are apm-gpt, a helpful assistant for performance analysis, optimisation and
-          root cause analysis of software. Answer as concisely as possible.`,
-        },
-      },
       {
         '@timestamp': now,
         message: {

x-pack/plugins/infra/public/alerting/log_threshold/components/alert_details_app_section/components/log_rate_analysis.tsx

@@ -233,14 +233,6 @@ export const LogRateAnalysis: FC<AlertDetailsLogRateAnalysisSectionProps> = ({ r
     const now = new Date().toISOString();
 
     return [
-      {
-        '@timestamp': now,
-        message: {
-          role: MessageRole.System,
-          content: `You are logs-gpt, a helpful assistant for logs-based observability. Answer as
-          concisely as possible.`,
-        },
-      },
       {
         '@timestamp': now,
         message: {

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

@@ -45,14 +45,6 @@ export const ContextualInsightProcessRow = ({ command }: { command: string }) =>
     }
     const now = new Date().toISOString();
     return [
-      {
-        '@timestamp': now,
-        message: {
-          role: MessageRole.System,
-          content: `You are infra-gpt, a helpful assistant for metrics-based infrastructure observability. Answer as
-          concisely as possible.`,
-        },
-      },
       {
         '@timestamp': now,
         message: {

x-pack/plugins/logs_shared/public/components/logging/log_entry_flyout/log_entry_flyout.tsx

@@ -39,12 +39,6 @@ import { DataSearchProgress } from '../../data_search_progress';
 import { LogEntryActionsMenu } from './log_entry_actions_menu';
 import { LogEntryFieldsTable } from './log_entry_fields_table';
 
-const LOGS_SYSTEM_MESSAGE = {
-  content: `You are logs-gpt, a helpful assistant for logs-based observability. Answer as
-    concisely as possible.`,
-  role: MessageRole.System,
-};
-
 export interface LogEntryFlyoutProps {
   logEntryId: string | null | undefined;
   onCloseFlyout: () => void;
@@ -144,10 +138,6 @@ export const LogEntryFlyout = ({
     const now = new Date().toISOString();
 
     return [
-      {
-        '@timestamp': now,
-        message: LOGS_SYSTEM_MESSAGE,
-      },
       {
         '@timestamp': now,
         message: {
@@ -170,10 +160,6 @@ export const LogEntryFlyout = ({
     const message = logEntry.fields.find((field) => field.field === 'message')?.value[0];
 
     return [
-      {
-        '@timestamp': now,
-        message: LOGS_SYSTEM_MESSAGE,
-      },
       {
         '@timestamp': now,
         message: {

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

@@ -4,11 +4,11 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { last } from 'lodash';
 import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
-import type { Subscription } from 'rxjs';
+import { isObservable, Subscription } from 'rxjs';
 import { MessageRole, type Message } from '../../../common/types';
 import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
 import { useKibana } from '../../hooks/use_kibana';
@@ -41,6 +41,9 @@ function ChatContent({
   const chatService = useObservabilityAIAssistantChatService();
 
   const [pendingMessage, setPendingMessage] = useState<PendingMessage | undefined>();
+
+  const [recalledMessages, setRecalledMessages] = useState<Message[] | undefined>(undefined);
+
   const [loading, setLoading] = useState(false);
   const [subscription, setSubscription] = useState<Subscription | undefined>();
 
@@ -56,37 +59,101 @@ function ChatContent({
   const conversationTitle = conversationId
     ? conversation.value?.conversation.title || ''
     : defaultTitle;
-  const reloadReply = useCallback(() => {
+
+  const controllerRef = useRef(new AbortController());
+
+  const reloadRecalledMessages = useCallback(async () => {
     setLoading(true);
 
+    setDisplayedMessages(initialMessages);
+
+    setRecalledMessages(undefined);
+
+    controllerRef.current.abort();
+
+    const controller = (controllerRef.current = new AbortController());
+
+    let appendedMessages: Message[] = [];
+
+    if (chatService.hasFunction('recall')) {
+      try {
+        // manually execute recall function and append to list of
+        // messages
+        const functionCall = {
+          name: 'recall',
+          args: JSON.stringify({ queries: [], contexts: [] }),
+        };
+
+        const response = await chatService.executeFunction({
+          ...functionCall,
+          messages: initialMessages,
+          signal: controller.signal,
+          connectorId,
+        });
+
+        if (isObservable(response)) {
+          throw new Error('Recall function unexpectedly returned an Observable');
+        }
+
+        appendedMessages = [
+          {
+            '@timestamp': new Date().toISOString(),
+            message: {
+              role: MessageRole.Assistant,
+              content: '',
+              function_call: {
+                name: functionCall.name,
+                arguments: functionCall.args,
+                trigger: MessageRole.User as const,
+              },
+            },
+          },
+          {
+            '@timestamp': new Date().toISOString(),
+            message: {
+              role: MessageRole.User,
+              name: functionCall.name,
+              content: JSON.stringify(response.content),
+            },
+          },
+        ];
+
+        setRecalledMessages(appendedMessages);
+      } catch (err) {
+        // eslint-disable-next-line no-console
+        console.error(err);
+        setRecalledMessages([]);
+      }
+    }
+  }, [chatService, connectorId, initialMessages, setDisplayedMessages]);
+
+  useEffect(() => {
     let lastPendingMessage: PendingMessage | undefined;
 
+    if (recalledMessages === undefined) {
+      // don't do anything, it's loading
+      return;
+    }
+
     const nextSubscription = chatService
-      .chat({ messages: initialMessages, connectorId, function: 'none' })
+      .chat({ messages: displayedMessages.concat(recalledMessages), connectorId, function: 'none' })
       .subscribe({
         next: (msg) => {
           lastPendingMessage = msg;
           setPendingMessage(() => msg);
         },
         complete: () => {
-          setDisplayedMessages((prevMessages) =>
-            prevMessages.concat({
-              '@timestamp': new Date().toISOString(),
-              message: {
-                ...lastPendingMessage!.message,
-              },
-            })
-          );
+          setPendingMessage(lastPendingMessage);
           setLoading(false);
         },
       });
 
     setSubscription(nextSubscription);
-  }, [initialMessages, setDisplayedMessages, connectorId, chatService]);
+  }, [chatService, connectorId, displayedMessages, setDisplayedMessages, recalledMessages]);
 
   useEffect(() => {
-    reloadReply();
-  }, [reloadReply]);
+    reloadRecalledMessages();
+  }, [reloadRecalledMessages]);
 
   useEffect(() => {
     setDisplayedMessages(initialMessages);
@@ -96,23 +163,25 @@ function ChatContent({
 
   const messagesWithPending = useMemo(() => {
     return pendingMessage
-      ? displayedMessages.concat({
+      ? displayedMessages.concat(recalledMessages || []).concat({
           '@timestamp': new Date().toISOString(),
           message: {
             ...pendingMessage.message,
           },
         })
-      : displayedMessages;
-  }, [pendingMessage, displayedMessages]);
+      : displayedMessages.concat(recalledMessages || []);
+  }, [pendingMessage, displayedMessages, recalledMessages]);
 
-  const lastMessage = last(messagesWithPending);
+  const lastAssistantMessage = last(
+    messagesWithPending.filter((message) => message.message.role === MessageRole.Assistant)
+  );
 
   return (
     <>
       <MessagePanel
         body={
           <MessageText
-            content={lastMessage?.message.content ?? ''}
+            content={lastAssistantMessage?.message.content ?? ''}
             loading={loading}
             onActionClick={async () => {}}
           />
@@ -147,7 +216,7 @@ function ChatContent({
               <EuiFlexItem grow={false}>
                 <RegenerateResponseButton
                   onClick={() => {
-                    reloadReply();
+                    reloadRecalledMessages();
                   }}
                 />
               </EuiFlexItem>
@@ -168,7 +237,7 @@ function ChatContent({
         onClose={() => {
           setIsOpen(() => false);
         }}
-        messages={displayedMessages}
+        messages={messagesWithPending}
         conversationId={conversationId}
         startedFrom="contextualInsight"
         onChatComplete={(nextMessages) => {

x-pack/plugins/observability_ai_assistant/public/hooks/use_conversation.ts

@@ -6,9 +6,10 @@
  */
 import { i18n } from '@kbn/i18n';
 import { merge, omit } from 'lodash';
-import { Dispatch, SetStateAction, useState } from 'react';
+import { Dispatch, SetStateAction, useMemo, useState } from 'react';
 import { type Conversation, type Message } from '../../common';
-import type { ConversationCreateRequest } from '../../common/types';
+import { ConversationCreateRequest, MessageRole } from '../../common/types';
+import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
 import { ObservabilityAIAssistantChatService } from '../types';
 import { useAbortableAsync, type AbortableAsyncState } from './use_abortable_async';
 import { useKibana } from './use_kibana';
@@ -21,7 +22,7 @@ export function useConversation({
   connectorId,
 }: {
   conversationId?: string;
-  chatService?: ObservabilityAIAssistantChatService;
+  chatService?: ObservabilityAIAssistantChatService; // will eventually resolve to a non-nullish value
   connectorId: string | undefined;
 }): {
   conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined>;
@@ -41,6 +42,19 @@ export function useConversation({
 
   const [displayedMessages, setDisplayedMessages] = useState<Message[]>([]);
 
+  const displayedMessagesWithHardcodedSystemMessage = useMemo(() => {
+    if (!chatService) {
+      return displayedMessages;
+    }
+    const systemMessage = getAssistantSetupMessage({ contexts: chatService?.getContexts() || [] });
+
+    if (displayedMessages[0]?.message.role === MessageRole.User) {
+      return [systemMessage, ...displayedMessages];
+    }
+
+    return [systemMessage, ...displayedMessages.slice(1)];
+  }, [displayedMessages, chatService]);
+
   const conversation: AbortableAsyncState<ConversationCreateRequest | Conversation | undefined> =
     useAbortableAsync(
       ({ signal }) => {
@@ -71,7 +85,7 @@ export function useConversation({
 
   return {
     conversation,
-    displayedMessages,
+    displayedMessages: displayedMessagesWithHardcodedSystemMessage,
     setDisplayedMessages,
     save: (messages: Message[], handleRefreshConversations?: () => void) => {
       const conversationObject = conversation.value!;

x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.test.ts

@@ -171,6 +171,8 @@ describe('useTimeline', () => {
             return subject;
           }),
           executeFunction: jest.fn(),
+          hasFunction: jest.fn(),
+          hasRenderFunction: jest.fn(),
         },
         onChatUpdate: jest.fn().mockImplementation((messages) => {
           props = { ...props, messages };