Run chat commands on message submission (#231)

* run chat commands on message submission & fix mention commands

* replace commands immediately on type

* add spaceOnAccept property

* fix multiple mentions

* hide users already mentioned

* simplify replacement logic

* add more info to spaceOnAccept docstring

* remove optional colon in mention regex

Author: dlqqq

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

@@ -23,7 +23,7 @@ const SEND_INCLUDE_LI_CLASS = 'jp-chat-send-include';
 export function SendButton(
   props: InputToolbarRegistry.IToolbarItemProps
 ): JSX.Element {
-  const { model } = props;
+  const { model, chatCommandRegistry } = props;
   const { activeCellManager, selectionWatcher } = model;
   const hideIncludeSelection = !activeCellManager || !selectionWatcher;
 
@@ -98,9 +98,17 @@ export function SendButton(
     };
   }, [activeCellManager, selectionWatcher]);
 
-  function sendWithSelection() {
+  async function send() {
+    await chatCommandRegistry?.onSubmit(model);
+    model.send(model.value);
+  }
+
+  async function sendWithSelection() {
     let source = '';
 
+    // Run all chat command providers
+    await chatCommandRegistry?.onSubmit(model);
+
     if (selectionWatcher?.selection) {
       // Append the selected text if exists.
       source = selectionWatcher.selection.text;
@@ -125,7 +133,7 @@ ${source}
   return (
     <>
       <TooltippedButton
-        onClick={() => model.send(model.value)}
+        onClick={send}
         disabled={disabled}
         tooltip={tooltip}
         buttonProps={{

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

@@ -105,7 +105,7 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
    * "Enter". This also handles many of the edge cases in the MUI Autocomplete
    * component.
    */
-  function handleKeyDown(event: React.KeyboardEvent<HTMLInputElement>) {
+  async function handleKeyDown(event: React.KeyboardEvent<HTMLInputElement>) {
     /**
      * IMPORTANT: This statement ensures that arrow keys can be used to navigate
      * the multiline input when the chat commands menu is closed.
@@ -157,6 +157,8 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
       (sendWithShiftEnter && event.shiftKey) ||
       (!sendWithShiftEnter && !event.shiftKey)
     ) {
+      // Run all command providers
+      await props.chatCommandRegistry?.onSubmit(model);
       model.send(input);
       event.stopPropagation();
       event.preventDefault();
@@ -218,7 +220,10 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
               endAdornment: (
                 <InputAdornment position="end" className={INPUT_TOOLBAR_CLASS}>
                   {toolbarElements.map(item => (
-                    <item.element model={model} />
+                    <item.element
+                      model={model}
+                      chatCommandRegistry={props.chatCommandRegistry}
+                    />
                   ))}
                 </InputAdornment>
               )

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

@@ -7,6 +7,7 @@ import * as React from 'react';
 import { AttachButton, CancelButton, SendButton } from './buttons';
 import { IInputModel } from '../../input-model';
 import { ISignal, Signal } from '@lumino/signaling';
+import { IChatCommandRegistry } from '../../registers';
 
 /**
  * The toolbar registry interface.
@@ -137,6 +138,11 @@ export namespace InputToolbarRegistry {
      * The input model of the input component including the button.
      */
     model: IInputModel;
+    /**
+     * Chat command registry. Should be used by the "Send" button to run
+     * `onSubmit()` on all command providers before sending the message.
+     */
+    chatCommandRegistry?: IChatCommandRegistry;
   }
 
   /**

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

@@ -45,6 +45,9 @@ export function useChatCommands(
   const [commands, setCommands] = useState<ChatCommand[]>([]);
 
   useEffect(() => {
+    /**
+     * Callback that runs whenever the current word changes.
+     */
     async function getCommands(_: IInputModel, currentWord: string | null) {
       const providers = chatCommandRegistry?.getProviders();
       if (!providers) {
@@ -58,12 +61,12 @@ export function useChatCommands(
         return;
       }
 
-      let newCommands: ChatCommand[] = [];
+      let commandCompletions: ChatCommand[] = [];
       for (const provider of providers) {
         // TODO: optimize performance when this method is truly async
         try {
-          newCommands = newCommands.concat(
-            await provider.getChatCommands(inputModel)
+          commandCompletions = commandCompletions.concat(
+            await provider.listCommandCompletions(inputModel)
           );
         } catch (e) {
           console.error(
@@ -72,10 +75,23 @@ export function useChatCommands(
           );
         }
       }
-      if (newCommands) {
-        setOpen(true);
+
+      // Immediately replace the current word if it exactly matches one command
+      // and 'replaceWith' is set.
+      if (
+        commandCompletions.length === 1 &&
+        commandCompletions[0].name === inputModel.currentWord &&
+        commandCompletions[0].replaceWith !== undefined
+      ) {
+        const replacement = commandCompletions[0].replaceWith;
+        inputModel.replaceCurrentWord(replacement);
+        return;
       }
-      setCommands(newCommands);
+
+      // Otherwise, open/close the menu based on the presence of command
+      // completions and set the menu entries.
+      setOpen(!!commandCompletions.length);
+      setCommands(commandCompletions);
     }
 
     inputModel.currentWordChanged.connect(getCommands);
@@ -87,7 +103,8 @@ export function useChatCommands(
 
   /**
    * onChange(): the callback invoked when a command is selected from the chat
-   * commands menu by the user.
+   * commands menu. When a command `cmd` is selected, this function replaces the
+   * current word with `cmd.replaceWith` if set, `cmd.name` otherwise.
    */
   const onChange: AutocompleteProps['onChange'] = (
     e: unknown,
@@ -109,14 +126,12 @@ export function useChatCommands(
       return;
     }
 
-    // if replaceWith is set, handle the command immediately
-    if (command.replaceWith) {
-      inputModel.replaceCurrentWord(command.replaceWith);
-      return;
+    let replacement =
+      command.replaceWith === undefined ? command.name : command.replaceWith;
+    if (command.spaceOnAccept) {
+      replacement += ' ';
     }
-
-    // otherwise, defer handling to the command provider
-    chatCommandRegistry.handleChatCommand(command, inputModel);
+    inputModel.replaceCurrentWord(replacement);
   };
 
   return {

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

@@ -21,14 +21,25 @@ export const IChatCommandRegistry = new Token<IChatCommandRegistry>(
  * and define how commands are handled when accepted in the chat commands menu.
  */
 export interface IChatCommandRegistry {
+  /**
+   * Adds a chat command provider to the registry.
+   */
   addProvider(provider: IChatCommandProvider): void;
+
+  /**
+   * Lists all chat command providers previously added via `addProvider()`.
+   */
   getProviders(): IChatCommandProvider[];
 
   /**
-   * Handles a chat command by calling `handleChatCommand()` on the provider
-   * corresponding to this chat command.
+   * Calls `onSubmit()` on each command provider in serial. Each command
+   * provider's `onSubmit()` method is responsible for checking the entire input
+   * for command calls and handling them accordingly.
+   *
+   * This method is called by the application after the user presses the "Send"
+   * button but before the message is sent to server.
    */
-  handleChatCommand(command: ChatCommand, inputModel: IInputModel): void;
+  onSubmit(inputModel: IInputModel): Promise<void>;
 }
 
 export type ChatCommand = {
@@ -56,13 +67,24 @@ export type ChatCommand = {
   description?: string;
 
   /**
-   * If set, Jupyter Chat will replace the current word with this string after
-   * the command is run from the chat commands menu.
+   * If set, Jupyter Chat will replace the current word with this string immediately when
+   * the command is accepted from the chat commands menu or the current word
+   * matches the command's `name` exactly.
    *
-   * If all commands from a provider have this property set, then
-   * `handleChatCommands()` can just return on the first line.
+   * This is generally used by "shortcut command" providers, e.g. the emoji
+   * command provider.
    */
   replaceWith?: string;
+
+  /**
+   * Specifies whether the application should add a space ' ' after the command
+   * is accepted from the menu. This should be set to `true` if the command that
+   * replaces the current word needs to be handled on submit, and the command is
+   * valid on its own.
+   *
+   * Defaults to `false`.
+   */
+  spaceOnAccept?: boolean;
 };
 
 /**
@@ -75,20 +97,23 @@ export interface IChatCommandProvider {
   id: string;
 
   /**
-   * 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.
+   * A method that should return the list of valid chat commands whose names
+   * complete the current word.
+   *
+   * The current word should be accessed from `inputModel.currentWord`.
    */
-  getChatCommands(inputModel: IInputModel): Promise<ChatCommand[]>;
+  listCommandCompletions(inputModel: IInputModel): Promise<ChatCommand[]>;
 
   /**
-   * Function called when a chat command is run by the user through the chat
-   * commands menu.
+   * A method that should identify and handle *all* command calls within a
+   * message that the user intends to submit. This method is called after a user
+   * presses the "Send" button but before the message is sent to the server.
+   *
+   * The entire message should be read from `inputModel.value`. This method may
+   * modify the new message before submission by setting `inputModel.value` or
+   * by calling other methods available on `inputModel`.
    */
-  handleChatCommand(
-    command: ChatCommand,
-    inputModel: IInputModel
-  ): Promise<void>;
+  onSubmit(inputModel: IInputModel): Promise<void>;
 }
 
 /**
@@ -107,17 +132,10 @@ export class ChatCommandRegistry implements IChatCommandRegistry {
     return Array.from(this._providers.values());
   }
 
-  handleChatCommand(command: ChatCommand, inputModel: IInputModel) {
-    const provider = this._providers.get(command.providerId);
-    if (!provider) {
-      console.error(
-        'Error in handling chat command: No command provider has an ID of ' +
-          command.providerId
-      );
-      return;
+  async onSubmit(inputModel: IInputModel) {
+    for (const provider of this._providers.values()) {
+      await provider.onSubmit(inputModel);
     }
-
-    provider.handleChatCommand(command, inputModel);
   }
 
   private _providers: Map<string, IChatCommandProvider>;

packages/jupyterlab-chat-extension/src/chat-commands/providers/emoji.ts

@@ -16,28 +16,28 @@ export class EmojiCommandProvider implements IChatCommandProvider {
   private _slash_commands: ChatCommand[] = [
     {
       name: ':heart:',
-      replaceWith: '❤ ',
+      replaceWith: '❤',
       providerId: this.id,
       description: 'Emoji',
       icon: '❤'
     },
     {
       name: ':smile:',
-      replaceWith: '🙂 ',
+      replaceWith: '🙂',
       providerId: this.id,
       description: 'Emoji',
       icon: '🙂'
     },
     {
       name: ':thinking:',
-      replaceWith: '🤔 ',
+      replaceWith: '🤔',
       providerId: this.id,
       description: 'Emoji',
       icon: '🤔'
     },
     {
       name: ':cool:',
-      replaceWith: '😎 ',
+      replaceWith: '😎',
       providerId: this.id,
       description: 'Emoji',
       icon: '😎'
@@ -47,7 +47,7 @@ export class EmojiCommandProvider implements IChatCommandProvider {
   // regex used to test the current word
   private _regex: RegExp = /^:\w*:?/;
 
-  async getChatCommands(inputModel: IInputModel) {
+  async listCommandCompletions(inputModel: IInputModel) {
     const match = inputModel.currentWord?.match(this._regex)?.[0];
     if (!match) {
       return [];
@@ -59,11 +59,10 @@ export class EmojiCommandProvider implements IChatCommandProvider {
     return commands;
   }
 
-  async handleChatCommand(
-    command: ChatCommand,
-    inputModel: IInputModel
-  ): Promise<void> {
-    // no handling needed because `replaceWith` is set in each command.
+  async onSubmit(inputModel: IInputModel) {
+    // this provider only provides commands that replace the current word
+    // when typed / selected from the menu. this method does not need to handle
+    // anything on submission.
     return;
   }
 }