Autocomplete commands (#57)

* Add optional autocomplete to the input of the chat * Add a registry to register autocompletion and provide the plugin in extensions * Fix tests and padding on send button * Add tests on autocompletion * Renaming for clarity, and make the registry explicitly implement its interface * Try to fix #outofband test * Fix #outofband test * Explicit plugin ID Co-authored-by: Jeremy Tuloup <[email protected]> --------- Co-authored-by: Jeremy Tuloup <[email protected]>
jupyterlab · Jul 4, 2024 · 9df1d48 · 9df1d48
1 parent 2a9d68f
commit 9df1d48
Show file tree

Hide file tree

Showing 11 changed files with 833 additions and 81 deletions.
diff --git a/packages/jupyter-chat/src/components/chat-input.tsx b/packages/jupyter-chat/src/components/chat-input.tsx
@@ -3,28 +3,99 @@
  * Distributed under the terms of the Modified BSD License.
  */
 
-import React, { useState } from 'react';
+import React, { useEffect, useRef, useState } from 'react';
 
 import {
+  Autocomplete,
   Box,
+  IconButton,
+  InputAdornment,
   SxProps,
   TextField,
-  Theme,
-  IconButton,
-  InputAdornment
+  Theme
 } from '@mui/material';
 import { Send, Cancel } from '@mui/icons-material';
 import clsx from 'clsx';
+import { AutocompleteCommand, IAutocompletionCommandsProps } from '../types';
+import { IAutocompletionRegistry } from '../registry';
 
 const INPUT_BOX_CLASS = 'jp-chat-input-container';
 const SEND_BUTTON_CLASS = 'jp-chat-send-button';
 const CANCEL_BUTTON_CLASS = 'jp-chat-cancel-button';
 
 export function ChatInput(props: ChatInput.IProps): JSX.Element {
+  const { autocompletionName, autocompletionRegistry, sendWithShiftEnter } =
+    props;
+  const autocompletion = useRef<IAutocompletionCommandsProps>();
   const [input, setInput] = useState<string>(props.value || '');
-  const { sendWithShiftEnter } = props;
+
+  // 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);
+
+  /**
+   * 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.
+   */
+  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 (event.key !== 'Enter') {
+      return;
+    }
+
+    // do not send the message if the user was selecting a suggested command from the
+    // Autocomplete component.
+    if (highlighted) {
+      return;
+    }
+
     if (
       event.key === 'Enter' &&
       ((sendWithShiftEnter && event.shiftKey) ||
@@ -65,49 +136,106 @@ export function ChatInput(props: ChatInput.IProps): JSX.Element {
 
   return (
     <Box sx={props.sx} className={clsx(INPUT_BOX_CLASS)}>
-      <Box sx={{ display: 'flex' }}>
-        <TextField
-          value={input}
-          onChange={e => setInput(e.target.value)}
-          fullWidth
-          variant="outlined"
-          multiline
-          onKeyDown={handleKeyDown}
-          placeholder="Start chatting"
-          InputProps={{
-            endAdornment: (
-              <InputAdornment position="end">
-                {props.onCancel && (
+      <Autocomplete
+        options={commandOptions}
+        value={props.value}
+        open={open}
+        autoHighlight
+        freeSolo
+        // ensure the autocomplete popup always renders on top
+        componentsProps={{
+          popper: {
+            placement: 'top'
+          },
+          paper: {
+            sx: {
+              border: '1px solid lightgray'
+            }
+          }
+        }}
+        ListboxProps={{
+          sx: {
+            '& .MuiAutocomplete-option': {
+              padding: 2
+            }
+          }
+        }}
+        renderInput={params => (
+          <TextField
+            {...params}
+            fullWidth
+            variant="outlined"
+            multiline
+            onKeyDown={handleKeyDown}
+            placeholder="Start chatting"
+            InputProps={{
+              ...params.InputProps,
+              endAdornment: (
+                <InputAdornment position="end">
+                  {props.onCancel && (
+                    <IconButton
+                      size="small"
+                      color="primary"
+                      onClick={onCancel}
+                      title={'Cancel edition'}
+                      className={clsx(CANCEL_BUTTON_CLASS)}
+                    >
+                      <Cancel />
+                    </IconButton>
+                  )}
                   <IconButton
                     size="small"
                     color="primary"
-                    onClick={onCancel}
-                    disabled={!input.trim().length}
-                    title={'Cancel edition'}
-                    className={clsx(CANCEL_BUTTON_CLASS)}
+                    onClick={onSend}
+                    disabled={
+                      props.onCancel
+                        ? input === props.value
+                        : !input.trim().length
+                    }
+                    title={`Send message ${sendWithShiftEnter ? '(SHIFT+ENTER)' : '(ENTER)'}`}
+                    className={clsx(SEND_BUTTON_CLASS)}
                   >
-                    <Cancel />
+                    <Send />
                   </IconButton>
-                )}
-                <IconButton
-                  size="small"
-                  color="primary"
-                  onClick={onSend}
-                  disabled={!input.trim().length}
-                  title={`Send message ${sendWithShiftEnter ? '(SHIFT+ENTER)' : '(ENTER)'}`}
-                  className={clsx(SEND_BUTTON_CLASS)}
-                >
-                  <Send />
-                </IconButton>
-              </InputAdornment>
-            )
-          }}
-          FormHelperTextProps={{
-            sx: { marginLeft: 'auto', marginRight: 0 }
-          }}
-          helperText={input.length > 2 ? helperText : ' '}
-        />
-      </Box>
+                </InputAdornment>
+              )
+            }}
+            FormHelperTextProps={{
+              sx: { marginLeft: 'auto', marginRight: 0 }
+            }}
+            helperText={input.length > 2 ? helperText : ' '}
+          />
+        )}
+        {...autocompletion.current?.props}
+        inputValue={input}
+        onInputChange={(_, newValue: string) => {
+          setInput(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>
   );
 }
@@ -140,5 +268,13 @@ export namespace ChatInput {
      * Custom mui/material styles.
      */
     sx?: SxProps<Theme>;
+    /**
+     * Autocompletion properties.
+     */
+    autocompletionRegistry?: IAutocompletionRegistry;
+    /**
+     * Autocompletion name.
+     */
+    autocompletionName?: string;
   }
 }
diff --git a/packages/jupyter-chat/src/components/chat.tsx b/packages/jupyter-chat/src/components/chat.tsx
@@ -15,16 +15,14 @@ import { JlThemeProvider } from './jl-theme-provider';
 import { ChatMessages } from './chat-messages';
 import { ChatInput } from './chat-input';
 import { IChatModel } from '../model';
+import { IAutocompletionRegistry } from '../registry';
 
-type ChatBodyProps = {
-  model: IChatModel;
-  rmRegistry: IRenderMimeRegistry;
-};
-
-function ChatBody({
-  model,
-  rmRegistry: renderMimeRegistry
-}: ChatBodyProps): JSX.Element {
+export function ChatBody(props: Chat.IChatBodyProps): JSX.Element {
+  const {
+    model,
+    rmRegistry: renderMimeRegistry,
+    autocompletionRegistry
+  } = props;
   // no need to append to messageGroups imperatively here. all of that is
   // handled by the listeners registered in the effect hooks above.
   const onSend = async (input: string) => {
@@ -45,6 +43,7 @@ function ChatBody({
           borderTop: '1px solid var(--jp-border-color1)'
         }}
         sendWithShiftEnter={model.config.sendWithShiftEnter ?? false}
+        autocompletionRegistry={autocompletionRegistry}
       />
     </>
   );
@@ -85,7 +84,11 @@ export function Chat(props: Chat.IOptions): JSX.Element {
         </Box>
         {/* body */}
         {view === Chat.View.chat && (
-          <ChatBody model={props.model} rmRegistry={props.rmRegistry} />
+          <ChatBody
+            model={props.model}
+            rmRegistry={props.rmRegistry}
+            autocompletionRegistry={props.autocompletionRegistry}
+          />
         )}
         {view === Chat.View.settings && props.settingsPanel && (
           <props.settingsPanel />
@@ -100,9 +103,9 @@ export function Chat(props: Chat.IOptions): JSX.Element {
  */
 export namespace Chat {
   /**
-   * The options to build the Chat UI.
+   * The props for the chat body component.
    */
-  export interface IOptions {
+  export interface IChatBodyProps {
     /**
      * The chat model.
      */
@@ -111,6 +114,20 @@ export namespace Chat {
      * The rendermime registry.
      */
     rmRegistry: IRenderMimeRegistry;
+    /**
+     * Autocompletion registry.
+     */
+    autocompletionRegistry?: IAutocompletionRegistry;
+    /**
+     * Autocompletion name.
+     */
+    autocompletionName?: string;
+  }
+
+  /**
+   * The options to build the Chat UI.
+   */
+  export interface IOptions extends IChatBodyProps {
     /**
      * The theme manager.
      */

diff --git a/packages/jupyter-chat/src/index.ts b/packages/jupyter-chat/src/index.ts
@@ -5,6 +5,7 @@
 
 export * from './icons';
 export * from './model';
+export * from './registry';
 export * from './types';
 export * from './widgets/chat-error';
 export * from './widgets/chat-sidebar';