[Observability AI Assistant] Make sure insight > chat flyout works

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

@@ -34,11 +34,12 @@ export function ObservabilityAIAssistantActionMenuItem() {
 
   const [conversationId, setConversationId] = useState<string>();
 
-  const { conversation, displayedMessages, setDisplayedMessages, save } = useConversation({
-    conversationId,
-    connectorId: connectors.selectedConnector,
-    chatService: chatService.value,
-  });
+  const { conversation, displayedMessages, setDisplayedMessages, save, saveTitle } =
+    useConversation({
+      conversationId,
+      connectorId: connectors.selectedConnector,
+      chatService: chatService.value,
+    });
 
   if (!service.isEnabled()) {
     return null;
@@ -74,6 +75,7 @@ export function ObservabilityAIAssistantActionMenuItem() {
             title={conversation.value?.conversation.title ?? EMPTY_CONVERSATION_TITLE}
             messages={displayedMessages}
             conversationId={conversationId}
+            startedFrom="appTopNavbar"
             onClose={() => {
               setIsOpen(() => false);
             }}
@@ -87,6 +89,9 @@ export function ObservabilityAIAssistantActionMenuItem() {
             onChatUpdate={(nextMessages) => {
               setDisplayedMessages(nextMessages);
             }}
+            onChatTitleSave={(newTitle) => {
+              saveTitle(newTitle);
+            }}
           />
         </ObservabilityAIAssistantChatServiceProvider>
       ) : null}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_body.tsx

@@ -30,6 +30,7 @@ import { IncorrectLicensePanel } from './incorrect_license_panel';
 import { ChatHeader } from './chat_header';
 import { ChatPromptEditor } from './chat_prompt_editor';
 import { ChatTimeline } from './chat_timeline';
+import { StartedFrom } from '../../utils/get_timeline_items_from_conversation';
 
 const containerClassName = css`
   max-height: 100%;
@@ -58,6 +59,7 @@ export function ChatBody({
   connectorsManagementHref,
   conversationId,
   currentUser,
+  startedFrom,
   onChatUpdate,
   onChatComplete,
   onSaveTitle,
@@ -70,6 +72,7 @@ export function ChatBody({
   connectorsManagementHref: string;
   conversationId?: string;
   currentUser?: Pick<AuthenticatedUser, 'full_name' | 'username'>;
+  startedFrom?: StartedFrom;
   onChatUpdate: (messages: Message[]) => void;
   onChatComplete: (messages: Message[]) => void;
   onSaveTitle: (title: string) => void;
@@ -84,6 +87,7 @@ export function ChatBody({
     connectors,
     currentUser,
     chatService,
+    startedFrom,
     onChatUpdate,
     onChatComplete,
   });
@@ -210,8 +214,8 @@ export function ChatBody({
       <EuiFlexItem grow={false}>
         <ChatHeader
           connectors={connectors}
-          licenseInvalid={!hasCorrectLicense && !conversationId}
           knowledgeBase={knowledgeBase}
+          licenseInvalid={!hasCorrectLicense && !conversationId}
           loading={loading}
           title={title}
           onSaveTitle={onSaveTitle}

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.stories.tsx

@@ -31,7 +31,11 @@ const defaultProps: ChatFlyoutProps = {
   isOpen: true,
   title: 'How is this working',
   messages: [getAssistantSetupMessage({ contexts: [] })],
+  startedFrom: 'appTopNavbar',
   onClose: () => {},
+  onChatComplete: () => {},
+  onChatTitleSave: () => {},
+  onChatUpdate: () => {},
 };
 
 export const ChatFlyout = Template.bind({});

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_flyout.tsx

@@ -9,15 +9,13 @@ import { css } from '@emotion/css';
 import { i18n } from '@kbn/i18n';
 import React from 'react';
 import type { Message } from '../../../common/types';
-import { useAbortableAsync } from '../../hooks/use_abortable_async';
-import { useConversation } from '../../hooks/use_conversation';
 import { useCurrentUser } from '../../hooks/use_current_user';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
 import { useKibana } from '../../hooks/use_kibana';
 import { useKnowledgeBase } from '../../hooks/use_knowledge_base';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { useObservabilityAIAssistantRouter } from '../../hooks/use_observability_ai_assistant_router';
 import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
+import { StartedFrom } from '../../utils/get_timeline_items_from_conversation';
 import { ChatBody } from './chat_body';
 
 const containerClassName = css`
@@ -33,17 +31,21 @@ export function ChatFlyout({
   messages,
   conversationId,
   isOpen,
+  startedFrom,
   onClose,
   onChatUpdate,
   onChatComplete,
+  onChatTitleSave,
 }: {
   title: string;
   messages: Message[];
   conversationId?: string;
   isOpen: boolean;
+  startedFrom: StartedFrom;
   onClose: () => void;
-  onChatUpdate?: (messages: Message[]) => void;
-  onChatComplete?: (messages: Message[]) => void;
+  onChatUpdate: (messages: Message[]) => void;
+  onChatComplete: (messages: Message[]) => void;
+  onChatTitleSave: (title: string) => void;
 }) {
   const { euiTheme } = useEuiTheme();
   const {
@@ -54,25 +56,10 @@ export function ChatFlyout({
 
   const connectors = useGenAIConnectors();
 
-  const service = useObservabilityAIAssistant();
-
-  const chatService = useAbortableAsync(
-    ({ signal }) => {
-      return service.start({ signal });
-    },
-    [service]
-  );
-
   const router = useObservabilityAIAssistantRouter();
 
   const knowledgeBase = useKnowledgeBase();
 
-  const { saveTitle } = useConversation({
-    conversationId,
-    chatService: chatService.value,
-    connectorId: connectors.selectedConnector,
-  });
-
   return isOpen ? (
     <EuiFlyout onClose={onClose}>
       <EuiFlexGroup
@@ -116,6 +103,7 @@ export function ChatFlyout({
             currentUser={currentUser}
             connectorsManagementHref={getConnectorsManagementHref(http)}
             knowledgeBase={knowledgeBase}
+            startedFrom={startedFrom}
             onChatUpdate={(nextMessages) => {
               if (onChatUpdate) {
                 onChatUpdate(nextMessages);
@@ -127,7 +115,7 @@ export function ChatFlyout({
               }
             }}
             onSaveTitle={(newTitle) => {
-              saveTitle(newTitle);
+              onChatTitleSave(newTitle);
             }}
           />
         </EuiFlexItem>

x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx

@@ -54,7 +54,7 @@ export function ChatPromptEditor({
 
   const { model, initialJsonString } = useJsonEditorModel({
     functionName: selectedFunctionName,
-    initialJson: initialFunctionPayload,
+    initialJson: functionPayload,
   });
 
   const textAreaRef = useRef<HTMLTextAreaElement>(null);

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

@@ -4,36 +4,38 @@
  * 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 { last } from 'lodash';
 import { EuiFlexGroup, EuiFlexItem } from '@elastic/eui';
-import { useKibana } from '@kbn/kibana-react-plugin/public';
 import { AbortError } from '@kbn/kibana-utils-plugin/common';
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
 import type { 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';
+import { useAbortableAsync } from '../../hooks/use_abortable_async';
+import { useConversation } from '../../hooks/use_conversation';
 import { useGenAIConnectors } from '../../hooks/use_genai_connectors';
+import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
+import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
 import type { PendingMessage } from '../../types';
+import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
+import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
+import { StartChatButton } from '../buttons/start_chat_button';
+import { StopGeneratingButton } from '../buttons/stop_generating_button';
 import { ChatFlyout } from '../chat/chat_flyout';
 import { ConnectorSelectorBase } from '../connector_selector/connector_selector_base';
 import { MessagePanel } from '../message_panel/message_panel';
 import { MessageText } from '../message_panel/message_text';
-import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
-import { StartChatButton } from '../buttons/start_chat_button';
-import { StopGeneratingButton } from '../buttons/stop_generating_button';
-import { InsightBase } from './insight_base';
 import { MissingCredentialsCallout } from '../missing_credentials_callout';
-import { getConnectorsManagementHref } from '../../utils/get_connectors_management_href';
-import { useObservabilityAIAssistantChatService } from '../../hooks/use_observability_ai_assistant_chat_service';
-import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
-import { useAbortableAsync } from '../../hooks/use_abortable_async';
-import { ObservabilityAIAssistantChatServiceProvider } from '../../context/observability_ai_assistant_chat_service_provider';
+import { InsightBase } from './insight_base';
 
 function ChatContent({
-  title,
-  messages,
+  title: defaultTitle,
+  initialMessages,
   connectorId,
 }: {
   title: string;
-  messages: Message[];
+  initialMessages: Message[];
   connectorId: string;
 }) {
   const chatService = useObservabilityAIAssistantChatService();
@@ -42,51 +44,88 @@ function ChatContent({
   const [loading, setLoading] = useState(false);
   const [subscription, setSubscription] = useState<Subscription | undefined>();
 
+  const [conversationId, setConversationId] = useState<string>();
+
+  const { conversation, displayedMessages, setDisplayedMessages, save, saveTitle } =
+    useConversation({
+      conversationId,
+      connectorId,
+      chatService,
+    });
+
+  const conversationTitle = conversationId
+    ? conversation.value?.conversation.title || ''
+    : defaultTitle;
   const reloadReply = useCallback(() => {
     setLoading(true);
 
+    let lastPendingMessage: PendingMessage | undefined;
+
     const nextSubscription = chatService
-      .chat({ messages, connectorId, function: 'none' })
+      .chat({ messages: initialMessages, connectorId, function: 'none' })
       .subscribe({
         next: (msg) => {
+          lastPendingMessage = msg;
           setPendingMessage(() => msg);
         },
         complete: () => {
+          setDisplayedMessages((prevMessages) =>
+            prevMessages.concat({
+              '@timestamp': new Date().toISOString(),
+              message: {
+                ...lastPendingMessage!.message,
+              },
+            })
+          );
           setLoading(false);
         },
       });
 
     setSubscription(nextSubscription);
-  }, [messages, connectorId, chatService]);
+  }, [initialMessages, setDisplayedMessages, connectorId, chatService]);
 
   useEffect(() => {
     reloadReply();
   }, [reloadReply]);
 
+  useEffect(() => {
+    setDisplayedMessages(initialMessages);
+  }, [initialMessages, setDisplayedMessages]);
+
   const [isOpen, setIsOpen] = useState(false);
 
-  const displayedMessages = useMemo(() => {
+  const messagesWithPending = useMemo(() => {
     return pendingMessage
-      ? messages.concat({
+      ? displayedMessages.concat({
           '@timestamp': new Date().toISOString(),
           message: {
             ...pendingMessage.message,
           },
         })
-      : messages;
-  }, [pendingMessage, messages]);
+      : displayedMessages;
+  }, [pendingMessage, displayedMessages]);
+
+  const lastMessage = last(messagesWithPending);
 
   return (
     <>
       <MessagePanel
-        body={<MessageText content={pendingMessage?.message.content ?? ''} loading={loading} />}
+        body={<MessageText content={lastMessage?.message.content ?? ''} loading={loading} />}
         error={pendingMessage?.error}
         controls={
           loading ? (
             <StopGeneratingButton
               onClick={() => {
                 subscription?.unsubscribe();
                 setLoading(false);
+                setDisplayedMessages((prevMessages) =>
+                  prevMessages.concat({
+                    '@timestamp': new Date().toISOString(),
+                    message: {
+                      ...pendingMessage!.message,
+                    },
+                  })
+                );
                 setPendingMessage((prev) => ({
                   message: {
                     role: MessageRole.Assistant,
@@ -118,12 +157,27 @@ function ChatContent({
         }
       />
       <ChatFlyout
-        title={title}
+        title={conversationTitle}
         isOpen={isOpen}
         onClose={() => {
           setIsOpen(() => false);
         }}
         messages={displayedMessages}
+        conversationId={conversationId}
+        startedFrom="contextualInsight"
+        onChatComplete={(nextMessages) => {
+          save(nextMessages)
+            .then((nextConversation) => {
+              setConversationId(nextConversation.conversation.id);
+            })
+            .catch(() => {});
+        }}
+        onChatUpdate={(nextMessages) => {
+          setDisplayedMessages(nextMessages);
+        }}
+        onChatTitleSave={(newTitle) => {
+          saveTitle(newTitle);
+        }}
       />
     </>
   );
@@ -151,7 +205,11 @@ export function Insight({ messages, title }: { messages: Message[]; title: strin
 
   if (hasOpened && connectors.selectedConnector) {
     children = (
-      <ChatContent title={title} messages={messages} connectorId={connectors.selectedConnector} />
+      <ChatContent
+        title={title}
+        initialMessages={messages}
+        connectorId={connectors.selectedConnector}
+      />
     );
   } else if (!connectors.loading && !connectors.connectors?.length) {
     children = (

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

@@ -55,6 +55,8 @@ export const useJsonEditorModel = ({
 
     if (model === null) {
       model = editor.createModel(initialJsonString, 'json', modelUri);
+    } else {
+      model.setValue(initialJsonString);
     }
 
     return { model, initialJsonString };