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 a3c8c8f42437a5..e66a32035a0d5f 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'};
@@ -106,7 +114,7 @@ export function ChatItem({
 
   const handleToggleEdit = () => {
     if (collapsed) {
-      setExpanded(false);
+      setExpanded(!expanded);
     }
     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 7190ccca73dccf..79240a03bb3143 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 d1e08e5c4a6988..8d1559c34c1983 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
@@ -103,7 +103,7 @@ export function ChatPromptEditor({
         await onSubmit({
           '@timestamp': new Date().toISOString(),
           message: {
-            role: MessageRole.Function,
+            role: MessageRole.User,
             function_call: {
               name: selectedFunctionName,
               trigger: MessageRole.User,
@@ -111,6 +111,9 @@ export function ChatPromptEditor({
             },
           },
         });
+
+        setFunctionPayload(undefined);
+        setSelectedFunctionName(undefined);
       } else {
         await onSubmit({
           '@timestamp': new Date().toISOString(),
@@ -126,8 +129,8 @@ export function ChatPromptEditor({
   useEffect(() => {
     const keyboardListener = (event: KeyboardEvent) => {
       if (!event.shiftKey && event.key === keys.ENTER) {
-        handleSubmit();
         event.preventDefault();
+        handleSubmit();
       }
     };
 
@@ -188,8 +191,6 @@ export function ChatPromptEditor({
                   fullWidth
                   height="120px"
                   languageId="json"
-                  value={functionPayload || ''}
-                  onChange={handleChangeFunctionPayload}
                   isCopyable
                   languageConfiguration={{
                     autoClosingPairs: [
@@ -199,6 +200,9 @@ export function ChatPromptEditor({
                       },
                     ],
                   }}
+                  editorDidMount={(editor) => {
+                    editor.focus();
+                  }}
                   options={{
                     accessibilitySupport: 'off',
                     acceptSuggestionOnEnter: 'on',
@@ -223,18 +227,20 @@ export function ChatPromptEditor({
                     wrappingIndent: 'indent',
                   }}
                   transparentBackground
+                  value={functionPayload || ''}
+                  onChange={handleChangeFunctionPayload}
                 />
               </EuiPanel>
             ) : (
               <EuiTextArea
-                resize="vertical"
-                rows={1}
                 fullWidth
+                inputRef={textAreaRef}
                 placeholder={i18n.translate('xpack.observabilityAiAssistant.prompt.placeholder', {
                   defaultMessage: 'Press ‘$’ for function recommendations',
                 })}
+                resize="vertical"
+                rows={1}
                 value={prompt}
-                inputRef={textAreaRef}
                 onChange={handleChange}
               />
             )}
@@ -245,12 +251,12 @@ export function ChatPromptEditor({
         <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"
+          isLoading={loading}
           size="m"
           onClick={handleSubmit}
         />
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 2f550fe1f965c5..56d1d35c9178f2 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 81159709431357..a030c3d74d96db 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,9 +8,10 @@ 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 `\`\`\`
@@ -22,10 +23,12 @@ export function getTimelineItemsfromConversation({
   currentUser,
   messages,
   hasConnector,
+  functions,
 }: {
   currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
   messages: Message[];
   hasConnector: boolean;
+  functions: FunctionDefinition[];
 }): ChatTimelineItem[] {
   return [
     {
@@ -45,90 +48,130 @@ export function getTimelineItemsfromConversation({
       }),
     },
     ...messages.map((message, index) => {
-      const hasFunction = !!message.message.function_call?.name;
-      const isSystemPrompt = message.message.role === MessageRole.System;
+      const id = v4();
+      const role = message.message.role;
 
-      let title: string;
+      let title: string = '';
       let content: string | undefined;
       let element: React.ReactNode | undefined;
+
+      const functionCall = message.message.name
+        ? 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:
+          // is a prompt by the user
+          if (!message.message.name) {
+            title = '';
+            content = message.message.content;
+
+            canCopy = true;
+            canEdit = hasConnector;
+            canGiveFeedback = false;
+            canRegenerate = false;
+            collapsed = false;
+            hide = false;
+          } else {
+            // user has executed a function
+            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 || '{}')
+            );
+
+            const fn = functions.find((func) => func.options.name === message.message.name);
 
-      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;
+            element = fn?.render ? (
+              <RenderFunction
+                name={message.message.name}
+                arguments={prevMessage.message.function_call.arguments}
+                response={message.message}
+              />
+            ) : null;
+
+            canCopy = true;
+            canEdit = hasConnector;
+            canGiveFeedback = true;
+            canRegenerate = hasConnector;
+            collapsed = !Boolean(element);
+            hide = false;
+          }
+          break;
+
+        case MessageRole.Assistant:
+          // is a function suggestion by the assistant
+          if (!!message.message.function_call?.name) {
+            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 || '{}'),
+            });
+
+            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;
     }),
   ];
 }