Add new InputModel class for managing input state (#171)

* Add an input model

* Set the input value using the input model

* Create a new input model when editing a message

* Use the 'inputChange' signal in the jupyterlab-chat model

* Use the chat model with the chat commands

* Remove a unused module, its functions have been moved to the chat model

* Allow adding or removing attachment when editing a message

* Passes the input model to the getChatCommand() method

* fix regression in how commands are handled

* Update the autocomplete if the cursor is at the begining of the input value

---------

Co-authored-by: David L. Qiu <[email protected]>

Author: brichet

docs/jupyter-chat-example/src/index.ts

@@ -34,10 +34,10 @@ class MyChatModel extends ChatModel {
       type: 'msg',
       time: Date.now() / 1000,
       sender: { username: 'me' },
-      attachments: this.inputAttachments
+      attachments: this.input.attachments
     };
     this.messageAdded(message);
-    this.clearAttachments();
+    this.input.clearAttachments();
   }
 }
 

packages/jupyter-chat/src/chat-commands/registry.ts

@@ -5,6 +5,7 @@
 
 import { Token } from '@lumino/coreutils';
 import { ChatCommand, IChatCommandProvider } from './types';
+import { IInputModel } from '../input-model';
 
 /**
  * Interface of a chat command registry, which tracks a list of chat command
@@ -19,11 +20,7 @@ export interface IChatCommandRegistry {
    * Handles a chat command by calling `handleChatCommand()` on the provider
    * corresponding to this chat command.
    */
-  handleChatCommand(
-    command: ChatCommand,
-    currentWord: string,
-    replaceCurrentWord: (newWord: string) => void
-  ): void;
+  handleChatCommand(command: ChatCommand, inputModel: IInputModel): void;
 }
 
 /**
@@ -42,11 +39,7 @@ export class ChatCommandRegistry implements IChatCommandRegistry {
     return Array.from(this._providers.values());
   }
 
-  handleChatCommand(
-    command: ChatCommand,
-    currentWord: string,
-    replaceCurrentWord: (newWord: string) => void
-  ) {
+  handleChatCommand(command: ChatCommand, inputModel: IInputModel) {
     const provider = this._providers.get(command.providerId);
     if (!provider) {
       console.error(
@@ -56,7 +49,7 @@ export class ChatCommandRegistry implements IChatCommandRegistry {
       return;
     }
 
-    provider.handleChatCommand(command, currentWord, replaceCurrentWord);
+    provider.handleChatCommand(command, inputModel);
   }
 
   private _providers: Map<string, IChatCommandProvider>;

packages/jupyter-chat/src/chat-commands/types.ts

@@ -4,6 +4,7 @@
  */
 
 import { LabIcon } from '@jupyterlab/ui-components';
+import { IInputModel } from '../input-model';
 
 export type ChatCommand = {
   /**
@@ -49,25 +50,18 @@ export interface IChatCommandProvider {
   id: string;
 
   /**
-   * Async function which accepts the current word and returns a list of
+   * Async function which accepts the input model and returns a list of
    * valid chat commands that match the current word. The current word is
    * space-separated word at the user's cursor.
-   *
-   * TODO: Pass a ChatModel/InputModel instance here to give the command access
-   * to more than the current word.
    */
-  getChatCommands(currentWord: string): Promise<ChatCommand[]>;
+  getChatCommands(inputModel: IInputModel): Promise<ChatCommand[]>;
 
   /**
    * Function called when a chat command is run by the user through the chat
    * commands menu.
-   *
-   * TODO: Pass a ChatModel/InputModel instance here to provide a function to
-   * replace the current word.
    */
   handleChatCommand(
     command: ChatCommand,
-    currentWord: string,
-    replaceCurrentWord: (newWord: string) => void
+    inputModel: IInputModel
   ): Promise<void>;
 }

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

@@ -6,6 +6,7 @@
 import { IDocumentManager } from '@jupyterlab/docmanager';
 import {
   Autocomplete,
+  AutocompleteInputChangeReason,
   Box,
   InputAdornment,
   SxProps,
@@ -17,33 +18,27 @@ import React, { useEffect, useRef, useState } from 'react';
 
 import { AttachmentPreviewList } from './attachments';
 import { AttachButton, CancelButton, SendButton } from './input';
-import { IChatModel } from '../model';
+import { IInputModel, InputModel } from '../input-model';
 import { IAutocompletionRegistry } from '../registry';
-import { IAttachment, IConfig, Selection } from '../types';
+import { IAttachment, Selection } from '../types';
 import { useChatCommands } from './input/use-chat-commands';
 import { IChatCommandRegistry } from '../chat-commands';
 
 const INPUT_BOX_CLASS = 'jp-chat-input-container';
 
 export function ChatInput(props: ChatInput.IProps): JSX.Element {
   const { documentManager, model } = props;
-  const [input, setInput] = useState<string>(props.value || '');
+  const [input, setInput] = useState<string>(model.value);
   const inputRef = useRef<HTMLInputElement>();
 
-  const chatCommands = useChatCommands(
-    input,
-    setInput,
-    inputRef,
-    props.chatCommandRegistry
-  );
+  const chatCommands = useChatCommands(model, props.chatCommandRegistry);
 
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState<boolean>(
     model.config.sendWithShiftEnter ?? false
   );
-  const [typingNotification, setTypingNotification] = useState<boolean>(
-    model.config.sendTypingNotification ?? false
+  const [attachments, setAttachments] = useState<IAttachment[]>(
+    model.attachments
   );
-  const [attachments, setAttachments] = useState<IAttachment[]>([]);
 
   // Display the include selection menu if it is not explicitly hidden, and if at least
   // one of the tool to check for text or cell selection is enabled.
@@ -53,9 +48,13 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
   }
 
   useEffect(() => {
-    const configChanged = (_: IChatModel, config: IConfig) => {
+    const inputChanged = (_: IInputModel, value: string) => {
+      setInput(value);
+    };
+    model.valueChanged.connect(inputChanged);
+
+    const configChanged = (_: IInputModel, config: InputModel.IConfig) => {
       setSendWithShiftEnter(config.sendWithShiftEnter ?? false);
-      setTypingNotification(config.sendTypingNotification ?? false);
     };
     model.configChanged.connect(configChanged);
 
@@ -66,15 +65,15 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
     };
     model.focusInputSignal?.connect(focusInputElement);
 
-    const attachmentChanged = (_: IChatModel, attachments: IAttachment[]) => {
+    const attachmentChanged = (_: IInputModel, attachments: IAttachment[]) => {
       setAttachments([...attachments]);
     };
-    model.inputAttachmentsChanged?.connect(attachmentChanged);
+    model.attachmentsChanged?.connect(attachmentChanged);
 
     return () => {
       model.configChanged?.disconnect(configChanged);
       model.focusInputSignal?.disconnect(focusInputElement);
-      model.inputAttachmentsChanged?.disconnect(attachmentChanged);
+      model.attachmentsChanged?.disconnect(attachmentChanged);
     };
   }, [model]);
 
@@ -160,15 +159,14 @@ ${selection.source}
 `;
     }
     props.onSend(content);
-    setInput('');
+    model.value = '';
   }
 
   /**
    * Triggered when cancelling edition.
    */
   function onCancel() {
-    setInput(props.value || '');
-    props.onCancel!();
+    props.onCancel?.();
   }
 
   // Set the helper text based on whether Shift+Enter is used for sending.
@@ -218,6 +216,9 @@ ${selection.source}
             placeholder="Start chatting"
             inputRef={inputRef}
             sx={{ marginTop: '1px' }}
+            onSelect={() =>
+              (model.cursorIndex = inputRef.current?.selectionStart ?? null)
+            }
             InputProps={{
               ...params.InputProps,
               endAdornment: (
@@ -247,10 +248,16 @@ ${selection.source}
           />
         )}
         inputValue={input}
-        onInputChange={(_, newValue: string) => {
-          setInput(newValue);
-          if (typingNotification && model.inputChanged) {
-            model.inputChanged(newValue);
+        onInputChange={(
+          _,
+          newValue: string,
+          reason: AutocompleteInputChangeReason
+        ) => {
+          // Do not update the value if the reason is 'reset', which should occur only
+          // if an autocompletion command has been selected. In this case, the value is
+          // set in the 'onChange()' callback of the autocompletion (to avoid conflicts).
+          if (reason !== 'reset') {
+            model.value = newValue;
           }
         }}
       />
@@ -269,11 +276,7 @@ export namespace ChatInput {
     /**
      * The chat model.
      */
-    model: IChatModel;
-    /**
-     * The initial value of the input (default to '')
-     */
-    value?: string;
+    model: IInputModel;
     /**
      * The function to be called to send the message.
      */

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

@@ -4,6 +4,7 @@
  */
 
 import { Button } from '@jupyter/react-components';
+import { IDocumentManager } from '@jupyterlab/docmanager';
 import { IRenderMimeRegistry } from '@jupyterlab/rendermime';
 import {
   LabIcon,
@@ -16,12 +17,14 @@ import type { SxProps, Theme } from '@mui/material';
 import clsx from 'clsx';
 import React, { useEffect, useState, useRef, forwardRef } from 'react';
 
+import { AttachmentPreviewList } from './attachments';
 import { ChatInput } from './chat-input';
 import { MarkdownRenderer } from './markdown-renderer';
 import { ScrollContainer } from './scroll-container';
+import { IChatCommandRegistry } from '../chat-commands';
+import { IInputModel, InputModel } from '../input-model';
 import { IChatModel } from '../model';
 import { IChatMessage, IUser } from '../types';
-import { AttachmentPreviewList } from './attachments';
 
 const MESSAGES_BOX_CLASS = 'jp-chat-messages-container';
 const MESSAGE_CLASS = 'jp-chat-message';
@@ -40,6 +43,8 @@ const NAVIGATION_BOTTOM_CLASS = 'jp-chat-navigation-bottom';
 type BaseMessageProps = {
   rmRegistry: IRenderMimeRegistry;
   model: IChatModel;
+  chatCommandRegistry?: IChatCommandRegistry;
+  documentManager?: IDocumentManager;
 };
 
 /**
@@ -338,6 +343,7 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
     const [deleted, setDeleted] = useState<boolean>(false);
     const [canEdit, setCanEdit] = useState<boolean>(false);
     const [canDelete, setCanDelete] = useState<boolean>(false);
+    const [inputModel, setInputModel] = useState<IInputModel | null>(null);
 
     // Look if the message can be deleted or edited.
     useEffect(() => {
@@ -353,6 +359,25 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       }
     }, [model, message]);
 
+    // Create an input model only if the message is edited.
+    useEffect(() => {
+      if (edit && canEdit) {
+        setInputModel(
+          new InputModel({
+            value: message.body,
+            activeCellManager: model.activeCellManager,
+            selectionWatcher: model.selectionWatcher,
+            config: {
+              sendWithShiftEnter: model.config.sendWithShiftEnter
+            },
+            attachments: message.attachments
+          })
+        );
+      } else {
+        setInputModel(null);
+      }
+    }, [edit]);
+
     // Cancel the current edition of the message.
     const cancelEdition = (): void => {
       setEdit(false);
@@ -366,6 +391,7 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       // Update the message
       const updatedMessage = { ...message };
       updatedMessage.body = input;
+      updatedMessage.attachments = inputModel?.attachments;
       model.updateMessage!(id, updatedMessage);
       setEdit(false);
     };
@@ -383,13 +409,14 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
       <div ref={ref} data-index={props.index}></div>
     ) : (
       <div ref={ref} data-index={props.index}>
-        {edit && canEdit ? (
+        {edit && canEdit && inputModel ? (
           <ChatInput
-            value={message.body}
             onSend={(input: string) => updateMessage(message.id, input)}
             onCancel={() => cancelEdition()}
-            model={model}
+            model={inputModel}
             hideIncludeSelection={true}
+            chatCommandRegistry={props.chatCommandRegistry}
+            documentManager={props.documentManager}
           />
         ) : (
           <MarkdownRenderer
@@ -401,7 +428,9 @@ export const ChatMessage = forwardRef<HTMLDivElement, ChatMessageProps>(
             rendered={props.renderedPromise}
           />
         )}
-        {message.attachments && (
+        {message.attachments && !edit && (
+          // Display the attachments only if message is not edited, otherwise the
+          // input component display them.
           <AttachmentPreviewList attachments={message.attachments} />
         )}
       </div>

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

@@ -32,7 +32,12 @@ export function ChatBody(props: Chat.IChatBodyProps): JSX.Element {
 
   return (
     <AttachmentOpenerContext.Provider value={props.attachmentOpenerRegistry}>
-      <ChatMessages rmRegistry={props.rmRegistry} model={model} />
+      <ChatMessages
+        rmRegistry={props.rmRegistry}
+        model={model}
+        chatCommandRegistry={props.chatCommandRegistry}
+        documentManager={props.documentManager}
+      />
       <ChatInput
         onSend={onSend}
         sx={{
@@ -42,7 +47,7 @@ export function ChatBody(props: Chat.IChatBodyProps): JSX.Element {
           paddingBottom: 0,
           borderTop: '1px solid var(--jp-border-color1)'
         }}
-        model={model}
+        model={model.input}
         documentManager={props.documentManager}
         autocompletionRegistry={props.autocompletionRegistry}
         chatCommandRegistry={props.chatCommandRegistry}

packages/jupyter-chat/src/components/input/send-button.tsx

@@ -8,9 +8,9 @@ import SendIcon from '@mui/icons-material/Send';
 import { Box, Menu, MenuItem, Typography } from '@mui/material';
 import React, { useCallback, useEffect, useState } from 'react';
 
-import { IChatModel } from '../../model';
 import { TooltippedButton } from '../mui-extras/tooltipped-button';
 import { includeSelectionIcon } from '../../icons';
+import { IInputModel } from '../../input-model';
 import { Selection } from '../../types';
 
 const SEND_BUTTON_CLASS = 'jp-chat-send-button';
@@ -21,7 +21,7 @@ const SEND_INCLUDE_LI_CLASS = 'jp-chat-send-include';
  * The send button props.
  */
 export type SendButtonProps = {
-  model: IChatModel;
+  model: IInputModel;
   sendWithShiftEnter: boolean;
   inputExists: boolean;
   onSend: (selection?: Selection) => unknown;