diff --git a/x-pack/plugins/observability_ai_assistant/common/types.ts b/x-pack/plugins/observability_ai_assistant/common/types.ts
index 59e0b85cdc5d5..8ef052cb8e517 100644
--- a/x-pack/plugins/observability_ai_assistant/common/types.ts
+++ b/x-pack/plugins/observability_ai_assistant/common/types.ts
@@ -28,7 +28,7 @@ export interface Message {
     function_call?: {
       name: string;
       arguments?: string;
-      trigger: MessageRole.Assistant | MessageRole.User | MessageRole.Elastic;
+      trigger: MessageRole;
     };
     data?: string;
   };
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx
index a3c8c8f42437a..f7ef7b338009f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item.tsx
@@ -32,6 +32,10 @@ export interface ChatItemProps extends ChatTimelineItem {
 }
 
 const normalMessageClassName = css`
+  .euiCommentEvent__header {
+    padding: 4px 8px;
+  }
+
   .euiCommentEvent__body {
     padding: 0;
   }
@@ -42,17 +46,17 @@ const normalMessageClassName = css`
 `;
 
 const noPanelMessageClassName = css`
+  .euiCommentEvent {
+    border: none;
+  }
+
   .euiCommentEvent__header {
     background: transparent;
     border-block-end: none;
   }
 
   .euiCommentEvent__body {
-    padding: 0;
-  }
-
-  .euiCommentEvent {
-    border: none;
+    display: none;
   }
 `;
 
@@ -89,6 +93,10 @@ export function ChatItem({
   const actions = [canCopy, collapsed, canCopy].filter(Boolean);
 
   const noBodyMessageClassName = css`
+    .euiCommentEvent__header {
+      padding: 4px 8px;
+    }
+
     .euiCommentEvent__body {
       padding: 0;
       height: ${expanded ? 'fit-content' : '0px'};
@@ -105,8 +113,8 @@ export function ChatItem({
   };
 
   const handleToggleEdit = () => {
-    if (collapsed) {
-      setExpanded(false);
+    if (collapsed && !expanded) {
+      setExpanded(true);
     }
     setEditing(!editing);
   };
@@ -155,9 +163,10 @@ export function ChatItem({
       actions={
         <ChatItemActions
           canCopy={canCopy}
-          collapsed={collapsed}
           canEdit={canEdit}
-          isCollapsed={expanded}
+          collapsed={collapsed}
+          editing={editing}
+          expanded={expanded}
           onCopyToClipboard={handleCopyToClipboard}
           onToggleEdit={handleToggleEdit}
           onToggleExpand={handleToggleExpand}
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_actions.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_actions.tsx
index 7190ccca73dcc..79240a03bb314 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_actions.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_actions.tsx
@@ -9,26 +9,21 @@ import React, { useEffect, useState } from 'react';
 import { i18n } from '@kbn/i18n';
 import { EuiButtonIcon, EuiPopover, EuiText } from '@elastic/eui';
 
-export interface ChatItemAction {
-  id: string;
-  label: string;
-  icon: string;
-  handler: () => void;
-}
-
 export function ChatItemActions({
+  canCopy,
   canEdit,
   collapsed,
-  canCopy,
-  isCollapsed,
+  editing,
+  expanded,
   onToggleEdit,
   onToggleExpand,
   onCopyToClipboard,
 }: {
+  canCopy: boolean;
   canEdit: boolean;
   collapsed: boolean;
-  canCopy: boolean;
-  isCollapsed: boolean;
+  editing: boolean;
+  expanded: boolean;
   onToggleEdit: () => void;
   onToggleExpand: () => void;
   onCopyToClipboard: () => void;
@@ -47,85 +42,73 @@ export function ChatItemActions({
     };
   }, [isPopoverOpen]);
 
-  const actions: ChatItemAction[] = [
-    ...(canEdit
-      ? [
-          {
-            id: 'edit',
-            icon: 'documentEdit',
-            label: '',
-            handler: () => {
-              onToggleEdit();
-            },
-          },
-        ]
-      : []),
-    ...(collapsed
-      ? [
-          {
-            id: 'expand',
-            icon: isCollapsed ? 'eyeClosed' : 'eye',
-            label: '',
-            handler: () => {
-              onToggleExpand();
-            },
-          },
-        ]
-      : []),
-    ...(canCopy
-      ? [
-          {
-            id: 'copy',
-            icon: 'copyClipboard',
-            label: i18n.translate(
-              'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessage',
-              {
-                defaultMessage: 'Copied message',
-              }
-            ),
-            handler: () => {
-              onCopyToClipboard();
-            },
-          },
-        ]
-      : []),
-  ];
   return (
     <>
-      {actions.map(({ id, icon, label, handler }) =>
-        label ? (
-          <EuiPopover
-            key={id}
-            button={
-              <EuiButtonIcon
-                aria-label={label}
-                key={id}
-                iconType={icon}
-                onClick={() => {
-                  setIsPopoverOpen(id);
-                  handler();
-                }}
-                color="text"
-              />
+      {canEdit ? (
+        <EuiButtonIcon
+          aria-label={i18n.translate(
+            'xpack.observabilityAiAssistant.chatTimeline.actions.editPrompt',
+            {
+              defaultMessage: 'Edit prompt',
+            }
+          )}
+          color="text"
+          display={editing ? 'fill' : 'empty'}
+          iconType="documentEdit"
+          onClick={onToggleEdit}
+        />
+      ) : null}
+
+      {collapsed ? (
+        <EuiButtonIcon
+          aria-label={i18n.translate(
+            'xpack.observabilityAiAssistant.chatTimeline.actions.inspectPrompt',
+            {
+              defaultMessage: 'Inspect prompt',
             }
-            isOpen={isPopoverOpen === id}
-            closePopover={() => setIsPopoverOpen(undefined)}
-            panelPaddingSize="s"
-          >
-            <EuiText size="s">
-              <p>{label}</p>
-            </EuiText>
-          </EuiPopover>
-        ) : (
-          <EuiButtonIcon
-            aria-label={label}
-            key={id}
-            iconType={icon}
-            onClick={handler}
-            color="text"
-          />
-        )
-      )}
+          )}
+          color="text"
+          display={expanded ? 'fill' : 'empty'}
+          iconType={expanded ? 'eyeClosed' : 'eye'}
+          onClick={onToggleExpand}
+        />
+      ) : null}
+
+      {canCopy ? (
+        <EuiPopover
+          button={
+            <EuiButtonIcon
+              aria-label={i18n.translate(
+                'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessage',
+                {
+                  defaultMessage: 'Copy message',
+                }
+              )}
+              color="text"
+              iconType="copyClipboard"
+              display={isPopoverOpen === 'copy' ? 'fill' : 'empty'}
+              onClick={() => {
+                setIsPopoverOpen('copy');
+                onCopyToClipboard();
+              }}
+            />
+          }
+          isOpen={isPopoverOpen === 'copy'}
+          panelPaddingSize="s"
+          closePopover={() => setIsPopoverOpen(undefined)}
+        >
+          <EuiText size="s">
+            <p>
+              {i18n.translate(
+                'xpack.observabilityAiAssistant.chatTimeline.actions.copyMessageSuccessful',
+                {
+                  defaultMessage: 'Copied message',
+                }
+              )}
+            </p>
+          </EuiText>
+        </EuiPopover>
+      ) : null}
     </>
   );
 }
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
index d1e08e5c4a698..827251990f429 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor.tsx
@@ -15,6 +15,7 @@ import {
   EuiSpacer,
   EuiPanel,
   keys,
+  EuiFocusTrap,
 } from '@elastic/eui';
 import { CodeEditor } from '@kbn/kibana-react-plugin/public';
 import { i18n } from '@kbn/i18n';
@@ -43,6 +44,8 @@ export function ChatPromptEditor({
   const { getFunctions } = useObservabilityAIAssistant();
   const functions = getFunctions();
 
+  const isFocusTrapEnabled = Boolean(initialPrompt);
+
   const [prompt, setPrompt] = useState(initialPrompt);
 
   const [selectedFunctionName, setSelectedFunctionName] = useState<string | undefined>(
@@ -103,7 +106,8 @@ export function ChatPromptEditor({
         await onSubmit({
           '@timestamp': new Date().toISOString(),
           message: {
-            role: MessageRole.Function,
+            role: MessageRole.Assistant,
+            content: '',
             function_call: {
               name: selectedFunctionName,
               trigger: MessageRole.User,
@@ -111,6 +115,9 @@ export function ChatPromptEditor({
             },
           },
         });
+
+        setFunctionPayload(undefined);
+        setSelectedFunctionName(undefined);
       } else {
         await onSubmit({
           '@timestamp': new Date().toISOString(),
@@ -125,9 +132,9 @@ export function ChatPromptEditor({
 
   useEffect(() => {
     const keyboardListener = (event: KeyboardEvent) => {
-      if (!event.shiftKey && event.key === keys.ENTER) {
-        handleSubmit();
+      if (!event.shiftKey && event.key === keys.ENTER && (prompt || selectedFunctionName)) {
         event.preventDefault();
+        handleSubmit();
       }
     };
 
@@ -136,7 +143,7 @@ export function ChatPromptEditor({
     return () => {
       window.removeEventListener('keypress', keyboardListener);
     };
-  }, [handleSubmit]);
+  }, [handleSubmit, prompt, selectedFunctionName]);
 
   useEffect(() => {
     const textarea = textAreaRef.current;
@@ -152,109 +159,114 @@ export function ChatPromptEditor({
   });
 
   return (
-    <EuiFlexGroup gutterSize="s" responsive={false}>
-      <EuiFlexItem grow>
-        <EuiFlexGroup direction="column" gutterSize="s">
-          <EuiFlexItem>
-            <EuiFlexGroup responsive={false}>
-              <EuiFlexItem grow>
-                <FunctionListPopover
-                  functions={functions}
-                  selectedFunctionName={selectedFunctionName}
-                  onSelectFunction={handleSelectFunction}
-                />
-              </EuiFlexItem>
-              <EuiFlexItem grow={false}>
-                {selectedFunctionName ? (
-                  <EuiButtonEmpty
-                    iconType="cross"
-                    iconSide="right"
-                    size="xs"
-                    onClick={handleClearSelection}
-                  >
-                    {i18n.translate('xpack.observabilityAiAssistant.prompt.emptySelection', {
-                      defaultMessage: 'Empty selection',
-                    })}
-                  </EuiButtonEmpty>
-                ) : null}
-              </EuiFlexItem>
-            </EuiFlexGroup>
-          </EuiFlexItem>
-          <EuiFlexItem>
-            {selectedFunctionName ? (
-              <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
-                <CodeEditor
-                  aria-label="payloadEditor"
+    <EuiFocusTrap disabled={!isFocusTrapEnabled}>
+      <EuiFlexGroup gutterSize="s" responsive={false}>
+        <EuiFlexItem grow>
+          <EuiFlexGroup direction="column" gutterSize="s">
+            <EuiFlexItem>
+              <EuiFlexGroup responsive={false}>
+                <EuiFlexItem grow>
+                  <FunctionListPopover
+                    functions={functions}
+                    selectedFunctionName={selectedFunctionName}
+                    onSelectFunction={handleSelectFunction}
+                  />
+                </EuiFlexItem>
+                <EuiFlexItem grow={false}>
+                  {selectedFunctionName ? (
+                    <EuiButtonEmpty
+                      iconType="cross"
+                      iconSide="right"
+                      size="xs"
+                      onClick={handleClearSelection}
+                    >
+                      {i18n.translate('xpack.observabilityAiAssistant.prompt.emptySelection', {
+                        defaultMessage: 'Empty selection',
+                      })}
+                    </EuiButtonEmpty>
+                  ) : null}
+                </EuiFlexItem>
+              </EuiFlexGroup>
+            </EuiFlexItem>
+            <EuiFlexItem>
+              {selectedFunctionName ? (
+                <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
+                  <CodeEditor
+                    aria-label="payloadEditor"
+                    fullWidth
+                    height="120px"
+                    languageId="json"
+                    isCopyable
+                    languageConfiguration={{
+                      autoClosingPairs: [
+                        {
+                          open: '{',
+                          close: '}',
+                        },
+                      ],
+                    }}
+                    editorDidMount={(editor) => {
+                      editor.focus();
+                    }}
+                    options={{
+                      accessibilitySupport: 'off',
+                      acceptSuggestionOnEnter: 'on',
+                      automaticLayout: true,
+                      autoClosingQuotes: 'always',
+                      autoIndent: 'full',
+                      contextmenu: true,
+                      fontSize: 12,
+                      formatOnPaste: true,
+                      formatOnType: true,
+                      inlineHints: { enabled: true },
+                      lineNumbers: 'on',
+                      minimap: { enabled: false },
+                      model,
+                      overviewRulerBorder: false,
+                      quickSuggestions: true,
+                      scrollbar: { alwaysConsumeMouseWheel: false },
+                      scrollBeyondLastLine: false,
+                      suggestOnTriggerCharacters: true,
+                      tabSize: 2,
+                      wordWrap: 'on',
+                      wrappingIndent: 'indent',
+                    }}
+                    transparentBackground
+                    value={functionPayload || ''}
+                    onChange={handleChangeFunctionPayload}
+                  />
+                </EuiPanel>
+              ) : (
+                <EuiTextArea
                   fullWidth
-                  height="120px"
-                  languageId="json"
-                  value={functionPayload || ''}
-                  onChange={handleChangeFunctionPayload}
-                  isCopyable
-                  languageConfiguration={{
-                    autoClosingPairs: [
-                      {
-                        open: '{',
-                        close: '}',
-                      },
-                    ],
-                  }}
-                  options={{
-                    accessibilitySupport: 'off',
-                    acceptSuggestionOnEnter: 'on',
-                    automaticLayout: true,
-                    autoClosingQuotes: 'always',
-                    autoIndent: 'full',
-                    contextmenu: true,
-                    fontSize: 12,
-                    formatOnPaste: true,
-                    formatOnType: true,
-                    inlineHints: { enabled: true },
-                    lineNumbers: 'on',
-                    minimap: { enabled: false },
-                    model,
-                    overviewRulerBorder: false,
-                    quickSuggestions: true,
-                    scrollbar: { alwaysConsumeMouseWheel: false },
-                    scrollBeyondLastLine: false,
-                    suggestOnTriggerCharacters: true,
-                    tabSize: 2,
-                    wordWrap: 'on',
-                    wrappingIndent: 'indent',
-                  }}
-                  transparentBackground
+                  inputRef={textAreaRef}
+                  placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
+                    defaultMessage: 'Press ‘$’ for function recommendations',
+                  })}
+                  resize="vertical"
+                  rows={1}
+                  value={prompt}
+                  onChange={handleChange}
                 />
-              </EuiPanel>
-            ) : (
-              <EuiTextArea
-                resize="vertical"
-                rows={1}
-                fullWidth
-                placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
-                  defaultMessage: 'Press ‘$’ for function recommendations',
-                })}
-                value={prompt}
-                inputRef={textAreaRef}
-                onChange={handleChange}
-              />
-            )}
-          </EuiFlexItem>
-        </EuiFlexGroup>
-      </EuiFlexItem>
-      <EuiFlexItem grow={false}>
-        <EuiSpacer size="xl" />
-        <EuiButtonIcon
-          aria-label="Submit"
-          isLoading={loading}
-          disabled={selectedFunctionName ? false : !prompt || loading || disabled}
-          display={
-            selectedFunctionName ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
-          }
-          iconType="kqlFunction"
-          size="m"
-          onClick={handleSubmit}
-        />
-      </EuiFlexItem>
-    </EuiFlexGroup>
+              )}
+            </EuiFlexItem>
+          </EuiFlexGroup>
+        </EuiFlexItem>
+        <EuiFlexItem grow={false}>
+          <EuiSpacer size="xl" />
+          <EuiButtonIcon
+            aria-label="Submit"
+            disabled={selectedFunctionName ? false : !prompt || loading || disabled}
+            display={
+              selectedFunctionName ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
+            }
+            iconType="kqlFunction"
+            isLoading={loading}
+            size="m"
+            onClick={handleSubmit}
+          />
+        </EuiFlexItem>
+      </EuiFlexGroup>
+    </EuiFocusTrap>
   );
 }
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx
index 45d1dac788871..c953ce9186795 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_timeline.tsx
@@ -7,7 +7,7 @@
 
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import React from 'react';
+import React, { ReactNode } from 'react';
 import { type Message } from '../../../common';
 import type { Feedback } from '../feedback_buttons';
 import { ChatItem } from './chat_item';
@@ -16,7 +16,7 @@ export interface ChatTimelineItem
   extends Pick<Message, '@timestamp'>,
     Pick<Message['message'], 'role' | 'content' | 'function_call'> {
   id: string;
-  title: string;
+  title: ReactNode;
   loading: boolean;
   canCopy: boolean;
   canEdit: boolean;
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts
index 2f550fe1f965c..56d1d35c9178f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/hooks/use_timeline.ts
@@ -17,7 +17,7 @@ import type { ObservabilityAIAssistantService, PendingMessage } from '../types';
 import { getTimelineItemsfromConversation } from '../utils/get_timeline_items_from_conversation';
 import type { UseGenAIConnectorsResult } from './use_genai_connectors';
 import { getAssistantSetupMessage } from '../service/get_assistant_setup_message';
-
+import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
 export function createNewConversation(): ConversationCreateRequest {
   return {
     '@timestamp': new Date().toISOString(),
@@ -54,6 +54,9 @@ export function useTimeline({
   onChatComplete: (messages: Message[]) => void;
   knowledgeBaseAvailable: boolean;
 }): UseTimelineResult {
+  const { getFunctions } = useObservabilityAIAssistant();
+  const functions = getFunctions();
+
   const connectorId = connectors.selectedConnector;
 
   const hasConnector = !!connectorId;
@@ -63,10 +66,11 @@ export function useTimeline({
       messages,
       currentUser,
       hasConnector,
+      functions,
     });
 
     return items;
-  }, [messages, currentUser, hasConnector]);
+  }, [messages, currentUser, hasConnector, functions]);
 
   const [subscription, setSubscription] = useState<Subscription | undefined>();
 
diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx b/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx
index 8115970943135..6bd60f081445f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/utils/get_timeline_items_from_conversation.tsx
@@ -8,12 +8,14 @@ import React from 'react';
 import { v4 } from 'uuid';
 import { i18n } from '@kbn/i18n';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import { type Message, MessageRole } from '../../common';
-import type { ChatTimelineItem } from '../components/chat/chat_timeline';
 import { RenderFunction } from '../components/render_function';
+import type { ChatTimelineItem } from '../components/chat/chat_timeline';
+import { type Message, MessageRole } from '../../common';
+import type { FunctionDefinition } from '../../common/types';
 
 function convertFunctionParamsToMarkdownCodeBlock(object: Record<string, string | number>) {
-  return `\`\`\`
+  return `
+\`\`\`
 ${JSON.stringify(object, null, 4)}
 \`\`\``;
 }
@@ -22,10 +24,12 @@ export function getTimelineItemsfromConversation({
   currentUser,
   messages,
   hasConnector,
+  functions,
 }: {
   currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
   messages: Message[];
   hasConnector: boolean;
+  functions: FunctionDefinition[];
 }): ChatTimelineItem[] {
   return [
     {
@@ -45,90 +49,132 @@ export function getTimelineItemsfromConversation({
       }),
     },
     ...messages.map((message, index) => {
-      const hasFunction = !!message.message.function_call?.name;
-      const isSystemPrompt = message.message.role === MessageRole.System;
+      const id = v4();
 
-      let title: string;
+      let title: string = '';
       let content: string | undefined;
       let element: React.ReactNode | undefined;
+
+      const role = message.message.name ? message.message.role : message.message.role;
+      const functionCall =
+        message.message.name && messages[index - 1] && messages[index - 1].message.function_call
+          ? messages[index - 1].message.function_call
+          : message.message.function_call;
+
+      let canCopy: boolean = false;
+      let canEdit: boolean = false;
+      let canGiveFeedback: boolean = false;
+      let canRegenerate: boolean = false;
       let collapsed: boolean = false;
+      let hide: boolean = false;
+
+      switch (role) {
+        case MessageRole.System:
+          hide = true;
+          break;
+
+        case MessageRole.User:
+          // User executed a function:
+          if (functionCall) {
+            title = i18n.translate('xpack.observabilityAiAssistant.executedFunctionEvent', {
+              defaultMessage: 'executed the function {functionName}',
+              values: {
+                functionName: functionCall.name,
+              },
+            });
+
+            content = convertFunctionParamsToMarkdownCodeBlock({
+              name: functionCall.name,
+              arguments: JSON.parse(functionCall.arguments || '{}'),
+            });
+
+            const fn = functions.find((func) => func.options.name === message.message.name);
+
+            element = fn?.render ? (
+              <RenderFunction
+                name={functionCall.name}
+                arguments={functionCall?.arguments}
+                response={message.message}
+              />
+            ) : null;
+
+            canCopy = true;
+            canEdit = hasConnector;
+            canGiveFeedback = true;
+            canRegenerate = hasConnector;
+            collapsed = !Boolean(element);
+            hide = false;
+          } else {
+            // is a prompt by the user
+            title = '';
+            content = message.message.content;
 
-      if (hasFunction) {
-        title = i18n.translate('xpack.observabilityAiAssistant.suggestedFunctionEvent', {
-          defaultMessage: 'suggested to use function {functionName}',
-          values: {
-            functionName: message.message.function_call!.name,
-          },
-        });
-
-        content = convertFunctionParamsToMarkdownCodeBlock({
-          name: message.message.function_call!.name,
-          arguments: JSON.parse(message.message.function_call?.arguments || '{}'),
-        });
-
-        collapsed = true;
-      } else if (isSystemPrompt) {
-        title = i18n.translate('xpack.observabilityAiAssistant.addedSystemPromptEvent', {
-          defaultMessage: 'added a prompt',
-        });
-        content = '';
-        collapsed = true;
-      } else if (message.message.name) {
-        const prevMessage = messages[index - 1];
-        if (!prevMessage || !prevMessage.message.function_call) {
-          throw new Error('Could not find preceding message with function_call');
-        }
-
-        title = i18n.translate('xpack.observabilityAiAssistant.executedFunctionEvent', {
-          defaultMessage: 'executed the function {functionName}',
-          values: {
-            functionName: prevMessage.message.function_call!.name,
-          },
-        });
-
-        content = convertFunctionParamsToMarkdownCodeBlock(
-          JSON.parse(message.message.content || '{}')
-        );
-
-        element = (
-          <RenderFunction
-            name={message.message.name}
-            arguments={prevMessage.message.function_call.arguments}
-            response={message.message}
-          />
-        );
-        collapsed = true;
-      } else {
-        title = '';
-        content = message.message.content;
-        collapsed = false;
+            canCopy = true;
+            canEdit = hasConnector;
+            canGiveFeedback = false;
+            canRegenerate = false;
+            collapsed = false;
+            hide = false;
+          }
+
+          break;
+
+        case MessageRole.Assistant:
+          // is a function suggestion by the assistant
+          if (!!functionCall?.name) {
+            title = i18n.translate('xpack.observabilityAiAssistant.suggestedFunctionEvent', {
+              defaultMessage: 'suggested to use function {functionName}',
+              values: {
+                functionName: functionCall!.name,
+              },
+            });
+            content =
+              i18n.translate('xpack.observabilityAiAssistant.responseWas', {
+                defaultMessage: 'Suggested the payload: ',
+              }) +
+              convertFunctionParamsToMarkdownCodeBlock({
+                name: functionCall!.name,
+                arguments: JSON.parse(functionCall?.arguments || '{}'),
+              });
+
+            canCopy = true;
+            canEdit = false;
+            canGiveFeedback = true;
+            canRegenerate = false;
+            collapsed = true;
+            hide = false;
+          } else {
+            // is an assistant response
+            title = '';
+            content = message.message.content;
+
+            canCopy = true;
+            canEdit = false;
+            canGiveFeedback = true;
+            canRegenerate = hasConnector;
+            collapsed = false;
+            hide = false;
+          }
+          break;
       }
 
-      const props = {
-        id: v4(),
+      return {
+        id,
         '@timestamp': message['@timestamp'],
-        canCopy: true,
-        canEdit: hasConnector && (message.message.role === MessageRole.User || hasFunction),
-        canGiveFeedback:
-          message.message.role === MessageRole.Assistant ||
-          message.message.role === MessageRole.Elastic,
-        canRegenerate:
-          (hasConnector && message.message.role === MessageRole.Assistant) ||
-          message.message.role === MessageRole.Elastic,
-        collapsed,
+        role,
+        title,
         content,
-        currentUser,
         element,
-        functionCall: message.message.name
-          ? messages[index - 1].message.function_call
-          : message.message.function_call,
-        hide: message.message.role === MessageRole.System,
+        canCopy,
+        canEdit,
+        canGiveFeedback,
+        canRegenerate,
+        collapsed,
+        currentUser,
+        function_call: functionCall,
+        hide,
         loading: false,
-        role: message.message.role,
-        title,
       };
-
-      return props;
     }),
   ];
 }