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 69ad3efc6..7dcf09ace 100644 --- a/packages/jupyter-ai/src/components/chat-input/send-button.tsx +++ b/packages/jupyter-ai/src/components/chat-input/send-button.tsx @@ -85,7 +85,8 @@ export function SendButton(props: SendButtonProps): JSX.Element { if (activeCell.exists) { props.onSend({ type: 'cell', - source: activeCell.manager.getContent(false).source + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + source: activeCell.manager.getContent(false)!.source }); closeMenu(); return; diff --git a/packages/jupyter-ai/src/contexts/active-cell-context.tsx b/packages/jupyter-ai/src/contexts/active-cell-context.tsx index a8a6ebcaf..72e93a8ca 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; + getContent(withError: false): CellContent | null; getContent(withError: true): CellWithErrorContent | null; getContent(withError = false): CellContent | CellWithErrorContent | null { const sharedModel = this._activeCell?.model.sharedModel; diff --git a/packages/jupyter-ai/src/index.ts b/packages/jupyter-ai/src/index.ts index e42091980..d6e52b576 100644 --- a/packages/jupyter-ai/src/index.ts +++ b/packages/jupyter-ai/src/index.ts @@ -18,10 +18,11 @@ import { ChatHandler } from './chat_handler'; import { buildErrorWidget } from './widgets/chat-error'; import { completionPlugin } from './completions'; import { statusItemPlugin } from './status'; -import { IJaiCompletionProvider, IJaiMessageFooter } from './tokens'; +import { IJaiCompletionProvider, IJaiCore, IJaiMessageFooter } from './tokens'; import { IRenderMimeRegistry } from '@jupyterlab/rendermime'; import { ActiveCellManager } from './contexts/active-cell-context'; import { Signal } from '@lumino/signaling'; +import { menuPlugin } from './plugins/menu-plugin'; export type DocumentTracker = IWidgetTracker; @@ -35,9 +36,10 @@ export namespace CommandIDs { /** * Initialization data for the jupyter_ai extension. */ -const plugin: JupyterFrontEndPlugin = { +const plugin: JupyterFrontEndPlugin = { id: '@jupyter-ai/core:plugin', autoStart: true, + requires: [IRenderMimeRegistry], optional: [ IGlobalAwareness, ILayoutRestorer, @@ -45,7 +47,7 @@ const plugin: JupyterFrontEndPlugin = { IJaiCompletionProvider, IJaiMessageFooter ], - requires: [IRenderMimeRegistry], + provides: IJaiCore, activate: async ( app: JupyterFrontEnd, rmRegistry: IRenderMimeRegistry, @@ -114,7 +116,14 @@ const plugin: JupyterFrontEndPlugin = { }, label: 'Focus the jupyter-ai chat' }); + + return { + activeCellManager, + chatHandler, + chatWidget, + selectionWatcher + }; } }; -export default [plugin, statusItemPlugin, completionPlugin]; +export default [plugin, statusItemPlugin, completionPlugin, menuPlugin]; diff --git a/packages/jupyter-ai/src/plugins/menu-plugin.ts b/packages/jupyter-ai/src/plugins/menu-plugin.ts new file mode 100644 index 000000000..8994a552d --- /dev/null +++ b/packages/jupyter-ai/src/plugins/menu-plugin.ts @@ -0,0 +1,158 @@ +import { + JupyterFrontEnd, + JupyterFrontEndPlugin +} from '@jupyterlab/application'; + +import { IJaiCore } from '../tokens'; +import { AiService } from '../handler'; +import { Menu } from '@lumino/widgets'; +import { CommandRegistry } from '@lumino/commands'; + +export namespace CommandIDs { + export const explain = 'jupyter-ai:explain'; + export const fix = 'jupyter-ai:fix'; + export const optimize = 'jupyter-ai:optimize'; + export const refactor = 'jupyter-ai:refactor'; +} + +/** + * Optional plugin that adds a "Generative AI" submenu to the context menu. + * These implement UI shortcuts that explain, fix, refactor, or optimize code in + * a notebook or file. + * + * **This plugin is experimental and may be removed in a future release.** + */ +export const menuPlugin: JupyterFrontEndPlugin = { + id: '@jupyter-ai/core:menu-plugin', + autoStart: true, + requires: [IJaiCore], + activate: (app: JupyterFrontEnd, jaiCore: IJaiCore) => { + const { activeCellManager, chatHandler, chatWidget, selectionWatcher } = + jaiCore; + + function activateChatSidebar() { + app.shell.activateById(chatWidget.id); + } + + function getSelection(): AiService.Selection | null { + const textSelection = selectionWatcher.selection; + const activeCell = activeCellManager.getContent(false); + const selection: AiService.Selection | null = textSelection + ? { type: 'text', source: textSelection.text } + : activeCell + ? { type: 'cell', source: activeCell.source } + : null; + + return selection; + } + + function buildLabelFactory(baseLabel: string): () => string { + return () => { + const textSelection = selectionWatcher.selection; + const activeCell = activeCellManager.getContent(false); + + return textSelection + ? `${baseLabel} (${textSelection.numLines} lines selected)` + : activeCell + ? `${baseLabel} (1 active cell)` + : baseLabel; + }; + } + + // register commands + const menuCommands = new CommandRegistry(); + menuCommands.addCommand(CommandIDs.explain, { + execute: () => { + const selection = getSelection(); + if (!selection) { + return; + } + + activateChatSidebar(); + chatHandler.sendMessage({ + prompt: 'Explain the code below.', + selection + }); + }, + label: buildLabelFactory('Explain code'), + isEnabled: () => !!getSelection() + }); + menuCommands.addCommand(CommandIDs.fix, { + execute: () => { + const activeCellWithError = activeCellManager.getContent(true); + if (!activeCellWithError) { + return; + } + + chatHandler.sendMessage({ + prompt: '/fix', + selection: { + type: 'cell-with-error', + error: activeCellWithError.error, + source: activeCellWithError.source + } + }); + }, + label: () => { + const activeCellWithError = activeCellManager.getContent(true); + return activeCellWithError + ? 'Fix code cell (1 error cell)' + : 'Fix code cell (no error cell)'; + }, + isEnabled: () => { + const activeCellWithError = activeCellManager.getContent(true); + return !!activeCellWithError; + } + }); + menuCommands.addCommand(CommandIDs.optimize, { + execute: () => { + const selection = getSelection(); + if (!selection) { + return; + } + + activateChatSidebar(); + chatHandler.sendMessage({ + prompt: 'Optimize the code below.', + selection + }); + }, + label: buildLabelFactory('Optimize code'), + isEnabled: () => !!getSelection() + }); + menuCommands.addCommand(CommandIDs.refactor, { + execute: () => { + const selection = getSelection(); + if (!selection) { + return; + } + + activateChatSidebar(); + chatHandler.sendMessage({ + prompt: 'Refactor the code below.', + selection + }); + }, + label: buildLabelFactory('Refactor code'), + isEnabled: () => !!getSelection() + }); + + // add commands as a context menu item containing a "Generative AI" submenu + const submenu = new Menu({ + commands: menuCommands + }); + submenu.id = 'jupyter-ai:submenu'; + submenu.title.label = 'Generative AI'; + submenu.addItem({ command: CommandIDs.explain }); + submenu.addItem({ command: CommandIDs.fix }); + submenu.addItem({ command: CommandIDs.optimize }); + submenu.addItem({ command: CommandIDs.refactor }); + + app.contextMenu.addItem({ + type: 'submenu', + selector: '.jp-Editor', + rank: 1, + submenu + }); + } +}; diff --git a/packages/jupyter-ai/src/selection-watcher.ts b/packages/jupyter-ai/src/selection-watcher.ts index 8dd7df586..9cbb67f31 100644 --- a/packages/jupyter-ai/src/selection-watcher.ts +++ b/packages/jupyter-ai/src/selection-watcher.ts @@ -76,6 +76,7 @@ function getTextSelection(widget: Widget | null): Selection | null { start, end, text, + numLines: text.split('\n').length, widgetId: widget.id, ...(cellId && { cellId @@ -88,6 +89,10 @@ export type Selection = CodeEditor.ITextSelection & { * The text within the selection as a string. */ text: string; + /** + * Number of lines contained by the text selection. + */ + numLines: number; /** * The ID of the document widget in which the selection was made. */ @@ -109,6 +114,10 @@ export class SelectionWatcher { setInterval(this._poll.bind(this), 200); } + get selection(): Selection | null { + return this._selection; + } + get selectionChanged(): Signal { return this._selectionChanged; } diff --git a/packages/jupyter-ai/src/tokens.ts b/packages/jupyter-ai/src/tokens.ts index efcada10f..9caf1362d 100644 --- a/packages/jupyter-ai/src/tokens.ts +++ b/packages/jupyter-ai/src/tokens.ts @@ -1,8 +1,12 @@ import React from 'react'; import { Token } from '@lumino/coreutils'; import { ISignal } from '@lumino/signaling'; -import type { IRankedMenu } from '@jupyterlab/ui-components'; +import type { IRankedMenu, ReactWidget } from '@jupyterlab/ui-components'; + import { AiService } from './handler'; +import { ChatHandler } from './chat_handler'; +import { ActiveCellManager } from './contexts/active-cell-context'; +import { SelectionWatcher } from './selection-watcher'; export interface IJaiStatusItem { addItem(item: IRankedMenu.IItemOptions): void; @@ -46,3 +50,20 @@ export const IJaiMessageFooter = new Token( 'jupyter_ai:IJaiMessageFooter', 'Optional component that is used to render a footer on each Jupyter AI chat message, when provided.' ); + +export interface IJaiCore { + chatWidget: ReactWidget; + chatHandler: ChatHandler; + activeCellManager: ActiveCellManager; + selectionWatcher: SelectionWatcher; +} + +/** + * The Jupyter AI core provider token. Frontend plugins that want to extend the + * Jupyter AI frontend by adding features which send messages or observe the + * current text selection & active cell should require this plugin. + */ +export const IJaiCore = new Token( + 'jupyter_ai:core', + 'The core implementation of the frontend.' +);