[8.12] Fix for code rendering in messages, code prompt function selec…

…ting (elastic#173920) (elastic#173928)

# Backport

This will backport the following commits from `main` to `8.12`:
- [Fix for code rendering in messages, code prompt function selecting
(elastic#173920)](elastic#173920)

<!--- Backport version: 8.9.7 -->

### Questions ?
Please refer to the [Backport tool
documentation](https://github.com/sqren/backport)

<!--BACKPORT [{"author":{"name":"Coen
Warmer","email":"coen.warmer@gmail.com"},"sourceCommit":{"committedDate":"2023-12-22T15:56:59Z","message":"Fix
for code rendering in messages, code prompt function selecting
(elastic#173920)\n\n## Summary\r\n\r\nFixes code rendering in Chat
items.\r\n\r\nAlso fixed bug with selecting
functions.","sha":"e6d1c3d320feaffa062ef911e91222193ac73e42","branchLabelMapping":{"^v8.13.0$":"main","^v(\\d+).(\\d+).\\d+$":"$1.$2"}},"sourcePullRequest":{"labels":["release_note:skip","backport:prev-minor","v8.12.0","v8.13.0"],"number":173920,"url":"https://github.com/elastic/kibana/pull/173920","mergeCommit":{"message":"Fix
for code rendering in messages, code prompt function selecting
(elastic#173920)\n\n## Summary\r\n\r\nFixes code rendering in Chat
items.\r\n\r\nAlso fixed bug with selecting
functions.","sha":"e6d1c3d320feaffa062ef911e91222193ac73e42"}},"sourceBranch":"main","suggestedTargetBranches":["8.12"],"targetPullRequestStates":[{"branch":"8.12","label":"v8.12.0","labelRegex":"^v(\\d+).(\\d+).\\d+$","isSourceBranch":false,"state":"NOT_CREATED"},{"branch":"main","label":"v8.13.0","labelRegex":"^v8.13.0$","isSourceBranch":true,"state":"MERGED","url":"https://github.com/elastic/kibana/pull/173920","number":173920,"mergeCommit":{"message":"Fix
for code rendering in messages, code prompt function selecting
(elastic#173920)\n\n## Summary\r\n\r\nFixes code rendering in Chat
items.\r\n\r\nAlso fixed bug with selecting
functions.","sha":"e6d1c3d320feaffa062ef911e91222193ac73e42"}}]}]
BACKPORT-->

Co-authored-by: Coen Warmer <coen.warmer@gmail.com>

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,18 +156,15 @@ export function ChatItem({
         forceState={expanded ? 'open' : 'closed'}
         onToggle={handleToggleExpand}
       >
-        <EuiSpacer size="s" />
         {contentElement}
       </EuiAccordion>
     );
   }
 
   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

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>
   );
 }

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>
     </>

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']) => {

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',

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);

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' });
           }

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,15 +26,22 @@ 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 {};
     }
 
     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]);
 };