Define a new framework for chat commands

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

@@ -0,0 +1,7 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+export * from './types';
+export * from './registry';

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

@@ -0,0 +1,67 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { Token } from '@lumino/coreutils';
+import { ChatCommand, IChatCommandProvider } from './types';
+
+/**
+ * Interface of a chat command registry, which tracks a list of chat command
+ * providers. Providers provide a list of commands given a user's partial input,
+ * and define how commands are handled when accepted in the chat commands menu.
+ */
+export interface IChatCommandRegistry {
+  addProvider(provider: IChatCommandProvider): void;
+  getProviders(): IChatCommandProvider[];
+
+  /**
+   * 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;
+}
+
+/**
+ * Default chat command registry implementation.
+ */
+export class ChatCommandRegistry implements IChatCommandRegistry {
+  constructor() {
+    this._providers = new Map<string, IChatCommandProvider>();
+  }
+
+  addProvider(provider: IChatCommandProvider): void {
+    this._providers.set(provider.id, provider);
+  }
+
+  getProviders(): IChatCommandProvider[] {
+    return Array.from(this._providers.values());
+  }
+
+  handleChatCommand(
+    command: ChatCommand,
+    currentWord: string,
+    replaceCurrentWord: (newWord: string) => void
+  ) {
+    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;
+    }
+
+    provider.handleChatCommand(command, currentWord, replaceCurrentWord);
+  }
+
+  private _providers: Map<string, IChatCommandProvider>;
+}
+
+export const IChatCommandRegistry = new Token<IChatCommandRegistry>(
+  '@jupyter/chat:IChatCommandRegistry'
+);

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

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) Jupyter Development Team.
+ * Distributed under the terms of the Modified BSD License.
+ */
+
+import { LabIcon } from '@jupyterlab/ui-components';
+
+export type ChatCommand = {
+  /**
+   * The name of the command. This defines what the user should type in the
+   * input to have the command appear in the chat commands menu.
+   */
+  name: string;
+
+  /**
+   * ID of the provider the command originated from.
+   */
+  providerId: string;
+
+  /**
+   * If set, this will be rendered as the icon for the command in the chat
+   * commands menu. Jupyter Chat will choose a default if this is unset.
+   */
+  icon?: LabIcon;
+
+  /**
+   * If set, this will be rendered as the description for the command in the
+   * chat commands menu. Jupyter Chat will choose a default if this is unset.
+   */
+  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 all commands from a provider have this property set, then
+   * `handleChatCommands()` can just return on the first line.
+   */
+  replaceWith?: string;
+};
+
+/**
+ * Interface of a command provider.
+ */
+export interface IChatCommandProvider {
+  /**
+   * ID of this command provider.
+   */
+  id: string;
+
+  /**
+   * Async function which accepts the current word 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[]>;
+
+  /**
+   * 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
+  ): Promise<void>;
+}

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

@@ -19,19 +19,24 @@ import { CancelButton } from './input/cancel-button';
 import { SendButton } from './input/send-button';
 import { IChatModel } from '../model';
 import { IAutocompletionRegistry } from '../registry';
-import {
-  AutocompleteCommand,
-  IAutocompletionCommandsProps,
-  IConfig,
-  Selection
-} from '../types';
+import { IConfig, 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 { autocompletionName, autocompletionRegistry, model } = props;
-  const autocompletion = useRef<IAutocompletionCommandsProps>();
+  const { model } = props;
   const [input, setInput] = useState<string>(props.value || '');
+  const inputRef = useRef<HTMLInputElement>();
+
+  const chatCommands = useChatCommands(
+    input,
+    setInput,
+    inputRef,
+    props.chatCommandRegistry
+  );
+
   const [sendWithShiftEnter, setSendWithShiftEnter] = useState<boolean>(
     model.config.sendWithShiftEnter ?? false
   );
@@ -46,9 +51,6 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
     hideIncludeSelection = true;
   }
 
-  // store reference to the input element to enable focusing it easily
-  const inputRef = useRef<HTMLInputElement>();
-
   useEffect(() => {
     const configChanged = (_: IChatModel, config: IConfig) => {
       setSendWithShiftEnter(config.sendWithShiftEnter ?? false);
@@ -69,87 +71,62 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
     };
   }, [model]);
 
-  // The autocomplete commands options.
-  const [commandOptions, setCommandOptions] = useState<AutocompleteCommand[]>(
-    []
-  );
-  // whether any option is highlighted in the slash command autocomplete
-  const [highlighted, setHighlighted] = useState<boolean>(false);
-  // controls whether the slash command autocomplete is open
-  const [open, setOpen] = useState<boolean>(false);
-
   const inputExists = !!input.trim();
 
   /**
-   * Effect: fetch the list of available autocomplete commands.
-   */
-  useEffect(() => {
-    if (autocompletionRegistry === undefined) {
-      return;
-    }
-    autocompletion.current = autocompletionName
-      ? autocompletionRegistry.get(autocompletionName)
-      : autocompletionRegistry.getDefaultCompletion();
-
-    if (autocompletion.current === undefined) {
-      return;
-    }
-
-    if (Array.isArray(autocompletion.current.commands)) {
-      setCommandOptions(autocompletion.current.commands);
-    } else if (typeof autocompletion.current.commands === 'function') {
-      autocompletion.current
-        .commands()
-        .then((commands: AutocompleteCommand[]) => {
-          setCommandOptions(commands);
-        });
-    }
-  }, []);
-
-  /**
-   * Effect: Open the autocomplete when the user types the 'opener' string into an
-   * empty chat input. Close the autocomplete and reset the last selected value when
-   * the user clears the chat input.
+   * `handleKeyDown()`: callback invoked when the user presses any key in the
+   * `TextField` component. This is used to send the message when a user presses
+   * "Enter". This also handles many of the edge cases in the MUI Autocomplete
+   * component.
    */
-  useEffect(() => {
-    if (!autocompletion.current?.opener) {
-      return;
-    }
-
-    if (input === autocompletion.current?.opener) {
-      setOpen(true);
-      return;
-    }
-
-    if (input === '') {
-      setOpen(false);
-      return;
-    }
-  }, [input]);
-
   function handleKeyDown(event: React.KeyboardEvent<HTMLInputElement>) {
-    if (['ArrowDown', 'ArrowUp'].includes(event.key) && !open) {
+    /**
+     * IMPORTANT: This statement ensures that arrow keys can be used to navigate
+     * the multiline input when the chat commands menu is closed.
+     */
+    if (
+      ['ArrowDown', 'ArrowUp'].includes(event.key) &&
+      !chatCommands.menu.open
+    ) {
       event.stopPropagation();
       return;
     }
 
+    // remainder of this function only handles the "Enter" key.
     if (event.key !== 'Enter') {
       return;
     }
 
-    // Do not send the message if the user was selecting a suggested command from the
-    // Autocomplete component.
-    if (highlighted) {
+    /**
+     * IMPORTANT: This statement ensures that when the chat commands menu is
+     * open with a highlighted command, the "Enter" key should run that command
+     * instead of sending the message.
+     *
+     * This is done by returning early and letting the event propagate to the
+     * `Autocomplete` component.
+     */
+    if (chatCommands.menu.highlighted) {
       return;
     }
 
+    // remainder of this function only handles the "Enter" key pressed while the
+    // commands menu is closed.
+    /**
+     * IMPORTANT: This ensures that when the "Enter" key is pressed with the
+     * commands menu closed, the event is not propagated up to the
+     * `Autocomplete` component. Without this, `Autocomplete.onChange()` gets
+     * called with an invalid `string` instead of a `ChatCommand`.
+     */
+    event.stopPropagation();
+
     // Do not send empty messages, and avoid adding new line in empty message.
     if (!inputExists) {
       event.stopPropagation();
       event.preventDefault();
       return;
     }
 
+    // Finally, send the message when all other conditions are met.
     if (
       (sendWithShiftEnter && event.shiftKey) ||
       (!sendWithShiftEnter && !event.shiftKey)
@@ -201,11 +178,7 @@ ${selection.source}
   return (
     <Box sx={props.sx} className={clsx(INPUT_BOX_CLASS)}>
       <Autocomplete
-        options={commandOptions}
-        value={props.value}
-        open={open}
-        autoHighlight
-        freeSolo
+        {...chatCommands.autocompleteProps}
         // ensure the autocomplete popup always renders on top
         componentsProps={{
           popper: {
@@ -255,38 +228,13 @@ ${selection.source}
             helperText={input.length > 2 ? helperText : ' '}
           />
         )}
-        {...autocompletion.current?.props}
         inputValue={input}
         onInputChange={(_, newValue: string) => {
           setInput(newValue);
           if (typingNotification && model.inputChanged) {
             model.inputChanged(newValue);
           }
         }}
-        onHighlightChange={
-          /**
-           * On highlight change: set `highlighted` to whether an option is
-           * highlighted by the user.
-           *
-           * This isn't called when an option is selected for some reason, so we
-           * need to call `setHighlighted(false)` in `onClose()`.
-           */
-          (_, highlightedOption) => {
-            setHighlighted(!!highlightedOption);
-          }
-        }
-        onClose={
-          /**
-           * On close: set `highlighted` to `false` and close the popup by
-           * setting `open` to `false`.
-           */
-          () => {
-            setHighlighted(false);
-            setOpen(false);
-          }
-        }
-        // hide default extra right padding in the text field
-        disableClearable
       />
     </Box>
   );
@@ -332,5 +280,9 @@ export namespace ChatInput {
      * Autocompletion name.
      */
     autocompletionName?: string;
+    /**
+     * Chat command registry.
+     */
+    chatCommandRegistry?: IChatCommandRegistry;
   }
 }