diff --git a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py b/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py index f8c9f6f6b..318f9a5dd 100644 --- a/packages/jupyter-ai/jupyter_ai/chat_handlers/fix.py +++ b/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"): diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index 9f3fcee71..a08a87df9 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/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, ) diff --git a/packages/jupyter-ai/jupyter_ai/models.py b/packages/jupyter-ai/jupyter_ai/models.py index 331f4cd94..f2fa098bb 100644 --- a/packages/jupyter-ai/jupyter_ai/models.py +++ b/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): diff --git a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py b/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py index 6fab65763..d2e73ce6c 100644 --- a/packages/jupyter-ai/jupyter_ai/tests/test_handlers.py +++ b/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): diff --git a/packages/jupyter-ai/src/components/chat-input.tsx b/packages/jupyter-ai/src/components/chat-input.tsx index 64ed7447c..4ab9bcd36 100644 --- a/packages/jupyter-ai/src/components/chat-input.tsx +++ b/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; - toggleIncludeSelection: () => unknown; sendWithShiftEnter: boolean; sx?: SxProps; /** @@ -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) { @@ -244,7 +240,7 @@ export function ChatInput(props: ChatInputProps): JSX.Element { ); - 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 { { - 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 && ( - - - } - label="Include selection" - /> - - )} ); } diff --git a/packages/jupyter-ai/src/components/chat-input/send-button.tsx b/packages/jupyter-ai/src/components/chat-input/send-button.tsx index 6ab79a355..69ad3efc6 100644 --- a/packages/jupyter-ai/src/components/chat-input/send-button.tsx +++ b/packages/jupyter-ai/src/components/chat-input/send-button.tsx @@ -1,10 +1,18 @@ -import React from 'react'; +import React, { useCallback, useState } from 'react'; +import { Box, Menu, MenuItem, Typography } from '@mui/material'; +import KeyboardArrowDown from '@mui/icons-material/KeyboardArrowDown'; import SendIcon from '@mui/icons-material/Send'; -import { TooltippedIconButton } from '../mui-extras/tooltipped-icon-button'; +import { TooltippedButton } from '../mui-extras/tooltipped-button'; +import { includeSelectionIcon } from '../../icons'; +import { useActiveCellContext } from '../../contexts/active-cell-context'; +import { useSelectionContext } from '../../contexts/selection-context'; +import { AiService } from '../../handler'; + +const FIX_TOOLTIP = '/fix requires an active code cell with an error'; export type SendButtonProps = { - onSend: () => unknown; + onSend: (selection?: AiService.Selection) => unknown; sendWithShiftEnter: boolean; currSlashCommand: string | null; inputExists: boolean; @@ -12,34 +20,163 @@ export type SendButtonProps = { }; export function SendButton(props: SendButtonProps): JSX.Element { + const [menuAnchorEl, setMenuAnchorEl] = useState(null); + const [menuOpen, setMenuOpen] = useState(false); + const [textSelection] = useSelectionContext(); + const activeCell = useActiveCellContext(); + + const openMenu = useCallback((el: HTMLElement | null) => { + setMenuAnchorEl(el); + setMenuOpen(true); + }, []); + + const closeMenu = useCallback(() => { + setMenuOpen(false); + }, []); + const disabled = props.currSlashCommand === '/fix' ? !props.inputExists || !props.activeCellHasError : !props.inputExists; + const includeSelectionDisabled = !(activeCell.exists || textSelection); + + const includeSelectionTooltip = + props.currSlashCommand === '/fix' + ? FIX_TOOLTIP + : textSelection + ? `${textSelection.text.split('\n').length} lines selected` + : activeCell.exists + ? 'Code from 1 active cell' + : 'No selection or active cell'; + const defaultTooltip = props.sendWithShiftEnter ? 'Send message (SHIFT+ENTER)' : 'Send message (ENTER)'; const tooltip = props.currSlashCommand === '/fix' && !props.activeCellHasError - ? '/fix requires a code cell with an error output selected' + ? FIX_TOOLTIP : !props.inputExists ? 'Message must not be empty' : defaultTooltip; + function sendWithSelection() { + // if the current slash command is `/fix`, `props.onSend()` should always + // include the code cell with error output, so the `selection` argument does + // not need to be defined. + if (props.currSlashCommand === '/fix') { + props.onSend(); + closeMenu(); + return; + } + + // otherwise, parse the text selection or active cell, with the text + // selection taking precedence. + if (textSelection?.text) { + props.onSend({ + type: 'text', + source: textSelection.text + }); + closeMenu(); + return; + } + + if (activeCell.exists) { + props.onSend({ + type: 'cell', + source: activeCell.manager.getContent(false).source + }); + closeMenu(); + return; + } + } + return ( - props.onSend()} - disabled={disabled} - tooltip={tooltip} - iconButtonProps={{ - size: 'small', - color: 'primary', - title: defaultTooltip - }} - > - - + + props.onSend()} + disabled={disabled} + tooltip={tooltip} + buttonProps={{ + size: 'small', + title: defaultTooltip, + variant: 'contained' + }} + sx={{ + minWidth: 'unset', + borderRadius: '2px 0px 0px 2px' + }} + > + + + { + openMenu(e.currentTarget); + }} + disabled={disabled} + tooltip="" + buttonProps={{ + variant: 'contained', + onKeyDown: e => { + if (e.key !== 'Enter' && e.key !== ' ') { + return; + } + openMenu(e.currentTarget); + // stopping propagation of this event prevents the prompt from being + // sent when the dropdown button is selected and clicked via 'Enter'. + e.stopPropagation(); + } + }} + sx={{ + minWidth: 'unset', + padding: '4px 0px', + borderRadius: '0px 2px 2px 0px', + borderLeft: '1px solid white' + }} + > + + + + { + sendWithSelection(); + // prevent sending second message with no selection + e.stopPropagation(); + }} + disabled={includeSelectionDisabled} + > + + + Send message with selection + + {includeSelectionTooltip} + + + + + ); } diff --git a/packages/jupyter-ai/src/components/chat-messages.tsx b/packages/jupyter-ai/src/components/chat-messages.tsx index 37d1f02ac..ec2e0cf1a 100644 --- a/packages/jupyter-ai/src/components/chat-messages.tsx +++ b/packages/jupyter-ai/src/components/chat-messages.tsx @@ -201,11 +201,6 @@ export function ChatMessages(props: ChatMessagesProps): JSX.Element { }} > {sortedMessages.map(message => { - // render selection in HumanChatMessage, if any - const markdownStr = - message.type === 'human' && message.selection - ? message.body + '\n\n```\n' + message.selection.source + '\n```\n' - : message.body; return ( (false); - const [includeSelection, setIncludeSelection] = useState(true); - const [input, setInput] = useState(''); - const [textSelection] = useSelectionContext(); const [sendWithShiftEnter, setSendWithShiftEnter] = useState(true); /** @@ -107,25 +101,6 @@ function ChatBody({ }; }, [chatHandler]); - // no need to append to messageGroups imperatively here. all of that is - // handled by the listeners registered in the effect hooks above. - // TODO: unify how text selection & cell selection are handled - const onSend = async (selection?: AiService.Selection) => { - setInput(''); - - const prompt = - input + - (includeSelection && textSelection?.text - ? '\n\n```\n' + textSelection.text + '\n```' - : ''); - - // send message to backend - const messageId = await chatHandler.sendMessage({ prompt, selection }); - - // await reply from agent - await chatHandler.replyFor(messageId); - }; - const openSettingsView = () => { setShowWelcomeMessage(false); chatViewHandler(ChatView.Settings); @@ -168,15 +143,8 @@ function ChatBody({ - setIncludeSelection(includeSelection => !includeSelection) - } sx={{ paddingLeft: 4, paddingRight: 4, diff --git a/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx b/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx new file mode 100644 index 000000000..363a4e832 --- /dev/null +++ b/packages/jupyter-ai/src/components/mui-extras/tooltipped-button.tsx @@ -0,0 +1,87 @@ +import React from 'react'; +import { Button, ButtonProps, SxProps, TooltipProps } from '@mui/material'; + +import { ContrastingTooltip } from './contrasting-tooltip'; + +export type TooltippedButtonProps = { + onClick: React.MouseEventHandler; + tooltip: string; + children: JSX.Element; + disabled?: boolean; + placement?: TooltipProps['placement']; + /** + * The offset of the tooltip popup. + * + * The expected syntax is defined by the Popper library: + * https://popper.js.org/docs/v2/modifiers/offset/ + */ + offset?: [number, number]; + 'aria-label'?: string; + /** + * Props passed directly to the MUI `Button` component. + */ + buttonProps?: ButtonProps; + /** + * Styles applied to the MUI `Button` component. + */ + sx?: SxProps; +}; + +/** + * A component that renders an MUI `Button` with a high-contrast tooltip + * provided by `ContrastingTooltip`. This component differs from the MUI + * defaults in the following ways: + * + * - Shows the tooltip on hover even if disabled. + * - Renders the tooltip above the button by default. + * - Renders the tooltip closer to the button by default. + * - Lowers the opacity of the Button when disabled. + * - Renders the Button with `line-height: 0` to avoid showing extra + * vertical space in SVG icons. + * + * NOTE TO DEVS: Please keep this component's features synchronized with + * features available to `TooltippedIconButton`. + */ +export function TooltippedButton(props: TooltippedButtonProps): JSX.Element { + return ( + + {/* + By default, tooltips never appear when the Button is disabled. The + official way to support this feature in MUI is to wrap the child Button + element in a `span` element. + + See: https://mui.com/material-ui/react-tooltip/#disabled-elements + */} + + + + + ); +} diff --git a/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx b/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx index 6f97c7462..31279b4d9 100644 --- a/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx +++ b/packages/jupyter-ai/src/components/mui-extras/tooltipped-icon-button.tsx @@ -1,10 +1,15 @@ import React from 'react'; -import { IconButton, IconButtonProps, TooltipProps } from '@mui/material'; +import { + IconButton, + IconButtonProps, + SxProps, + TooltipProps +} from '@mui/material'; import { ContrastingTooltip } from './contrasting-tooltip'; export type TooltippedIconButtonProps = { - onClick: () => unknown; + onClick: React.MouseEventHandler; tooltip: string; children: JSX.Element; disabled?: boolean; @@ -21,6 +26,10 @@ export type TooltippedIconButtonProps = { * Props passed directly to the MUI `IconButton` component. */ iconButtonProps?: IconButtonProps; + /** + * Styles applied to the MUI `IconButton` component. + */ + sx?: SxProps; }; /** @@ -34,6 +43,9 @@ export type TooltippedIconButtonProps = { * - Lowers the opacity of the IconButton when disabled. * - Renders the IconButton with `line-height: 0` to avoid showing extra * vertical space in SVG icons. + * + * NOTE TO DEVS: Please keep this component's features synchronized with + * features available to `TooltippedButton`. */ export function TooltippedIconButton( props: TooltippedIconButtonProps @@ -62,12 +74,16 @@ export function TooltippedIconButton( See: https://mui.com/material-ui/react-tooltip/#disabled-elements */} - + {props.children} diff --git a/packages/jupyter-ai/src/contexts/active-cell-context.tsx b/packages/jupyter-ai/src/contexts/active-cell-context.tsx index 72e93a8ca..a8a6ebcaf 100644 --- a/packages/jupyter-ai/src/contexts/active-cell-context.tsx +++ b/packages/jupyter-ai/src/contexts/active-cell-context.tsx @@ -83,7 +83,7 @@ export class ActiveCellManager { * `ActiveCellContentWithError` object that describes both the active cell and * the error output. */ - getContent(withError: false): CellContent | null; + getContent(withError: false): CellContent; getContent(withError: true): CellWithErrorContent | null; getContent(withError = false): CellContent | CellWithErrorContent | null { const sharedModel = this._activeCell?.model.sharedModel; diff --git a/packages/jupyter-ai/src/handler.ts b/packages/jupyter-ai/src/handler.ts index b93a4571a..b69c6ecef 100644 --- a/packages/jupyter-ai/src/handler.ts +++ b/packages/jupyter-ai/src/handler.ts @@ -57,13 +57,26 @@ export namespace AiService { traceback: string[]; }; + export type TextSelection = { + type: 'text'; + source: string; + }; + + export type CellSelection = { + type: 'cell'; + source: string; + }; + export type CellWithErrorSelection = { type: 'cell-with-error'; source: string; error: CellError; }; - export type Selection = CellWithErrorSelection; + export type Selection = + | TextSelection + | CellSelection + | CellWithErrorSelection; export type ChatRequest = { prompt: string; @@ -101,9 +114,20 @@ export namespace AiService { type: 'human'; id: string; time: number; + /** + * The formatted body of the message to be rendered in the UI. Includes both + * `prompt` and `selection`. + */ body: string; - client: ChatClient; + /** + * The prompt typed into the chat input by the user. + */ + prompt: string; + /** + * The selection included with the prompt, if any. + */ selection?: Selection; + client: ChatClient; }; export type ConnectionMessage = { diff --git a/packages/jupyter-ai/src/icons.ts b/packages/jupyter-ai/src/icons.ts index 12d194260..29aee13ea 100644 --- a/packages/jupyter-ai/src/icons.ts +++ b/packages/jupyter-ai/src/icons.ts @@ -5,6 +5,7 @@ import { LabIcon } from '@jupyterlab/ui-components'; import chatSvgStr from '../style/icons/chat.svg'; import jupyternautSvg from '../style/icons/jupyternaut.svg'; import replaceCellSvg from '../style/icons/replace-cell.svg'; +import includeSelectionSvg from '../style/icons/include-selection.svg'; export const chatIcon = new LabIcon({ name: 'jupyter-ai::chat', @@ -21,6 +22,11 @@ export const replaceCellIcon = new LabIcon({ svgstr: replaceCellSvg }); +export const includeSelectionIcon = new LabIcon({ + name: 'jupyter-ai::include-selection', + svgstr: includeSelectionSvg +}); + // this icon is only used in the status bar. // to configure the icon shown on agent replies in the chat UI, please specify a // custom `Persona`. diff --git a/packages/jupyter-ai/style/icons/include-selection.svg b/packages/jupyter-ai/style/icons/include-selection.svg new file mode 100644 index 000000000..665ece235 --- /dev/null +++ b/packages/jupyter-ai/style/icons/include-selection.svg @@ -0,0 +1,5 @@ + + +