From a778f3654df42eb617593d4a3d60db18cc2df871 Mon Sep 17 00:00:00 2001
From: Coen Warmer <coen.warmer@gmail.com>
Date: Tue, 8 Aug 2023 17:27:37 +0200
Subject: [PATCH] Add InlinePrompt Editor, clean up Timeline

---
 .../public/components/chat/chat_item.tsx      | 211 +++++++++++-------
 ...chat_item_content_inline_prompt_editor.tsx |  43 ++++
 .../components/chat/chat_item_title.tsx       |  25 ---
 .../components/chat/chat_prompt_editor.tsx    |  57 +++--
 .../public/components/chat/chat_timeline.tsx  |  17 +-
 .../components/chat/conversation_list.tsx     |   4 +-
 .../components/chat/function_list_popover.tsx |  12 +-
 .../public/hooks/use_json_editor_model.ts     |  21 +-
 .../public/hooks/use_timeline.ts              |  17 +-
 .../conversations/conversation_view.tsx       |   1 +
 .../public/utils/builders.ts                  |   4 +-
 .../get_timeline_items_from_conversation.tsx  |  75 ++++---
 12 files changed, 305 insertions(+), 182 deletions(-)
 create mode 100644 x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
 delete mode 100644 x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_title.tsx

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 e5ff32d9ffcb0..6e1f1ea28ff87 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
@@ -5,6 +5,8 @@
  * 2.0.
  */
 
+import React, { useEffect, useState, useRef } from 'react';
+import { css } from '@emotion/css';
 import {
   EuiAccordion,
   EuiButtonIcon,
@@ -16,17 +18,17 @@ import {
   EuiSpacer,
   EuiText,
   useGeneratedHtmlId,
+  EuiHorizontalRule,
+  EuiPanel,
+  useEuiTheme,
 } from '@elastic/eui';
-import { css } from '@emotion/css';
 import { i18n } from '@kbn/i18n';
-import React, { useEffect, useState } from 'react';
 import { MessageRole } from '../../../common/types';
 import { Feedback, FeedbackButtons } from '../feedback_buttons';
-import { MessageText } from '../message_panel/message_text';
+import { ChatItemContentInlinePromptEditor } from './chat_item_content_inline_prompt_editor';
 import { RegenerateResponseButton } from '../buttons/regenerate_response_button';
 import { StopGeneratingButton } from '../buttons/stop_generating_button';
 import { ChatItemAvatar } from './chat_item_avatar';
-import { ChatItemTitle } from './chat_item_title';
 import { ChatTimelineItem } from './chat_timeline';
 
 export interface ChatItemAction {
@@ -43,22 +45,22 @@ export interface ChatItemProps extends ChatTimelineItem {
   onStopGeneratingClick: () => void;
 }
 
-const euiCommentClassName = css`
-  .euiCommentEvent__headerEvent {
-    flex-grow: 1;
-  }
-
-  > div:last-child {
-    overflow: hidden;
+const normalMessageClassName = css`
+  .euiCommentEvent__body {
+    padding: 0;
   }
 `;
 
-const expandableMessageClassName = css`
+const noPanelMessageClassName = css`
   .euiCommentEvent__header {
     background: transparent;
     border-block-end: none;
   }
 
+  .euiCommentEvent__body {
+    padding: 0;
+  }
+
   .euiCommentEvent {
     border: none;
   }
@@ -71,32 +73,67 @@ const accordionButtonClassName = css`
 `;
 
 export function ChatItem({
-  title,
-  content,
-  element,
   canCopy,
   canEdit,
-  canExpand,
   canGiveFeedback,
   canRegenerate,
-  role,
-  loading,
-  error,
+  collapsed,
+  content,
   currentUser,
+  element,
+  error,
+  functionCall,
+  loading,
   onEditSubmit,
+  onFeedbackClick,
   onRegenerateClick,
   onStopGeneratingClick,
-  onFeedbackClick,
+  role,
+  title,
 }: ChatItemProps) {
+  const { euiTheme } = useEuiTheme();
   const accordionId = useGeneratedHtmlId({ prefix: 'chat' });
 
+  const inlineEditInputRef = useRef<HTMLInputElement | null>(null);
+
+  const [editing, setEditing] = useState<boolean>(false);
   const [isPopoverOpen, setIsPopoverOpen] = useState<string | undefined>();
   const [isAccordionOpen, setIsAccordionOpen] = useState<boolean>(false);
 
+  const noBodyMessageClassName = css`
+    .euiCommentEvent__body {
+      padding: 0;
+      height: ${isAccordionOpen ? 'fit-content' : '0px'};
+      overflow: hidden;
+    }
+  `;
+
   const handleAccordionToggle = () => {
     setIsAccordionOpen(!isAccordionOpen);
+
+    if (editing) {
+      setEditing(false);
+    }
+  };
+
+  const handleToggleEdit = () => {
+    if (collapsed) {
+      setIsAccordionOpen(true);
+    }
+    setEditing(!editing);
+  };
+
+  const handleInlineEditSubmit = (newPrompt: string) => {
+    handleToggleEdit();
+    onEditSubmit(newPrompt);
   };
 
+  useEffect(() => {
+    if (editing) {
+      inlineEditInputRef.current?.focus();
+    }
+  }, [editing]);
+
   useEffect(() => {
     const timeout = setTimeout(() => {
       if (isPopoverOpen) {
@@ -115,26 +152,19 @@ export function ChatItem({
           {
             id: 'edit',
             icon: 'documentEdit',
-            label: i18n.translate(
-              'xpack.observabilityAiAssistant.chatTimeline.actions.editMessage',
-              {
-                defaultMessage: 'Edit message',
-              }
-            ),
+            label: '',
             handler: () => {
-              onEditSubmit(content || '');
+              handleToggleEdit();
             },
           },
         ]
       : []),
-    ...(canExpand
+    ...(collapsed
       ? [
           {
             id: 'expand',
             icon: isAccordionOpen ? 'eyeClosed' : 'eye',
-            label: i18n.translate('xpack.observabilityAiAssistant.chatTimeline.actions.inspect', {
-              defaultMessage: 'Inspect message',
-            }),
+            label: '',
             handler: () => {
               handleAccordionToggle();
             },
@@ -169,29 +199,39 @@ export function ChatItem({
     controls = <StopGeneratingButton onClick={onStopGeneratingClick} />;
   } else if (displayFeedback || displayRegenerate) {
     controls = (
-      <EuiFlexGroup justifyContent="flexEnd">
-        {displayFeedback ? (
-          <EuiFlexItem grow={true}>
-            <FeedbackButtons onClickFeedback={onFeedbackClick} />
-          </EuiFlexItem>
-        ) : null}
-        {displayRegenerate ? (
-          <EuiFlexItem grow={false} style={{ alignSelf: 'flex-end' }}>
-            <RegenerateResponseButton onClick={onRegenerateClick} />
-          </EuiFlexItem>
-        ) : null}
-      </EuiFlexGroup>
+      <>
+        <EuiSpacer size="s" />
+        <EuiHorizontalRule margin="none" color={euiTheme.colors.lightestShade} />
+        <EuiPanel hasShadow={false} paddingSize="s">
+          <EuiFlexGroup justifyContent="flexEnd">
+            {displayFeedback ? (
+              <EuiFlexItem grow={true}>
+                <FeedbackButtons onClickFeedback={onFeedbackClick} />
+              </EuiFlexItem>
+            ) : null}
+            {displayRegenerate ? (
+              <EuiFlexItem grow={false} style={{ alignSelf: 'flex-end' }}>
+                <RegenerateResponseButton onClick={onRegenerateClick} />
+              </EuiFlexItem>
+            ) : null}
+          </EuiFlexGroup>
+        </EuiPanel>
+      </>
     );
   }
 
   let contentElement: React.ReactNode =
-    content || error || controls ? <MessageText content={content || ''} loading={loading} /> : null;
-
-  if (element) {
-    element = <EuiErrorBoundary>{element}</EuiErrorBoundary>;
-  }
+    content || error || controls ? (
+      <ChatItemContentInlinePromptEditor
+        content={content}
+        functionCall={functionCall}
+        loading={loading}
+        editing={editing}
+        onSubmit={handleInlineEditSubmit}
+      />
+    ) : null;
 
-  if (canExpand) {
+  if (collapsed) {
     contentElement = (
       <EuiAccordion
         id={accordionId}
@@ -207,54 +247,57 @@ export function ChatItem({
 
   return (
     <EuiComment
-      event={
-        <ChatItemTitle
-          actions={actions.map(({ id, icon, label, handler }) =>
-            label ? (
-              <EuiPopover
-                button={
-                  <EuiButtonIcon
-                    aria-label={label}
-                    key={id}
-                    iconType={icon}
-                    onClick={() => {
-                      setIsPopoverOpen(id);
-                      handler();
-                    }}
-                    color="text"
-                  />
-                }
-                isOpen={isPopoverOpen === id}
-                closePopover={() => setIsPopoverOpen(undefined)}
-                panelPaddingSize="s"
-              >
-                <EuiText size="s">
-                  <p>{label}</p>
-                </EuiText>
-              </EuiPopover>
-            ) : (
+      actions={actions.map(({ id, icon, label, handler }) =>
+        label ? (
+          <EuiPopover
+            button={
               <EuiButtonIcon
                 aria-label={label}
                 key={id}
                 iconType={icon}
-                onClick={handler}
+                onClick={() => {
+                  setIsPopoverOpen(id);
+                  handler();
+                }}
                 color="text"
               />
-            )
-          )}
-          title={title}
-        />
-      }
+            }
+            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"
+          />
+        )
+      )}
+      event={title}
       className={
-        role === MessageRole.User && canExpand ? expandableMessageClassName : euiCommentClassName
+        actions.length === 0 && !content
+          ? noPanelMessageClassName
+          : collapsed
+          ? noBodyMessageClassName
+          : normalMessageClassName
       }
       timelineAvatar={
         <ChatItemAvatar loading={loading && !content} currentUser={currentUser} role={role} />
       }
       username={getRoleTranslation(role)}
     >
-      {element}
-      {contentElement}
+      <EuiPanel hasShadow={false} paddingSize="s">
+        {element ? <EuiErrorBoundary>{element}</EuiErrorBoundary> : null}
+        {contentElement}
+      </EuiPanel>
+      {controls}
     </EuiComment>
   );
 }
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
new file mode 100644
index 0000000000000..e7bb708705001
--- /dev/null
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_content_inline_prompt_editor.tsx
@@ -0,0 +1,43 @@
+/*
+ * 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 from 'react';
+import { MessageText } from '../message_panel/message_text';
+import { ChatPromptEditor } from './chat_prompt_editor';
+import { MessageRole } from '../../../common';
+
+interface Props {
+  content: string | undefined;
+  functionCall:
+    | {
+        name: string;
+        arguments?: string | undefined;
+        trigger: MessageRole;
+      }
+    | undefined;
+  loading: boolean;
+  editing: boolean;
+  onSubmit: (newPrompt: string) => void;
+}
+export function ChatItemContentInlinePromptEditor({
+  content,
+  functionCall,
+  editing,
+  loading,
+  onSubmit,
+}: Props) {
+  return !editing ? (
+    <MessageText content={content || ''} loading={loading} />
+  ) : (
+    <ChatPromptEditor
+      initialPrompt={content}
+      initialFunctionPayload={functionCall?.arguments}
+      initialSelectedFunctionName={functionCall?.name}
+      onSubmit={onSubmit}
+    />
+  );
+}
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_title.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_title.tsx
deleted file mode 100644
index 0fea4ccde1f68..0000000000000
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/chat_item_title.tsx
+++ /dev/null
@@ -1,25 +0,0 @@
-/*
- * 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 { euiThemeVars } from '@kbn/ui-theme';
-import React, { ReactNode } from 'react';
-
-interface ChatItemTitleProps {
-  actions?: ReactNode;
-  title: string;
-}
-
-export function ChatItemTitle({ actions, title }: ChatItemTitleProps) {
-  return (
-    <>
-      {title}
-      {actions ? (
-        <div css={{ position: 'absolute', top: 4, right: euiThemeVars.euiSizeS }}>{actions}</div>
-      ) : 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 8937afabf663e..cfa10a1ef4435 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
@@ -21,31 +21,47 @@ import { i18n } from '@kbn/i18n';
 import { useObservabilityAIAssistant } from '../../hooks/use_observability_ai_assistant';
 import { useJsonEditorModel } from '../../hooks/use_json_editor_model';
 import { type Message, MessageRole } from '../../../common';
-import type { FunctionDefinition } from '../../../common/types';
 import { FunctionListPopover } from './function_list_popover';
 
 export interface ChatPromptEditorProps {
   disabled: boolean;
   loading: boolean;
+  initialPrompt?: string;
+  initialSelectedFunctionName?: string;
+  initialFunctionPayload?: string;
   onSubmit: (message: Message) => Promise<void>;
 }
 
-export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEditorProps) {
+export function ChatPromptEditor({
+  disabled,
+  loading,
+  initialPrompt,
+  initialSelectedFunctionName,
+  initialFunctionPayload,
+  onSubmit,
+}: ChatPromptEditorProps) {
   const { getFunctions } = useObservabilityAIAssistant();
   const functions = getFunctions();
 
-  const [prompt, setPrompt] = useState('');
+  const [prompt, setPrompt] = useState(initialPrompt);
 
-  const [functionPayload, setFunctionPayload] = useState<string | undefined>('');
-  const [selectedFunction, setSelectedFunction] = useState<FunctionDefinition | undefined>();
+  const [selectedFunctionName, setSelectedFunctionName] = useState<string | undefined>(
+    initialSelectedFunctionName
+  );
+  const [functionPayload, setFunctionPayload] = useState<string | undefined>(
+    initialFunctionPayload
+  );
 
-  const { model, initialJsonString } = useJsonEditorModel(selectedFunction);
+  const { model, initialJsonString } = useJsonEditorModel({
+    functionName: selectedFunctionName,
+    initialJson: initialFunctionPayload,
+  });
 
   const textAreaRef = useRef<HTMLTextAreaElement>(null);
 
   useEffect(() => {
     setFunctionPayload(initialJsonString);
-  }, [initialJsonString, selectedFunction]);
+  }, [initialJsonString, selectedFunctionName]);
 
   const handleChange = (event: React.ChangeEvent<HTMLTextAreaElement>) => {
     setPrompt(event.currentTarget.value);
@@ -56,8 +72,15 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
   };
 
   const handleClearSelection = () => {
-    setSelectedFunction(undefined);
+    setSelectedFunctionName(undefined);
+    setFunctionPayload('');
+    setPrompt('');
+  };
+
+  const handleSelectFunction = (functionName: string) => {
+    setPrompt('');
     setFunctionPayload('');
+    setSelectedFunctionName(functionName);
   };
 
   const handleResizeTextArea = () => {
@@ -76,13 +99,13 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
     handleResizeTextArea();
 
     try {
-      if (selectedFunction) {
+      if (selectedFunctionName) {
         await onSubmit({
           '@timestamp': new Date().toISOString(),
           message: {
             role: MessageRole.Function,
             function_call: {
-              name: selectedFunction.options.name,
+              name: selectedFunctionName,
               trigger: MessageRole.User,
               arguments: currentPayload,
             },
@@ -98,7 +121,7 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
     } catch (_) {
       setPrompt(currentPrompt);
     }
-  }, [functionPayload, onSubmit, prompt, selectedFunction]);
+  }, [functionPayload, onSubmit, prompt, selectedFunctionName]);
 
   useEffect(() => {
     const keyboardListener = (event: KeyboardEvent) => {
@@ -136,12 +159,12 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
               <EuiFlexItem grow>
                 <FunctionListPopover
                   functions={functions}
-                  selectedFunction={selectedFunction}
-                  onSelectFunction={setSelectedFunction}
+                  selectedFunctionName={selectedFunctionName}
+                  onSelectFunction={handleSelectFunction}
                 />
               </EuiFlexItem>
               <EuiFlexItem grow={false}>
-                {selectedFunction ? (
+                {selectedFunctionName ? (
                   <EuiButtonEmpty
                     iconType="cross"
                     iconSide="right"
@@ -157,7 +180,7 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
             </EuiFlexGroup>
           </EuiFlexItem>
           <EuiFlexItem>
-            {selectedFunction ? (
+            {selectedFunctionName ? (
               <EuiPanel borderRadius="none" color="subdued" hasShadow={false} paddingSize="xs">
                 <CodeEditor
                   aria-label="payloadEditor"
@@ -222,9 +245,9 @@ export function ChatPromptEditor({ disabled, loading, onSubmit }: ChatPromptEdit
         <EuiButtonIcon
           aria-label="Submit"
           isLoading={loading}
-          disabled={selectedFunction ? false : !prompt || loading || disabled}
+          disabled={selectedFunctionName ? false : !prompt || loading || disabled}
           display={
-            selectedFunction ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
+            selectedFunctionName ? (functionPayload ? 'fill' : 'base') : prompt ? 'fill' : 'base'
           }
           iconType="kqlFunction"
           size="m"
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 664972a59515b..84e43dca9f251 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
@@ -5,24 +5,29 @@
  * 2.0.
  */
 
+import React from 'react';
 import { EuiCommentList } from '@elastic/eui';
 import type { AuthenticatedUser } from '@kbn/security-plugin/common';
-import React from 'react';
-import type { Message } from '../../../common';
+import { type Message, MessageRole } from '../../../common';
 import type { Feedback } from '../feedback_buttons';
 import { ChatItem } from './chat_item';
 
-export interface ChatTimelineItem
-  extends Pick<Message['message'], 'role' | 'content' | 'function_call'> {
+export interface ChatTimelineItem extends Pick<Message['message'], 'role' | 'content'> {
   id: string;
   title: string;
   loading: boolean;
   canCopy: boolean;
-  element?: React.ReactNode;
   canEdit: boolean;
-  canExpand: boolean;
   canGiveFeedback: boolean;
   canRegenerate: boolean;
+
+  collapsed: boolean;
+  functionCall?: {
+    name: string;
+    arguments?: string;
+    trigger: MessageRole.Assistant | MessageRole.User | MessageRole.Elastic;
+  };
+  element?: React.ReactNode;
   hide: boolean;
   currentUser?: Pick<AuthenticatedUser, 'username' | 'full_name'>;
   error?: any;
diff --git a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
index 27b1ff23d5c06..05408f62ea6ba 100644
--- a/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/components/chat/conversation_list.tsx
@@ -37,11 +37,11 @@ export function ConversationList({
   onClickDeleteConversation,
 }: {
   selected: string;
-  onClickConversation: (conversationId: string) => void;
-  onClickNewChat: () => void;
   loading: boolean;
   error?: any;
   conversations?: Array<{ id: string; label: string; href?: string }>;
+  onClickConversation: (conversationId: string) => void;
+  onClickNewChat: () => void;
   onClickDeleteConversation: (id: string) => void;
 }) {
   return (
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 c85a15fe71bd5..ae5f42a8557e8 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
@@ -19,12 +19,12 @@ import { FunctionDefinition } from '../../../common/types';
 
 export function FunctionListPopover({
   functions,
-  selectedFunction,
+  selectedFunctionName,
   onSelectFunction,
 }: {
   functions: FunctionDefinition[];
-  selectedFunction?: FunctionDefinition;
-  onSelectFunction: (func: FunctionDefinition) => void;
+  selectedFunctionName?: string;
+  onSelectFunction: (func: string) => void;
 }) {
   const [isFunctionListOpen, setIsFunctionListOpen] = useState(false);
 
@@ -34,7 +34,7 @@ export function FunctionListPopover({
 
   const handleSelectFunction = (func: FunctionDefinition) => {
     setIsFunctionListOpen(false);
-    onSelectFunction(func);
+    onSelectFunction(func.options.name);
   };
 
   useEffect(() => {
@@ -61,8 +61,8 @@ export function FunctionListPopover({
           size="xs"
           onClick={handleClickFunctionList}
         >
-          {selectedFunction
-            ? selectedFunction.options.name
+          {selectedFunctionName
+            ? selectedFunctionName
             : i18n.translate('xpack.observabilityAiAssistant.prompt.callFunction', {
                 defaultMessage: 'Call function',
               })}
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 f04cdf68e9c3e..fba0795bda391 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
@@ -6,14 +6,25 @@
  */
 import { useMemo } from 'react';
 import { monaco } from '@kbn/monaco';
-import { FunctionDefinition } from '../../common/types';
+import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
 
 const { editor, languages, Uri } = monaco;
 
 const SCHEMA_URI = 'http://elastic.co/foo.json';
 const modelUri = Uri.parse(SCHEMA_URI);
 
-export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
+export const useJsonEditorModel = ({
+  functionName,
+  initialJson,
+}: {
+  functionName: string | undefined;
+  initialJson?: string | undefined;
+}) => {
+  const { getFunctions } = useObservabilityAIAssistant();
+  const functions = getFunctions();
+
+  const functionDefinition = functions.find((func) => func.options.name === functionName);
+
   return useMemo(() => {
     if (!functionDefinition) {
       return {};
@@ -21,7 +32,9 @@ export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
 
     const schema = { ...functionDefinition.options.parameters };
 
-    const initialJsonString = functionDefinition.options.parameters.properties
+    const initialJsonString = initialJson
+      ? initialJson
+      : functionDefinition.options.parameters.properties
       ? Object.keys(functionDefinition.options.parameters.properties).reduce(
           (acc, curr, index, arr) => {
             const val = `${acc}  "${curr}": "",\n`;
@@ -49,5 +62,5 @@ export const useJsonEditorModel = (functionDefinition?: FunctionDefinition) => {
     }
 
     return { model, initialJsonString };
-  }, [functionDefinition]);
+  }, [functionDefinition, initialJson]);
 };
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 5091c47ffc03e..5ac0ce76ec34c 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
@@ -135,8 +135,8 @@ export function useTimeline({
             messagesAfterChat.concat({
               '@timestamp': new Date().toISOString(),
               message: {
-                role: MessageRole.User,
                 name,
+                role: MessageRole.User,
                 content: JSON.stringify(message.content),
                 data: JSON.stringify(message.data),
               },
@@ -169,18 +169,17 @@ export function useTimeline({
         id: '',
         canCopy: true,
         canEdit: false,
-        canExpand:
-          pendingMessage.message.role === MessageRole.User && Boolean(pendingMessage.message.name),
-        canRegenerate: pendingMessage.aborted || !!pendingMessage.error,
         canGiveFeedback: false,
-        title: '',
-        role: pendingMessage.message.role,
+        canRegenerate: pendingMessage.aborted || !!pendingMessage.error,
+        collapsed: false,
         content: pendingMessage.message.content,
-        loading: !pendingMessage.aborted && !pendingMessage.error,
-        function_call: pendingMessage.message.function_call,
-        hide: Boolean(pendingMessage.message.isAssistantSetupMessage),
         currentUser,
         error: pendingMessage.error,
+        functionCall: pendingMessage.message.function_call,
+        hide: Boolean(pendingMessage.message.isAssistantSetupMessage),
+        loading: !pendingMessage.aborted && !pendingMessage.error,
+        role: pendingMessage.message.role,
+        title: '',
       });
     }
 
diff --git a/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx b/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
index f46bf32246451..e278bdd6abec1 100644
--- a/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
+++ b/x-pack/plugins/observability_ai_assistant/public/routes/conversations/conversation_view.tsx
@@ -173,6 +173,7 @@ export function ConversationView() {
                 });
             }}
           />
+          <EuiSpacer size="m" />
         </EuiFlexItem>
         <EuiFlexItem
           grow
diff --git a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
index cfed691438a77..07f4b70d49c1f 100644
--- a/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
+++ b/x-pack/plugins/observability_ai_assistant/public/utils/builders.ts
@@ -19,7 +19,7 @@ export function buildChatItem(params: ChatItemBuildProps): ChatTimelineItem {
     canEdit: false,
     canGiveFeedback: false,
     canCopy: true,
-    canExpand: false,
+    collapsed: false,
     canRegenerate: params.role === MessageRole.User,
     currentUser: {
       username: 'elastic',
@@ -69,7 +69,7 @@ export function buildFunctionChatItem(params: Omit<ChatItemBuildProps, 'role'>)
   return buildChatItem({
     role: MessageRole.User,
     title: 'executed a function',
-    function_call: {
+    functionCall: {
       name: 'leftpad',
       arguments: '{ foo: "bar" }',
       trigger: MessageRole.Assistant,
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 f1a1bf29f86fe..bd5c12d59c339 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
@@ -12,6 +12,12 @@ import { type Message, MessageRole } from '../../common';
 import type { ChatTimelineItem } from '../components/chat/chat_timeline';
 import { RenderFunction } from '../components/render_function';
 
+function convertFunctionParamsToMarkdownCodeBlock(object: Record<string, string | number>) {
+  return `\`\`\`
+${JSON.stringify(object, null, 4)}
+\`\`\``;
+}
+
 export function getTimelineItemsfromConversation({
   currentUser,
   messages,
@@ -24,18 +30,18 @@ export function getTimelineItemsfromConversation({
   return [
     {
       id: v4(),
-      role: MessageRole.User,
-      title: i18n.translate('xpack.observabilityAiAssistant.conversationStartTitle', {
-        defaultMessage: 'started a conversation',
-      }),
       canCopy: false,
       canEdit: false,
-      canExpand: false,
       canGiveFeedback: false,
       canRegenerate: false,
+      collapsed: false,
+      currentUser,
       hide: false,
       loading: false,
-      currentUser,
+      role: MessageRole.User,
+      title: i18n.translate('xpack.observabilityAiAssistant.conversationStartTitle', {
+        defaultMessage: 'started a conversation',
+      }),
     },
     ...messages.map((message, index) => {
       const hasFunction = !!message.message.function_call?.name;
@@ -44,39 +50,45 @@ export function getTimelineItemsfromConversation({
       let title: string;
       let content: string | undefined;
       let element: React.ReactNode | undefined;
+      let collapsed: boolean = false;
 
       if (hasFunction) {
         title = i18n.translate('xpack.observabilityAiAssistant.suggestedFunctionEvent', {
-          defaultMessage: 'suggested to call {functionName}',
+          defaultMessage: 'suggested to use function {functionName}',
           values: {
             functionName: message.message.function_call!.name,
           },
         });
-        content =
-          '```\n' +
-          JSON.stringify(
-            {
-              name: message.message.function_call!.name,
-              arguments: JSON.parse(message.message.function_call?.arguments || ''),
-            },
-            null,
-            4
-          ) +
-          '\n```';
+
+        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 a function',
+          defaultMessage: 'executed the function {functionName}',
+          values: {
+            functionName: prevMessage.message.function_call!.name,
+          },
         });
-        content = message.message.content;
+
+        content = convertFunctionParamsToMarkdownCodeBlock(
+          JSON.parse(message.message.content || '{}')
+        );
+
         element = (
           <RenderFunction
             name={message.message.name}
@@ -84,25 +96,34 @@ export function getTimelineItemsfromConversation({
             response={message.message}
           />
         );
+        collapsed = true;
       } else {
         title = '';
         content = message.message.content;
+        collapsed = false;
       }
 
       const props = {
         id: v4(),
-        role: message.message.role,
         canCopy: true,
         canEdit: hasConnector && (message.message.role === MessageRole.User || hasFunction),
-        canExpand: Boolean(message.message.name) || Boolean(message.message.function_call),
-        canRegenerate: hasConnector && message.message.role === MessageRole.Assistant,
-        canGiveFeedback: message.message.role === MessageRole.Assistant,
-        loading: false,
-        hide: Boolean(message.message.isAssistantSetupMessage),
-        title,
+        canGiveFeedback:
+          message.message.role === MessageRole.Assistant ||
+          message.message.role === MessageRole.Elastic,
+        canRegenerate:
+          (hasConnector && message.message.role === MessageRole.Assistant) ||
+          message.message.role === MessageRole.Elastic,
+        collapsed,
         content,
         currentUser,
         element,
+        functionCall: message.message.name
+          ? messages[index - 1].message.function_call
+          : message.message.function_call,
+        hide: Boolean(message.message.isAssistantSetupMessage),
+        loading: false,
+        role: message.message.role,
+        title,
       };
 
       return props;