Rework selection inclusion; new Send button UX (#905)

* remove include selection checkbox

* add new selection types to backend

* add new TooltippedButton component

* add `prompt` field to `HumanChatMessage`

This allows chat handlers to distinguish between the message prompt,
selection, and body (prompt + selection).

- The backend builds the message body from each `ChatRequest`, which
  consists of a `prompt` and a `selection`.

- The body will be used by most chat handlers, and will be used to
  render the messages in the UI.

- `/fix` uses `prompt` and `selection` separately as it does additional
  processing on each component. `/fix` does not use `body`.

* add new include selection icon

* implement new send button

* pre-commit

* fix unit tests

* pre-commit

* edit cell selection tooltip to indicate output is not included

* allow Enter to open dropdown menu

* fix double-send bug when using keyboard

* re-enable auto focus for first item

packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py

@@ -89,7 +89,7 @@ async def process_message(self, message: HumanChatMessage):
         selection: CellWithErrorSelection = message.selection
 
         # parse additional instructions specified after `/fix`
-        extra_instructions = message.body[4:].strip() or "None."
+        extra_instructions = message.prompt[4:].strip() or "None."
 
         self.get_llm_chain()
         with self.pending("Analyzing error"):

packages/jupyter-ai/jupyter_ai/handlers.py

@@ -247,12 +247,17 @@ async def on_message(self, message):
             self.log.error(e)
             return
 
+        message_body = chat_request.prompt
+        if chat_request.selection:
+            message_body += f"\n\n```\n{chat_request.selection.source}\n```\n"
+
         # message broadcast to chat clients
         chat_message_id = str(uuid.uuid4())
         chat_message = HumanChatMessage(
             id=chat_message_id,
             time=time.time(),
-            body=chat_request.prompt,
+            body=message_body,
+            prompt=chat_request.prompt,
             selection=chat_request.selection,
             client=self.chat_client,
         )

packages/jupyter-ai/jupyter_ai/models.py

@@ -14,22 +14,28 @@ class CellError(BaseModel):
     traceback: List[str]
 
 
+class TextSelection(BaseModel):
+    type: Literal["text"] = "text"
+    source: str
+
+
+class CellSelection(BaseModel):
+    type: Literal["cell"] = "cell"
+    source: str
+
+
 class CellWithErrorSelection(BaseModel):
     type: Literal["cell-with-error"] = "cell-with-error"
     source: str
     error: CellError
 
 
-Selection = Union[CellWithErrorSelection]
+Selection = Union[TextSelection, CellSelection, CellWithErrorSelection]
 
 
 # the type of message used to chat with the agent
 class ChatRequest(BaseModel):
     prompt: str
-    # TODO: This currently is only used when a user runs the /fix slash command.
-    # In the future, the frontend should set the text selection on this field in
-    # the `HumanChatMessage` it sends to JAI, instead of appending the text
-    # selection to `body` in the frontend.
     selection: Optional[Selection]
 
 
@@ -88,8 +94,13 @@ class HumanChatMessage(BaseModel):
     id: str
     time: float
     body: str
-    client: ChatClient
+    """The formatted body of the message to be rendered in the UI. Includes both
+    `prompt` and `selection`."""
+    prompt: str
+    """The prompt typed into the chat input by the user."""
     selection: Optional[Selection]
+    """The selection included with the prompt, if any."""
+    client: ChatClient
 
 
 class ClearMessage(BaseModel):

packages/jupyter-ai/jupyter_ai/tests/test_handlers.py

@@ -89,7 +89,13 @@ def chat_client():
 
 @pytest.fixture
 def human_chat_message(chat_client):
-    return HumanChatMessage(id="test", time=0, body="test message", client=chat_client)
+    return HumanChatMessage(
+        id="test",
+        time=0,
+        body="test message",
+        prompt="test message",
+        client=chat_client,
+    )
 
 
 def test_learn_index_permissions(tmp_path):

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

@@ -6,9 +6,6 @@ import {
   SxProps,
   TextField,
   Theme,
-  FormGroup,
-  FormControlLabel,
-  Checkbox,
   InputAdornment,
   Typography
 } from '@mui/material';
@@ -27,15 +24,11 @@ import { ISignal } from '@lumino/signaling';
 import { AiService } from '../handler';
 import { SendButton, SendButtonProps } from './chat-input/send-button';
 import { useActiveCellContext } from '../contexts/active-cell-context';
+import { ChatHandler } from '../chat_handler';
 
 type ChatInputProps = {
-  value: string;
-  onChange: (newValue: string) => unknown;
-  onSend: (selection?: AiService.Selection) => unknown;
-  hasSelection: boolean;
-  includeSelection: boolean;
+  chatHandler: ChatHandler;
   focusInputSignal: ISignal<unknown, void>;
-  toggleIncludeSelection: () => unknown;
   sendWithShiftEnter: boolean;
   sx?: SxProps<Theme>;
   /**
@@ -105,6 +98,7 @@ function renderSlashCommandOption(
 }
 
 export function ChatInput(props: ChatInputProps): JSX.Element {
+  const [input, setInput] = useState('');
   const [slashCommandOptions, setSlashCommandOptions] = useState<
     SlashCommandOption[]
   >([]);
@@ -159,24 +153,24 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
    * chat input. Close the autocomplete when the user clears the chat input.
    */
   useEffect(() => {
-    if (props.value === '/') {
+    if (input === '/') {
       setOpen(true);
       return;
     }
 
-    if (props.value === '') {
+    if (input === '') {
       setOpen(false);
       return;
     }
-  }, [props.value]);
+  }, [input]);
 
   /**
    * Effect: Set current slash command
    */
   useEffect(() => {
-    const matchedSlashCommand = props.value.match(/^\s*\/\w+/);
+    const matchedSlashCommand = input.match(/^\s*\/\w+/);
     setCurrSlashCommand(matchedSlashCommand && matchedSlashCommand[0]);
-  }, [props.value]);
+  }, [input]);
 
   /**
    * Effect: ensure that the `highlighted` is never `true` when `open` is
@@ -190,25 +184,27 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
     }
   }, [open, highlighted]);
 
-  // TODO: unify the `onSend` implementation in `chat.tsx` and here once text
-  // selection is refactored.
-  function onSend() {
-    // case: /fix
+  function onSend(selection?: AiService.Selection) {
+    const prompt = input;
+    setInput('');
+
+    // if the current slash command is `/fix`, we always include a code cell
+    // with error output in the selection.
     if (currSlashCommand === '/fix') {
       const cellWithError = activeCell.manager.getContent(true);
       if (!cellWithError) {
         return;
       }
 
-      props.onSend({
-        ...cellWithError,
-        type: 'cell-with-error'
+      props.chatHandler.sendMessage({
+        prompt,
+        selection: { ...cellWithError, type: 'cell-with-error' }
       });
       return;
     }
 
-    // default case
-    props.onSend();
+    // otherwise, send a ChatRequest with the prompt and selection
+    props.chatHandler.sendMessage({ prompt, selection });
   }
 
   function handleKeyDown(event: React.KeyboardEvent<HTMLDivElement>) {
@@ -244,7 +240,7 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
     </span>
   );
 
-  const inputExists = !!props.value.trim();
+  const inputExists = !!input.trim();
   const sendButtonProps: SendButtonProps = {
     onSend,
     sendWithShiftEnter: props.sendWithShiftEnter,
@@ -258,9 +254,9 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
       <Autocomplete
         autoHighlight
         freeSolo
-        inputValue={props.value}
+        inputValue={input}
         onInputChange={(_, newValue: string) => {
-          props.onChange(newValue);
+          setInput(newValue);
         }}
         onHighlightChange={
           /**
@@ -322,23 +318,10 @@ export function ChatInput(props: ChatInputProps): JSX.Element {
             FormHelperTextProps={{
               sx: { marginLeft: 'auto', marginRight: 0 }
             }}
-            helperText={props.value.length > 2 ? helperText : ' '}
+            helperText={input.length > 2 ? helperText : ' '}
           />
         )}
       />
-      {props.hasSelection && (
-        <FormGroup sx={{ display: 'flex', flexDirection: 'row' }}>
-          <FormControlLabel
-            control={
-              <Checkbox
-                checked={props.includeSelection}
-                onChange={props.toggleIncludeSelection}
-              />
-            }
-            label="Include selection"
-          />
-        </FormGroup>
-      )}
     </Box>
   );
 }