[Observability AI Assistant] Assorted bug fixes (#168496)

Resolves https://github.com/elastic/obs-ai-assistant-team/issues/68
Resolves https://github.com/elastic/obs-ai-assistant-team/issues/43
Resolves https://github.com/elastic/obs-ai-assistant-team/issues/95
Resolves https://github.com/elastic/obs-ai-assistant-team/issues/96

## Summary


This fixes a number of issues:

* https://github.com/elastic/obs-ai-assistant-team/issues/68
* https://github.com/elastic/obs-ai-assistant-team/issues/43
* https://github.com/elastic/obs-ai-assistant-team/issues/95
* https://github.com/elastic/obs-ai-assistant-team/issues/96

It also cleans up some DOM nesting issues (using `<p>` tags inside
`<EuiText>` elements so the console shows less noise when in dev mode.

---------

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

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

@@ -170,28 +170,26 @@ export function ChatActionsMenu({
             content: (
               <EuiPanel>
                 <EuiText size="s">
-                  <p>
+                  {i18n.translate(
+                    'xpack.observabilityAiAssistant.chatHeader.actions.knowledgeBase.description.paragraph',
+                    {
+                      defaultMessage:
+                        'Using a knowledge base is optional but improves the experience of using the Assistant significantly.',
+                    }
+                  )}{' '}
+                  <EuiLink
+                    data-test-subj="observabilityAiAssistantChatActionsMenuLearnMoreLink"
+                    external
+                    target="_blank"
+                    href="https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-elser.html"
+                  >
                     {i18n.translate(
-                      'xpack.observabilityAiAssistant.chatHeader.actions.knowledgeBase.description.paragraph',
+                      'xpack.observabilityAiAssistant.chatHeader.actions.knowledgeBase.elser.learnMore',
                       {
-                        defaultMessage:
-                          'Using a knowledge base is optional but improves the experience of using the Assistant significantly.',
+                        defaultMessage: 'Learn more',
                       }
-                    )}{' '}
-                    <EuiLink
-                      data-test-subj="observabilityAiAssistantChatActionsMenuLearnMoreLink"
-                      external
-                      target="_blank"
-                      href="https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-elser.html"
-                    >
-                      {i18n.translate(
-                        'xpack.observabilityAiAssistant.chatHeader.actions.knowledgeBase.elser.learnMore',
-                        {
-                          defaultMessage: 'Learn more',
-                        }
-                      )}
-                    </EuiLink>
-                  </p>
+                    )}
+                  </EuiLink>
                 </EuiText>
 
                 <EuiSpacer size="l" />

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

@@ -6,7 +6,7 @@
  */
 
 import React, { useEffect, useRef, useState } from 'react';
-import { last } from 'lodash';
+import { flatten, last } from 'lodash';
 import {
   EuiFlexGroup,
   EuiFlexItem,
@@ -94,7 +94,7 @@ export function ChatBody({
   let footer: React.ReactNode;
 
   const isLoading = Boolean(
-    connectors.loading || knowledgeBase.status.loading || last(timeline.items)?.loading
+    connectors.loading || knowledgeBase.status.loading || last(flatten(timeline.items))?.loading
   );
 
   const containerClassName = css`

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

@@ -0,0 +1,136 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import React, { useState } from 'react';
+import { i18n } from '@kbn/i18n';
+import { EuiAvatar, EuiButtonIcon, EuiComment, EuiLink } from '@elastic/eui';
+import { css } from '@emotion/css';
+import { ChatItem } from './chat_item';
+import type { ChatTimelineItem, ChatTimelineProps } from './chat_timeline';
+
+const noPanelStyle = css`
+  .euiCommentEvent {
+    border: none;
+  }
+
+  .euiCommentEvent__header {
+    background: transparent;
+    border-block-end: none;
+  }
+
+  .euiCommentEvent__body {
+    display: none;
+  }
+
+  .euiLink {
+    padding: 0 8px;
+  }
+
+  .euiLink:focus {
+    text-decoration: none;
+  }
+
+  .euiLink:hover {
+    text-decoration: underline;
+  }
+`;
+
+const avatarStyle = css`
+  cursor: 'pointer';
+`;
+
+export function ChatConsolidatedItems({
+  consolidatedItem,
+  onFeedback,
+  onRegenerate,
+  onEditSubmit,
+  onStopGenerating,
+  onActionClick,
+}: {
+  consolidatedItem: ChatTimelineItem[];
+  onFeedback: ChatTimelineProps['onFeedback'];
+  onRegenerate: ChatTimelineProps['onRegenerate'];
+  onEditSubmit: ChatTimelineProps['onEdit'];
+  onStopGenerating: ChatTimelineProps['onStopGenerating'];
+  onActionClick: ChatTimelineProps['onActionClick'];
+}) {
+  const [expanded, setExpanded] = useState(false);
+
+  const handleToggleExpand = () => {
+    setExpanded(!expanded);
+  };
+
+  return (
+    <>
+      <EuiComment
+        className={noPanelStyle}
+        timelineAvatar={
+          <EuiAvatar
+            color="subdued"
+            css={avatarStyle}
+            name="inspect"
+            iconType="layers"
+            onClick={handleToggleExpand}
+          />
+        }
+        event={
+          <EuiLink
+            color="subdued"
+            data-test-subj="observabilityAiAssistantChatCollapsedItemsCollapsedItemsLink"
+            onClick={handleToggleExpand}
+          >
+            <em>
+              {!expanded
+                ? i18n.translate('xpack.observabilityAiAssistant.chatCollapsedItems.showEvents', {
+                    defaultMessage: 'Show {count} events',
+                    values: { count: consolidatedItem.length },
+                  })
+                : i18n.translate('xpack.observabilityAiAssistant.chatCollapsedItems.hideEvents', {
+                    defaultMessage: 'Hide {count} events',
+                    values: { count: consolidatedItem.length },
+                  })}
+            </em>
+          </EuiLink>
+        }
+        username=""
+        actions={
+          <EuiButtonIcon
+            aria-label={i18n.translate(
+              'xpack.observabilityAiAssistant.chatCollapsedItems.toggleButtonLabel',
+              {
+                defaultMessage: 'Show / hide items',
+              }
+            )}
+            color="text"
+            data-test-subj="observabilityAiAssistantChatCollapsedItemsButton"
+            iconType={expanded ? 'arrowUp' : 'arrowDown'}
+            onClick={handleToggleExpand}
+          />
+        }
+      />
+
+      {expanded
+        ? consolidatedItem.map((item, index) => (
+            <ChatItem
+              // use index, not id to prevent unmounting of component when message is persisted
+              key={index}
+              {...item}
+              onFeedbackClick={(feedback) => {
+                onFeedback(item, feedback);
+              }}
+              onRegenerateClick={() => {
+                onRegenerate(item);
+              }}
+              onEditSubmit={(message) => onEditSubmit(item, message)}
+              onStopGeneratingClick={onStopGenerating}
+              onActionClick={onActionClick}
+            />
+          ))
+        : null}
+    </>
+  );
+}

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

@@ -101,14 +101,12 @@ export function ChatItemActions({
           closePopover={() => setIsPopoverOpen(undefined)}
         >
           <EuiText size="s">
-            <p>
-              {i18n.translate(
-                'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessageSuccessful',
-                {
-                  defaultMessage: 'Copied message',
-                }
-              )}
-            </p>
+            {i18n.translate(
+              'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessageSuccessful',
+              {
+                defaultMessage: 'Copied message',
+              }
+            )}
           </EuiText>
         </EuiPopover>
       ) : null}

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

@@ -102,7 +102,7 @@ export function ChatPromptEditor({
   }, [functionEditorLineCount, model]);
 
   const handleSubmit = useCallback(async () => {
-    if (loading || !prompt?.trim()) {
+    if (loading || (!prompt?.trim() && !selectedFunctionName)) {
       return;
     }
     const currentPrompt = prompt;

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

@@ -6,16 +6,16 @@
  */
 
 import React, { ReactNode } from 'react';
-import { css } from '@emotion/react';
-import { compact } from 'lodash';
+import { css } from '@emotion/css';
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
 import { ChatItem } from './chat_item';
 import { ChatWelcomePanel } from './chat_welcome_panel';
+import { ChatConsolidatedItems } from './chat_consolidated_items';
 import type { Feedback } from '../feedback_buttons';
-import type { Message } from '../../../common';
-import { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
-import { ChatActionClickHandler } from './types';
+import { type Message } from '../../../common';
+import type { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
+import type { ChatActionClickHandler } from './types';
 
 export interface ChatTimelineItem
   extends Pick<Message['message'], 'role' | 'content' | 'function_call'> {
@@ -38,7 +38,7 @@ export interface ChatTimelineItem
 }
 
 export interface ChatTimelineProps {
-  items: ChatTimelineItem[];
+  items: Array<ChatTimelineItem | ChatTimelineItem[]>;
   knowledgeBase: UseKnowledgeBaseResult;
   onEdit: (item: ChatTimelineItem, message: Message) => Promise<void>;
   onFeedback: (item: ChatTimelineItem, feedback: Feedback) => void;
@@ -56,35 +56,44 @@ export function ChatTimeline({
   onStopGenerating,
   onActionClick,
 }: ChatTimelineProps) {
-  const filteredItems = items.filter((item) => !item.display.hide);
-
   return (
     <EuiCommentList
       css={css`
         padding-bottom: 32px;
       `}
     >
-      {compact(
-        filteredItems.map((item, index) => (
-          <ChatItem
-            // use index, not id to prevent unmounting of component when message is persisted
-            key={index}
-            {...item}
-            onFeedbackClick={(feedback) => {
-              onFeedback(item, feedback);
-            }}
-            onRegenerateClick={() => {
-              onRegenerate(item);
-            }}
-            onEditSubmit={(message) => {
-              return onEdit(item, message);
-            }}
-            onStopGeneratingClick={onStopGenerating}
-            onActionClick={onActionClick}
-          />
-        ))
+      {items.length <= 1 ? (
+        <ChatWelcomePanel knowledgeBase={knowledgeBase} />
+      ) : (
+        items.map((item, index) =>
+          Array.isArray(item) ? (
+            <ChatConsolidatedItems
+              key={index}
+              consolidatedItem={item}
+              onFeedback={onFeedback}
+              onRegenerate={onRegenerate}
+              onEditSubmit={onEdit}
+              onStopGenerating={onStopGenerating}
+              onActionClick={onActionClick}
+            />
+          ) : (
+            <ChatItem
+              // use index, not id to prevent unmounting of component when message is persisted
+              key={index}
+              {...item}
+              onFeedbackClick={(feedback) => {
+                onFeedback(item, feedback);
+              }}
+              onRegenerateClick={() => {
+                onRegenerate(item);
+              }}
+              onEditSubmit={(message) => onEdit(item, message)}
+              onStopGeneratingClick={onStopGenerating}
+              onActionClick={onActionClick}
+            />
+          )
+        )
       )}
-      {filteredItems.length === 1 ? <ChatWelcomePanel knowledgeBase={knowledgeBase} /> : null}
     </EuiCommentList>
   );
 }

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

@@ -36,17 +36,15 @@ export function ChatWelcomePanel({ knowledgeBase }: { knowledgeBase: UseKnowledg
           </h2>
         </EuiTitle>
         <EuiText color="subdued" textAlign="center">
-          <p>
-            {knowledgeBase.status.value?.ready
-              ? i18n.translate('xpack.observabilityAiAssistant.chatWelcomePanel.body.kbReady', {
-                  defaultMessage:
-                    'Keep in mind that Elastic AI Assistant is a technical preview feature. Please provide feedback at any time.',
-                })
-              : i18n.translate('xpack.observabilityAiAssistant.chatWelcomePanel.body.kbNotReady', {
-                  defaultMessage:
-                    'We recommend you enable the knowledge base for a better experience. It will provide the assistant with the ability to learn from your interaction with it. Keep in mind that Elastic AI Assistant is a technical preview feature. Please provide feedback at any time.',
-                })}
-          </p>
+          {knowledgeBase.status.value?.ready
+            ? i18n.translate('xpack.observabilityAiAssistant.chatWelcomePanel.body.kbReady', {
+                defaultMessage:
+                  'Keep in mind that Elastic AI Assistant is a technical preview feature. Please provide feedback at any time.',
+              })
+            : i18n.translate('xpack.observabilityAiAssistant.chatWelcomePanel.body.kbNotReady', {
+                defaultMessage:
+                  'We recommend you enable the knowledge base for a better experience. It will provide the assistant with the ability to learn from your interaction with it. Keep in mind that Elastic AI Assistant is a technical preview feature. Please provide feedback at any time.',
+              })}
         </EuiText>
 
         {!knowledgeBase.status.value?.ready ? (

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

@@ -119,6 +119,12 @@ export function ConversationList({
                         conversation.id
                           ? {
                               iconType: 'trash',
+                              'aria-label': i18n.translate(
+                                'xpack.observabilityAiAssistant.conversationList.deleteConversationIconLabel',
+                                {
+                                  defaultMessage: 'Delete',
+                                }
+                              ),
                               onClick: () => {
                                 onClickDeleteConversation(conversation.id);
                               },

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

@@ -88,17 +88,13 @@ export function FunctionListPopover({
     return (
       <>
         <EuiText size="s">
-          <p>
-            <strong>
-              <EuiHighlight search={searchValue}>{option.label}</EuiHighlight>{' '}
-            </strong>
-          </p>
+          <strong>
+            <EuiHighlight search={searchValue}>{option.label}</EuiHighlight>{' '}
+          </strong>
         </EuiText>
         <EuiSpacer size="xs" />
-        <EuiText size="s">
-          <p style={{ textOverflow: 'ellipsis', overflow: 'hidden' }}>
-            <EuiHighlight search={searchValue}>{option.searchableLabel || ''}</EuiHighlight>
-          </p>
+        <EuiText size="s" style={{ textOverflow: 'ellipsis', overflow: 'hidden' }}>
+          <EuiHighlight search={searchValue}>{option.searchableLabel || ''}</EuiHighlight>
         </EuiText>
       </>
     );

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

@@ -44,11 +44,9 @@ export function IncorrectLicensePanel() {
           <h2>{UPGRADE_LICENSE_TITLE}</h2>
         </EuiTitle>
         <EuiText color="subdued">
-          <p>
-            {i18n.translate('xpack.observabilityAiAssistant.incorrectLicense.body', {
-              defaultMessage: 'You need an Enterprise license to use the Elastic AI Assistant.',
-            })}
-          </p>
+          {i18n.translate('xpack.observabilityAiAssistant.incorrectLicense.body', {
+            defaultMessage: 'You need an Enterprise license to use the Elastic AI Assistant.',
+          })}
         </EuiText>
         <EuiFlexItem>
           <EuiFlexGroup>

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

@@ -44,12 +44,9 @@ export function InitialSetupPanel({
         <EuiSpacer size="s" />
 
         <EuiText color="subdued" size="s">
-          <p>
-            {i18n.translate('xpack.observabilityAiAssistant.initialSetupPanel.title', {
-              defaultMessage:
-                'Start your Al experience with Elastic by completing the steps below.',
-            })}
-          </p>
+          {i18n.translate('xpack.observabilityAiAssistant.initialSetupPanel.title', {
+            defaultMessage: 'Start your Al experience with Elastic by completing the steps below.',
+          })}
         </EuiText>
 
         <EuiSpacer size="l" />
@@ -65,8 +62,8 @@ export function InitialSetupPanel({
                 }
               )}
               description={
-                <EuiText size="s">
-                  <p>
+                <>
+                  <EuiText size="s">
                     {i18n.translate(
                       'xpack.observabilityAiAssistant.initialSetupPanel.knowledgeBase.description.paragraph1',
                       {
@@ -74,16 +71,16 @@ export function InitialSetupPanel({
                           'We recommend you enable the knowledge base for a better experience. It will provide the assistant with the ability to learn from your interaction with it.',
                       }
                     )}
-                  </p>
-                  <p>
+                  </EuiText>
+                  <EuiText size="s">
                     {i18n.translate(
                       'xpack.observabilityAiAssistant.initialSetupPanel.knowledgeBase.description.paragraph2',
                       {
                         defaultMessage: 'This step is optional, you can always do it later.',
                       }
                     )}
-                  </p>
-                </EuiText>
+                  </EuiText>
+                </>
               }
               footer={
                 knowledgeBase.status.value?.ready ? (
@@ -138,17 +135,17 @@ export function InitialSetupPanel({
               )}
               description={
                 !connectors.connectors?.length ? (
-                  <EuiText size="s">
-                    <p>
+                  <>
+                    <EuiText size="s">
                       {i18n.translate(
                         'xpack.observabilityAiAssistant.initialSetupPanel.setupConnector.description1',
                         {
                           defaultMessage: 'Set up an OpenAI connector with your AI provider.',
                         }
                       )}
-                    </p>
+                    </EuiText>
 
-                    <p>
+                    <EuiText size="s">
                       {i18n.translate(
                         'xpack.observabilityAiAssistant.initialSetupPanel.setupConnector.description2',
                         {
@@ -169,18 +166,16 @@ export function InitialSetupPanel({
                         iconType="iInCircle"
                         size="s"
                       />
-                    </p>
-                  </EuiText>
+                    </EuiText>
+                  </>
                 ) : connectors.connectors.length && !connectors.selectedConnector ? (
                   <EuiText size="s">
-                    <p>
-                      {i18n.translate(
-                        'xpack.observabilityAiAssistant.initialSetupPanel.setupConnector.description',
-                        {
-                          defaultMessage: 'Please select a provider.',
-                        }
-                      )}
-                    </p>
+                    {i18n.translate(
+                      'xpack.observabilityAiAssistant.initialSetupPanel.setupConnector.description',
+                      {
+                        defaultMessage: 'Please select a provider.',
+                      }
+                    )}
                   </EuiText>
                 ) : (
                   ''
@@ -212,12 +207,10 @@ export function InitialSetupPanel({
         <EuiSpacer size="xxl" />
 
         <EuiText color="subdued" size="s">
-          <p>
-            {i18n.translate('xpack.observabilityAiAssistant.initialSetupPanel.disclaimer', {
-              defaultMessage:
-                'The AI provider that is configured may collect telemetry when using the Elastic AI Assistant. Contact your AI provider for information on how data is collected.',
-            })}
-          </p>
+          {i18n.translate('xpack.observabilityAiAssistant.initialSetupPanel.disclaimer', {
+            defaultMessage:
+              'The AI provider that is configured may collect telemetry when using the Elastic AI Assistant. Contact your AI provider for information on how data is collected.',
+          })}
         </EuiText>
       </EuiPanel>
     </>

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

@@ -14,6 +14,7 @@ import {
 } from '@testing-library/react-hooks';
 import { BehaviorSubject, Subject } from 'rxjs';
 import { MessageRole } from '../../common';
+import { ChatTimelineItem } from '../components/chat/chat_timeline';
 import type { PendingMessage } from '../types';
 import { useTimeline, UseTimelineResult } from './use_timeline';
 
@@ -83,13 +84,35 @@ describe('useTimeline', () => {
             {
               message: {
                 role: MessageRole.User,
-                content: 'Hello',
+                content: 'hello',
               },
             },
             {
               message: {
                 role: MessageRole.Assistant,
-                content: 'Goodbye',
+                content: '',
+                function_call: {
+                  name: 'recall',
+                  trigger: MessageRole.User,
+                },
+              },
+            },
+            {
+              message: {
+                name: 'recall',
+                role: MessageRole.User,
+                content: '',
+              },
+            },
+            {
+              message: {
+                content: 'goodbye',
+                function_call: {
+                  name: '',
+                  arguments: '',
+                  trigger: MessageRole.Assistant,
+                },
+                role: MessageRole.Assistant,
               },
             },
           ],
@@ -98,48 +121,98 @@ describe('useTimeline', () => {
           },
           chatService: {
             chat: () => {},
+            hasRenderFunction: () => {},
           },
         } as unknown as HookProps,
       });
     });
     it('renders the correct timeline items', () => {
-      expect(hookResult.result.current.items.length).toEqual(3);
+      expect(hookResult.result.current.items.length).toEqual(4);
 
       expect(hookResult.result.current.items[1]).toEqual({
-        actions: {
-          canCopy: true,
-          canEdit: true,
-          canRegenerate: false,
-          canGiveFeedback: false,
-        },
-        display: {
-          collapsed: false,
-          hide: false,
-        },
-        role: MessageRole.User,
-        content: 'Hello',
-        loading: false,
+        actions: { canCopy: true, canEdit: true, canGiveFeedback: false, canRegenerate: false },
+        content: 'hello',
+        currentUser: undefined,
+        display: { collapsed: false, hide: false },
+        element: undefined,
+        function_call: undefined,
         id: expect.any(String),
+        loading: false,
+        role: MessageRole.User,
         title: '',
       });
 
-      expect(hookResult.result.current.items[2]).toEqual({
-        display: {
-          collapsed: false,
-          hide: false,
+      expect(hookResult.result.current.items[3]).toEqual({
+        actions: { canCopy: true, canEdit: false, canGiveFeedback: false, canRegenerate: true },
+        content: 'goodbye',
+        currentUser: undefined,
+        display: { collapsed: false, hide: false },
+        element: undefined,
+        function_call: {
+          arguments: '',
+          name: '',
+          trigger: MessageRole.Assistant,
         },
-        actions: {
-          canCopy: true,
-          canEdit: false,
-          canRegenerate: true,
-          canGiveFeedback: false,
-        },
-        role: MessageRole.Assistant,
-        content: 'Goodbye',
-        loading: false,
         id: expect.any(String),
+        loading: false,
+        role: MessageRole.Assistant,
         title: '',
       });
+
+      // Items that are function calls are collapsed into an array.
+
+      // 'title' is a <FormattedMessage /> component. This throws Jest for a loop.
+      const collapsedItemsWithoutTitle = (
+        hookResult.result.current.items[2] as ChatTimelineItem[]
+      ).map(({ title, ...rest }) => rest);
+
+      expect(collapsedItemsWithoutTitle).toEqual([
+        {
+          display: {
+            collapsed: true,
+            hide: false,
+          },
+          actions: {
+            canCopy: true,
+            canEdit: true,
+            canRegenerate: false,
+            canGiveFeedback: false,
+          },
+          currentUser: undefined,
+          function_call: {
+            name: 'recall',
+            trigger: MessageRole.User,
+          },
+          role: MessageRole.User,
+          content: `\`\`\`
+{
+  \"name\": \"recall\"
+}
+\`\`\``,
+          loading: false,
+          id: expect.any(String),
+        },
+        {
+          display: {
+            collapsed: true,
+            hide: false,
+          },
+          actions: {
+            canCopy: true,
+            canEdit: false,
+            canRegenerate: false,
+            canGiveFeedback: false,
+          },
+          currentUser: undefined,
+          function_call: undefined,
+          role: MessageRole.User,
+          content: `\`\`\`
+{}
+\`\`\``,
+          loading: false,
+          id: expect.any(String),
+        },
+      ]);
     });
   });
 
@@ -197,10 +270,16 @@ describe('useTimeline', () => {
       });
 
       it('adds two items of which the last one is loading', async () => {
-        expect(hookResult.result.current.items[0].role).toEqual(MessageRole.User);
-        expect(hookResult.result.current.items[1].role).toEqual(MessageRole.User);
+        expect((hookResult.result.current.items[0] as ChatTimelineItem).role).toEqual(
+          MessageRole.User
+        );
+        expect((hookResult.result.current.items[1] as ChatTimelineItem).role).toEqual(
+          MessageRole.User
+        );
 
-        expect(hookResult.result.current.items[2].role).toEqual(MessageRole.Assistant);
+        expect((hookResult.result.current.items[2] as ChatTimelineItem).role).toEqual(
+          MessageRole.Assistant
+        );
 
         expect(hookResult.result.current.items[1]).toMatchObject({
           role: MessageRole.User,
@@ -303,7 +382,9 @@ describe('useTimeline', () => {
         describe('and it being regenerated', () => {
           beforeEach(() => {
             act(() => {
-              hookResult.result.current.onRegenerate(hookResult.result.current.items[2]);
+              hookResult.result.current.onRegenerate(
+                hookResult.result.current.items[2] as ChatTimelineItem
+              );
               subject.next({ message: { role: MessageRole.Assistant, content: '' } });
             });
           });
@@ -335,7 +416,9 @@ describe('useTimeline', () => {
               });
 
               act(() => {
-                hookResult.result.current.onRegenerate(hookResult.result.current.items[2]);
+                hookResult.result.current.onRegenerate(
+                  hookResult.result.current.items[2] as ChatTimelineItem
+                );
               });
             });
 
@@ -445,7 +528,7 @@ describe('useTimeline', () => {
                 '@timestamp': expect.any(String),
                 message: {
                   content: 'Hello',
-                  role: 'user',
+                  role: MessageRole.User,
                 },
               },
             ],

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

@@ -19,7 +19,7 @@ import {
   type Message,
 } from '../../common/types';
 import type { ChatPromptEditorProps } from '../components/chat/chat_prompt_editor';
-import type { ChatTimelineProps } from '../components/chat/chat_timeline';
+import type { ChatTimelineItem, ChatTimelineProps } from '../components/chat/chat_timeline';
 import { ChatActionClickType } from '../components/chat/types';
 import { EMPTY_CONVERSATION_TITLE } from '../i18n';
 import type { ObservabilityAIAssistantChatService, PendingMessage } from '../types';
@@ -101,6 +101,7 @@ export function useTimeline({
   const [isFunctionLoading, setIsFunctionLoading] = useState(false);
 
   const prevConversationId = usePrevious(conversationId);
+
   useEffect(() => {
     if (prevConversationId !== conversationId && pendingMessage?.error) {
       setPendingMessage(undefined);
@@ -257,26 +258,27 @@ export function useTimeline({
     });
   }
 
-  const items = useMemo(() => {
-    if (pendingMessage) {
+  const itemsWithAddedLoadingStates = useMemo(() => {
+    // While we're loading we add an empty loading chat item:
+    if (pendingMessage || isFunctionLoading) {
       const nextItems = conversationItems.concat({
         id: '',
         actions: {
           canCopy: true,
           canEdit: false,
           canGiveFeedback: false,
-          canRegenerate: pendingMessage.aborted || !!pendingMessage.error,
+          canRegenerate: pendingMessage?.aborted || !!pendingMessage?.error,
         },
         display: {
           collapsed: false,
-          hide: pendingMessage.message.role === MessageRole.System,
+          hide: pendingMessage?.message.role === MessageRole.System,
         },
-        content: pendingMessage.message.content,
+        content: pendingMessage?.message.content,
         currentUser,
-        error: pendingMessage.error,
-        function_call: pendingMessage.message.function_call,
-        loading: !pendingMessage.aborted && !pendingMessage.error,
-        role: pendingMessage.message.role,
+        error: pendingMessage?.error,
+        function_call: pendingMessage?.message.function_call,
+        loading: !pendingMessage?.aborted && !pendingMessage?.error,
+        role: pendingMessage?.message.role || MessageRole.Assistant,
         title: '',
       });
 
@@ -288,6 +290,7 @@ export function useTimeline({
     }
 
     return conversationItems.map((item, index) => {
+      // When we're done loading we remove the placeholder item again
       if (index < conversationItems.length - 1) {
         return item;
       }
@@ -298,6 +301,29 @@ export function useTimeline({
     });
   }, [conversationItems, pendingMessage, currentUser, isFunctionLoading]);
 
+  const items = useMemo(() => {
+    const consolidatedChatItems: Array<ChatTimelineItem | ChatTimelineItem[]> = [];
+    let currentGroup: ChatTimelineItem[] | null = null;
+
+    for (const item of itemsWithAddedLoadingStates) {
+      if (item.display.hide || !item) continue;
+
+      if (item.display.collapsed) {
+        if (currentGroup) {
+          currentGroup.push(item);
+        } else {
+          currentGroup = [item];
+          consolidatedChatItems.push(currentGroup);
+        }
+      } else {
+        consolidatedChatItems.push(item);
+        currentGroup = null;
+      }
+    }
+
+    return consolidatedChatItems;
+  }, [itemsWithAddedLoadingStates]);
+
   useEffect(() => {
     return () => {
       subscription?.unsubscribe();

x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx

@@ -92,7 +92,7 @@ export function getTimelineItemsfromConversation({
           ? messages[index - 1].message.function_call
           : undefined;
 
-      const role = message.message.function_call?.trigger || message.message.role;
+      let role = message.message.function_call?.trigger || message.message.role;
 
       const actions = {
         canCopy: false,
@@ -159,8 +159,12 @@ export function getTimelineItemsfromConversation({
 
             content = !element ? convertMessageToMarkdownCodeBlock(message.message) : undefined;
 
+            if (prevFunctionCall?.trigger === MessageRole.Assistant) {
+              role = MessageRole.Assistant;
+            }
+
             actions.canEdit = false;
-            display.collapsed = !isError && !element;
+            display.collapsed = !element;
           } else if (message.message.function_call) {
             // User suggested a function
             title = (