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 fc5552d008fa7d..9f581b31795f14 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
@@ -12,7 +12,6 @@ import {
   EuiComment,
   EuiErrorBoundary,
   EuiPanel,
-  EuiSpacer,
   useGeneratedHtmlId,
 } from '@elastic/eui';
 import { ChatItemActions } from './chat_item_actions';
@@ -27,7 +26,7 @@ import type { Feedback } from '../feedback_buttons';
 import type { ChatActionClickHandler } from './types';
 import type { TelemetryEventTypeWithPayload } from '../../analytics';
 
-export interface ChatItemProps extends ChatTimelineItem {
+export interface ChatItemProps extends Omit<ChatTimelineItem, 'message'> {
   onActionClick: ChatActionClickHandler;
   onEditSubmit: (message: Message) => void;
   onFeedbackClick: (feedback: Feedback) => void;
@@ -36,15 +35,19 @@ export interface ChatItemProps extends ChatTimelineItem {
   onStopGeneratingClick: () => void;
 }
 
-const normalMessageClassName = css`
-  .euiCommentEvent__body {
-    padding: 0;
-  }
-
+const moreCompactHeaderClassName = css`
   .euiCommentEvent__header > .euiPanel {
     padding-top: 4px;
     padding-bottom: 4px;
   }
+`;
+
+const normalMessageClassName = css`
+  ${moreCompactHeaderClassName}
+
+  .euiCommentEvent__body {
+    padding: 0;
+  }
 
   /* targets .*euiTimelineItemEvent-top, makes sure text properly wraps and doesn't overflow */
   > :last-child {
@@ -74,12 +77,13 @@ const noPanelMessageClassName = css`
 export function ChatItem({
   actions: { canCopy, canEdit, canGiveFeedback, canRegenerate },
   content,
+  function_call: functionCall,
+  role,
   currentUser,
   display: { collapsed },
   element,
   error,
   loading,
-  message,
   title,
   onActionClick,
   onEditSubmit,
@@ -96,10 +100,13 @@ export function ChatItem({
   const actions = [canCopy, collapsed, canCopy].filter(Boolean);
 
   const noBodyMessageClassName = css`
+    ${moreCompactHeaderClassName}
+
     .euiCommentEvent__body {
       padding: 0;
       height: ${expanded ? 'fit-content' : '0px'};
       overflow: hidden;
+      border: none;
     }
   `;
 
@@ -132,7 +139,9 @@ export function ChatItem({
       <ChatItemContentInlinePromptEditor
         editing={editing}
         loading={loading}
-        message={message}
+        functionCall={functionCall}
+        content={content}
+        role={role}
         onSubmit={handleInlineEditSubmit}
         onActionClick={onActionClick}
         onSendTelemetry={onSendTelemetry}
@@ -147,7 +156,6 @@ export function ChatItem({
         forceState={expanded ? 'open' : 'closed'}
         onToggle={handleToggleExpand}
       >
-        <EuiSpacer size="s" />
         {contentElement}
       </EuiAccordion>
     );
@@ -155,10 +163,8 @@ export function ChatItem({
 
   return (
     <EuiComment
-      timelineAvatar={
-        <ChatItemAvatar loading={loading} currentUser={currentUser} role={message.message.role} />
-      }
-      username={getRoleTranslation(message.message.role)}
+      timelineAvatar={<ChatItemAvatar loading={loading} currentUser={currentUser} role={role} />}
+      username={getRoleTranslation(role)}
       event={title}
       actions={
         <ChatItemActions
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
index 05f21ba92bb629..a01032e8711c51 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
@@ -6,6 +6,8 @@
  */
 
 import React from 'react';
+import { EuiPanel } from '@elastic/eui';
+import { css } from '@emotion/css';
 import { MessageText } from '../message_panel/message_text';
 import { ChatPromptEditor } from './chat_prompt_editor';
 import type { Message } from '../../../common';
@@ -15,34 +17,59 @@ import type { TelemetryEventTypeWithPayload } from '../../analytics';
 interface Props {
   editing: boolean;
   loading: boolean;
-  message: Message;
+  role: Message['message']['role'];
+  content: Message['message']['content'];
+  functionCall: Message['message']['function_call'];
   onActionClick: ChatActionClickHandler;
   onSendTelemetry: (eventWithPayload: TelemetryEventTypeWithPayload) => void;
   onSubmit: (message: Message) => void;
 }
+
+const textContainerClassName = css`
+  padding: 4px 0;
+`;
+
+const editorContainerClassName = css`
+  padding: 12px 0;
+`;
+
 export function ChatItemContentInlinePromptEditor({
   editing,
   loading,
-  message,
+  functionCall,
+  content,
+  role,
   onActionClick,
   onSendTelemetry,
   onSubmit,
 }: Props) {
   return !editing ? (
-    <MessageText
-      content={message.message.content || ''}
-      loading={loading}
-      onActionClick={onActionClick}
-    />
+    <EuiPanel
+      paddingSize="none"
+      hasBorder={false}
+      hasShadow={false}
+      className={textContainerClassName}
+    >
+      <MessageText content={content || ''} loading={loading} onActionClick={onActionClick} />
+    </EuiPanel>
   ) : (
-    <ChatPromptEditor
-      disabled={false}
-      hidden={false}
-      loading={false}
-      initialMessage={message}
-      onChangeHeight={() => {}}
-      onSubmit={onSubmit}
-      onSendTelemetry={onSendTelemetry}
-    />
+    <EuiPanel
+      paddingSize="none"
+      hasBorder={false}
+      hasShadow={false}
+      className={editorContainerClassName}
+    >
+      <ChatPromptEditor
+        disabled={false}
+        hidden={false}
+        loading={false}
+        initialFunctionCall={functionCall}
+        initialContent={content}
+        initialRole={role}
+        onChangeHeight={() => {}}
+        onSubmit={onSubmit}
+        onSendTelemetry={onSendTelemetry}
+      />
+    </EuiPanel>
   );
 }
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_controls.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_controls.tsx
index d828b0330d69c6..2d2743488f766d 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_controls.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_controls.tsx
@@ -6,18 +6,17 @@
  */
 
 import React from 'react';
-import {
-  EuiFlexGroup,
-  EuiFlexItem,
-  EuiHorizontalRule,
-  EuiPanel,
-  EuiSpacer,
-  useEuiTheme,
-} from '@elastic/eui';
+import { EuiFlexGroup, EuiFlexItem, EuiHorizontalRule, EuiPanel, useEuiTheme } from '@elastic/eui';
+import { css } from '@emotion/css';
 import { Feedback, FeedbackButtons } from '../feedback_buttons';
 import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
 import { StopGeneratingButton } from '../buttons/stop_generating_button';
 
+const containerClassName = css`
+  padding-top: 4px;
+  padding-bottom: 4px;
+`;
+
 export function ChatItemControls({
   error,
   loading,
@@ -65,9 +64,8 @@ export function ChatItemControls({
 
   return controls ? (
     <>
-      <EuiSpacer size="s" />
       <EuiHorizontalRule margin="none" color={euiTheme.colors.lightestShade} />
-      <EuiPanel hasShadow={false} paddingSize="s">
+      <EuiPanel hasShadow={false} paddingSize="s" className={containerClassName}>
         {controls}
       </EuiPanel>
     </>
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 4d0c2e86fc0b36..4b35d1ee11432e 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
@@ -20,7 +20,9 @@ export interface ChatPromptEditorProps {
   disabled: boolean;
   hidden: boolean;
   loading: boolean;
-  initialMessage?: Message;
+  initialRole?: Message['message']['role'];
+  initialFunctionCall?: Message['message']['function_call'];
+  initialContent?: Message['message']['content'];
   onChangeHeight: (height: number) => void;
   onSendTelemetry: (eventWithPayload: TelemetryEventTypeWithPayload) => void;
   onSubmit: (message: Message) => void;
@@ -30,21 +32,31 @@ export function ChatPromptEditor({
   disabled,
   hidden,
   loading,
-  initialMessage,
+  initialRole,
+  initialFunctionCall,
+  initialContent,
   onChangeHeight,
   onSendTelemetry,
   onSubmit,
 }: ChatPromptEditorProps) {
   const containerRef = useRef<HTMLDivElement>(null);
 
-  const isFocusTrapEnabled = Boolean(initialMessage?.message);
+  const isFocusTrapEnabled = Boolean(initialContent || initialFunctionCall);
 
-  const [innerMessage, setInnerMessage] = useState<Message['message'] | undefined>(
-    initialMessage?.message
+  const [mode, setMode] = useState<'prompt' | 'function'>(
+    initialFunctionCall?.name ? 'function' : 'prompt'
   );
 
-  const [mode, setMode] = useState<'prompt' | 'function'>(
-    initialMessage?.message.function_call?.name ? 'function' : 'prompt'
+  const initialInnerMessage = initialRole
+    ? {
+        role: initialRole,
+        content: initialContent ?? '',
+        ...(initialFunctionCall ? { function_call: initialFunctionCall } : {}),
+      }
+    : undefined;
+
+  const [innerMessage, setInnerMessage] = useState<Message['message'] | undefined>(
+    initialInnerMessage
   );
 
   const handleChangeMessageInner = (newInnerMessage: Message['message']) => {
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor_function.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor_function.tsx
index b3ec4609e5d2de..33f467be26cb23 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor_function.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_prompt_editor_function.tsx
@@ -9,6 +9,7 @@ import { CodeEditor } from '@kbn/kibana-react-plugin/public';
 import { i18n } from '@kbn/i18n';
 import usePrevious from 'react-use/lib/usePrevious';
 import { EuiCode, EuiPanel } from '@elastic/eui';
+import { css } from '@emotion/css';
 import { useJsonEditorModel } from '../../hooks/use_json_editor_model';
 import { type Message, MessageRole } from '../../../common';
 
@@ -17,6 +18,11 @@ export interface Props {
   functionPayload?: string;
   onChange: (message: Message['message']) => void;
 }
+
+const functionNameClassName = css`
+  display: inline-block;
+`;
+
 export function ChatPromptEditorFunction({ functionName, functionPayload, onChange }: Props) {
   const [functionEditorLineCount, setFunctionEditorLineCount] = useState<number>(0);
 
@@ -67,8 +73,8 @@ export function ChatPromptEditorFunction({ functionName, functionPayload, onChan
   }, [functionName, functionPayload, initialJsonString, onChange, previousPayload]);
 
   return (
-    <EuiPanel paddingSize="none">
-      <EuiCode>{functionName}</EuiCode>
+    <EuiPanel paddingSize="none" hasShadow={false} hasBorder>
+      <EuiCode className={functionNameClassName}>{functionName}</EuiCode>
       <CodeEditor
         aria-label={i18n.translate(
           'xpack.observabilityAiAssistant.chatPromptEditor.codeEditor.payloadEditorLabel',
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 8b1067064e20c5..afee49619fd248 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
@@ -9,6 +9,7 @@ import React, { ReactNode, useMemo } from 'react';
 import { css } from '@emotion/css';
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
+import { omit } from 'lodash';
 import type { Feedback } from '../feedback_buttons';
 import type { Message } from '../../../common';
 import type { UseKnowledgeBaseResult } from '../../hooks/use_knowledge_base';
@@ -42,6 +43,7 @@ export interface ChatTimelineItem
   currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
   error?: any;
   message: Message;
+  functionCall?: Message['message']['function_call'];
 }
 
 export interface ChatTimelineProps {
@@ -128,7 +130,7 @@ export function ChatTimeline({
           <ChatItem
             // use index, not id to prevent unmounting of component when message is persisted
             key={index}
-            {...item}
+            {...omit(item, 'message')}
             onActionClick={onActionClick}
             onFeedbackClick={(feedback) => {
               onFeedback(item.message, feedback);
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
index fdc3171ae03cfe..e3c5f27e7bfe76 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/function_list_popover.tsx
@@ -180,7 +180,7 @@ export function FunctionListPopover({
         }}
         singleSelection
         onChange={(options) => {
-          const selectedFunction = options.filter((fn) => fn.checked !== 'off');
+          const selectedFunction = options.filter((fn) => !('checked' in fn));
           if (selectedFunction && selectedFunction.length === 1) {
             handleSelectFunction({ ...selectedFunction[0], checked: 'on' });
           }
diff --git a/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts b/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
index 44122614abd5da..230d0065e29c3b 100644
--- a/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/hooks/use_json_editor_model.ts
@@ -4,10 +4,11 @@
  * 2.0; you may not use this file except in compliance with the Elastic License
  * 2.0.
  */
+import { useEffect, useMemo, useState } from 'react';
 import { monaco } from '@kbn/monaco';
-import { useMemo } from 'react';
 import { createInitializedObject } from '../utils/create_initialized_object';
 import { useObservabilityAIAssistantChatService } from './use_observability_ai_assistant_chat_service';
+import { safeJsonParse } from '../utils/safe_json_parse';
 
 const { editor, languages, Uri } = monaco;
 
@@ -25,6 +26,13 @@ export const useJsonEditorModel = ({
 
   const functionDefinition = chatService.getFunctions().find((func) => func.name === functionName);
 
+  const [initialJsonValue, setInitialJsonValue] = useState<string | undefined>(initialJson);
+
+  useEffect(() => {
+    setInitialJsonValue(initialJson);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [functionName]);
+
   return useMemo(() => {
     if (!functionDefinition) {
       return {};
@@ -32,8 +40,8 @@ export const useJsonEditorModel = ({
 
     const schema = { ...functionDefinition.parameters };
 
-    const initialJsonString = initialJson
-      ? initialJson
+    const initialJsonString = initialJsonValue
+      ? JSON.stringify(safeJsonParse(initialJsonValue), null, 4) // prettify the json
       : functionDefinition.parameters.properties
       ? JSON.stringify(createInitializedObject(functionDefinition.parameters), null, 4)
       : '';
@@ -58,5 +66,5 @@ export const useJsonEditorModel = ({
     }
 
     return { model, initialJsonString };
-  }, [functionDefinition, initialJson]);
+  }, [functionDefinition, initialJsonValue]);
 };
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 e8851a7ceddd8c..2b910211ba5fd2 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
@@ -16,14 +16,7 @@ import type { ChatTimelineItem } from '../components/chat/chat_timeline';
 import { RenderFunction } from '../components/render_function';
 import type { ObservabilityAIAssistantChatService } from '../types';
 import { ChatState } from '../hooks/use_chat';
-
-function safeParse(jsonStr: string) {
-  try {
-    return JSON.parse(jsonStr);
-  } catch (err) {
-    return jsonStr;
-  }
-}
+import { safeJsonParse } from './safe_json_parse';
 
 function convertMessageToMarkdownCodeBlock(message: Message['message']) {
   let value: object;
@@ -31,7 +24,7 @@ function convertMessageToMarkdownCodeBlock(message: Message['message']) {
   if (!message.name) {
     const name = message.function_call?.name;
     const args = message.function_call?.arguments
-      ? safeParse(message.function_call.arguments)
+      ? safeJsonParse(message.function_call.arguments)
       : undefined;
 
     value = {
@@ -41,9 +34,9 @@ function convertMessageToMarkdownCodeBlock(message: Message['message']) {
   } else {
     const content =
       message.role !== MessageRole.Assistant && message.content
-        ? safeParse(message.content)
+        ? safeJsonParse(message.content)
         : message.content;
-    const data = message.data ? safeParse(message.data) : undefined;
+    const data = message.data ? safeJsonParse(message.data) : undefined;
     value = omitBy(
       {
         content,
diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/safe_json_parse.ts b/x-pack/plugins/observability_ai_assistant/public/utils/safe_json_parse.ts
new file mode 100644
index 00000000000000..a4f2dfa5c25036
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/utils/safe_json_parse.ts
@@ -0,0 +1,14 @@
+/*
+ * 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.
+ */
+
+export function safeJsonParse(jsonStr: string) {
+  try {
+    return JSON.parse(jsonStr);
+  } catch (err) {
+    return jsonStr;
+  }
+}