General UI/UX improvements

packages/jupyter-ai/src/components/chat-input.tsx

@@ -2,13 +2,14 @@ import React from 'react';
 
 import {
   Box,
-  IconButton,
-  Input,
   SxProps,
+  TextField,
   Theme,
   FormGroup,
   FormControlLabel,
-  Checkbox
+  Checkbox,
+  IconButton,
+  InputAdornment
 } from '@mui/material';
 import SendIcon from '@mui/icons-material/Send';
 
@@ -25,23 +26,33 @@ type ChatInputProps = {
 };
 
 export function ChatInput(props: ChatInputProps): JSX.Element {
+  function handleKeyDown(event: React.KeyboardEvent<HTMLInputElement>) {
+    if (event.key === 'Enter' && event.shiftKey) {
+      props.onSend();
+      event.stopPropagation();
+      event.preventDefault();
+    }
+  }
   return (
     <Box sx={props.sx}>
       <Box sx={{ display: 'flex' }}>
-        <Input
+        <TextField
           value={props.value}
           onChange={e => props.onChange(e.target.value)}
+          fullWidth
+          variant="outlined"
           multiline
-          sx={{ flexGrow: 1 }}
+          onKeyDown={handleKeyDown}
+          InputProps={{
+            endAdornment: (
+               <InputAdornment position="end">
+                  <IconButton size="small" color="primary" onClick={props.onSend} disabled={!props.value.trim().length}>
+                     <SendIcon />
+                  </IconButton>
+               </InputAdornment>
+            ),
+         }}
         />
-        <IconButton
-          size="large"
-          color="primary"
-          onClick={props.onSend}
-          disabled={!props.value.trim().length}
-        >
-          <SendIcon />
-        </IconButton>
       </Box>
       {props.hasSelection && (
         <FormGroup sx={{ display: 'flex', flexDirection: 'row' }}>

packages/jupyter-ai/src/components/chat-messages.tsx

@@ -44,7 +44,7 @@ export function ChatMessageHeader(props: ChatMessageHeaderProps) {
         <Typography
           sx={{
             fontSize: 'var(--jp-ui-font-size1)',
-            color: 'var(--jp-ui-font-color1)'
+            color: 'var(--jp-ui-inverse-font-color1)'
           }}
         >
           {props.message.client.initials}
@@ -85,8 +85,8 @@ export function ChatMessageHeader(props: ChatMessageHeaderProps) {
           alignItems: 'center'
         }}
       >
-        <Typography sx={{ fontWeight: 700 }}>{name}</Typography>
-        <Typography sx={{ fontSize: '0.8em', fontWeight: 300 }}>
+        <Typography sx={{ fontWeight: 700, color: 'var(--jp-ui-font-color1)' }}>{name}</Typography>
+        <Typography sx={{ fontSize: '0.8em', color: 'var(--jp-ui-font-color2)', fontWeight: 300 }}>
           {props.timestamp}
         </Typography>
       </Box>
@@ -120,17 +120,20 @@ export function ChatMessages(props: ChatMessagesProps) {
 
   return (
     <Box
-      sx={{ '& > :not(:last-child)': { borderBottom: '1px solid lightgrey' } }}
+      sx={{ '& > :not(:last-child)': { borderBottom: '1px solid var(--jp-border-color2)' } }}
     >
       {props.messages.map((message, i) => (
         // extra div needed to ensure each bubble is on a new line
-        <Box key={i} sx={{ padding: 2 }}>
+        <Box key={i} sx={{ padding: 4 }}>
           <ChatMessageHeader
             message={message}
             timestamp={timestamps[message.id]}
-            sx={{ marginBottom: '12px' }}
+            sx={{ marginBottom: 3 }}
           />
           <ReactMarkdown
+            // We are using the jp-RenderedHTMLCommon class here to get the default Jupyter
+            // markdown styling and then overriding any CSS to make it more compact.
+            className="jp-RenderedHTMLCommon jp-ai-react-markdown"
             components={{
               code: ChatCodeView
             }}

packages/jupyter-ai/src/components/chat.tsx

@@ -63,7 +63,7 @@ function ChatBody({ chatHandler }: ChatBodyProps): JSX.Element {
 
     const prompt =
       input +
-      (includeSelection && selection?.text ? '\n--\n' + selection.text : '');
+      (includeSelection && selection?.text ? '\n\n```\n' + selection.text + '```': '');
 
     // send message to backend
     const messageId = await chatHandler.sendMessage({ prompt });
@@ -90,7 +90,7 @@ function ChatBody({ chatHandler }: ChatBodyProps): JSX.Element {
         width: '100%',
         height: '100%',
         boxSizing: 'border-box',
-        background: 'white',
+        background: 'var(--jp-layout-color0)',
         display: 'flex',
         flexDirection: 'column'
       }}
@@ -114,7 +114,11 @@ function ChatBody({ chatHandler }: ChatBodyProps): JSX.Element {
           setReplaceSelection(replaceSelection => !replaceSelection)
         }
         sx={{
-          padding: 2
+          paddingLeft: 4,
+          paddingRight: 4,
+          paddingTop: 2,
+          paddingBottom: 2,
+          borderTop: '1px solid var(--jp-border-color1)'
         }}
       />
     </Box>

packages/jupyter-ai/src/icons.ts

@@ -3,8 +3,14 @@
 import { LabIcon } from '@jupyterlab/ui-components';
 
 import psychologySvgStr from '../style/icons/psychology.svg';
+import chatSvgStr from '../style/icons/chat.svg';
 
 export const psychologyIcon = new LabIcon({
   name: 'jupyter-ai::psychology',
   svgstr: psychologySvgStr
 });
+
+export const chatIcon = new LabIcon({
+  name: 'jupyter-ai::chat',
+  svgstr: chatSvgStr
+});

packages/jupyter-ai/src/index.ts

@@ -1,6 +1,7 @@
 import {
   JupyterFrontEnd,
-  JupyterFrontEndPlugin
+  JupyterFrontEndPlugin,
+  ILayoutRestorer
 } from '@jupyterlab/application';
 import { INotebookTracker } from '@jupyterlab/notebook';
 import { IEditorTracker } from '@jupyterlab/fileeditor';
@@ -50,12 +51,13 @@ const plugin: JupyterFrontEndPlugin<void> = {
   id: 'jupyter_ai:plugin',
   autoStart: true,
   requires: [INotebookTracker, IEditorTracker],
-  optional: [IGlobalAwareness],
+  optional: [IGlobalAwareness, ILayoutRestorer],
   activate: async (
     app: JupyterFrontEnd,
     notebookTracker: INotebookTracker,
     editorTracker: IEditorTracker,
-    globalAwareness: Awareness | null
+    globalAwareness: Awareness | null,
+    restorer: ILayoutRestorer
   ) => {
     const { commands, shell } = app;
 
@@ -92,14 +94,20 @@ const plugin: JupyterFrontEndPlugin<void> = {
     const chatHandler = new ChatHandler();
     await chatHandler.initialize();
 
+    const chatWidget = buildChatSidebar(selectionWatcher, chatHandler, globalAwareness);
+
     /**
      * Add Chat widget to right sidebar
      */
     shell.add(
-      buildChatSidebar(selectionWatcher, chatHandler, globalAwareness),
-      'right'
+      chatWidget,
+      'left', { rank: 2000 }
     );
 
+    if (restorer) {
+      restorer.add(chatWidget, 'jupyter-ai-chat');
+    }
+
     /**
      * Register inserters
      */

packages/jupyter-ai/src/widgets/chat-sidebar.tsx

@@ -3,7 +3,7 @@ import { ReactWidget } from '@jupyterlab/apputils';
 import type { Awareness } from 'y-protocols/awareness';
 
 import { Chat } from '../components/chat';
-import { psychologyIcon } from '../icons';
+import { chatIcon } from '../icons';
 import { SelectionWatcher } from '../selection-watcher';
 import { ChatHandler } from '../chat_handler';
 
@@ -20,7 +20,7 @@ export function buildChatSidebar(
     />
   );
   ChatWidget.id = 'jupyter-ai::chat';
-  ChatWidget.title.icon = psychologyIcon;
+  ChatWidget.title.icon = chatIcon;
   ChatWidget.title.caption = 'Jupyter AI Chat'; // TODO: i18n
   return ChatWidget;
 }

packages/jupyter-ai/style/base.css

@@ -5,3 +5,4 @@
 */
 
 @import './expandable-text-field.css';
+@import './react-markdown.css';

packages/jupyter-ai/style/react-markdown.css

@@ -0,0 +1,7 @@
+.jp-RenderedHTMLCommon.jp-ai-react-markdown > pre {
+    margin: 0px;
+}
+
+.jp-RenderedHTMLCommon.jp-ai-react-markdown {
+    padding: 0px;
+}