Code highlight in message (#226)

* Highlight code attached from a cell

* Highlight code attached from a selection

* Add test

* Apply suggestions from PR comments

Author: brichet

packages/jupyter-chat/package.json

@@ -48,6 +48,8 @@
     "@jupyter/react-components": "^0.15.2",
     "@jupyterlab/application": "^4.2.0",
     "@jupyterlab/apputils": "^4.3.0",
+    "@jupyterlab/codeeditor": "^4.2.0",
+    "@jupyterlab/codemirror": "^4.2.0",
     "@jupyterlab/docmanager": "^4.2.0",
     "@jupyterlab/filebrowser": "^4.2.0",
     "@jupyterlab/fileeditor": "^4.2.0",

packages/jupyter-chat/src/active-cell-manager.ts

@@ -13,11 +13,13 @@ import { ISignal, Signal } from '@lumino/signaling';
 type CellContent = {
   type: string;
   source: string;
+  language?: string;
 };
 
 type CellWithErrorContent = {
   type: 'code';
   source: string;
+  language?: string;
   error: {
     name: string;
     value: string;
@@ -148,6 +150,11 @@ export class ActiveCellManager implements IActiveCellManager {
   getContent(withError: true): CellWithErrorContent | null;
   getContent(withError = false): CellContent | CellWithErrorContent | null {
     const sharedModel = this._notebookTracker.activeCell?.model.sharedModel;
+    const language =
+      sharedModel?.cell_type === 'code'
+        ? this._notebookTracker.currentWidget?.model?.defaultKernelLanguage
+        : undefined;
+
     if (!sharedModel) {
       return null;
     }
@@ -156,7 +163,8 @@ export class ActiveCellManager implements IActiveCellManager {
     if (!withError) {
       return {
         type: sharedModel.cell_type,
-        source: sharedModel.getSource()
+        source: sharedModel.getSource(),
+        language
       };
     }
 
@@ -166,6 +174,7 @@ export class ActiveCellManager implements IActiveCellManager {
       return {
         type: 'code',
         source: sharedModel.getSource(),
+        language,
         error: {
           name: error.ename,
           value: error.evalue,

packages/jupyter-chat/src/components/input/buttons/send-button.tsx

@@ -109,18 +109,22 @@ export function SendButton(
     // Run all chat command providers
     await chatCommandRegistry?.onSubmit(model);
 
+    let language: string | undefined;
     if (selectionWatcher?.selection) {
       // Append the selected text if exists.
       source = selectionWatcher.selection.text;
+      language = selectionWatcher.selection.language;
     } else if (activeCellManager?.available) {
       // Append the active cell content if exists.
-      source = activeCellManager.getContent(false)!.source;
+      const content = activeCellManager.getContent(false);
+      source = content!.source;
+      language = content?.language;
     }
     let content = model.value;
     if (source) {
       content += `
 
-\`\`\`
+\`\`\`${language ?? ''}
 ${source}
 \`\`\`
 `;

packages/jupyter-chat/src/selection-watcher.ts

@@ -6,6 +6,10 @@
 import { JupyterFrontEnd } from '@jupyterlab/application';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { CodeEditor } from '@jupyterlab/codeeditor';
+import {
+  EditorLanguageRegistry,
+  IEditorLanguageRegistry
+} from '@jupyterlab/codemirror';
 import { Notebook } from '@jupyterlab/notebook';
 
 import { find } from '@lumino/algorithm';
@@ -26,6 +30,10 @@ export namespace SelectionWatcher {
      * The current shell of the application.
      */
     shell: JupyterFrontEnd.IShell;
+    /**
+     * Editor language registry.
+     */
+    languages?: IEditorLanguageRegistry;
   }
 
   /**
@@ -44,6 +52,10 @@ export namespace SelectionWatcher {
      * The ID of the document widget in which the selection was made.
      */
     widgetId: string;
+    /**
+     * The language of the selection.
+     */
+    language?: string;
     /**
      * The ID of the cell in which the selection was made, if the original widget
      * was a notebook.
@@ -70,6 +82,7 @@ export interface ISelectionWatcher {
 export class SelectionWatcher {
   constructor(options: SelectionWatcher.IOptions) {
     this._shell = options.shell;
+    this._languages = options.languages || new EditorLanguageRegistry();
     this._shell.currentChanged?.connect((sender, args) => {
       // Do not change the main area widget if the new one has no editor, for example
       // a chat panel. However, the selected text is only available if the main area
@@ -149,12 +162,15 @@ export class SelectionWatcher {
     editor.setSelection({ start: newPosition, end: newPosition });
   }
 
-  protected _poll(): void {
+  protected async _poll(): Promise<void> {
     let currSelection: SelectionWatcher.Selection | null = null;
     const prevSelection = this._selection;
     // Do not return selected text if the main area widget is hidden.
     if (this._mainAreaDocumentWidget?.isVisible) {
-      currSelection = getTextSelection(this._mainAreaDocumentWidget);
+      currSelection = await getTextSelection(
+        this._mainAreaDocumentWidget,
+        this._languages
+      );
     }
     if (prevSelection?.text !== currSelection?.text) {
       this._selection = currSelection;
@@ -169,14 +185,16 @@ export class SelectionWatcher {
     this,
     SelectionWatcher.Selection | null
   >(this);
+  private _languages: IEditorLanguageRegistry;
 }
 
 /**
  * Gets a Selection object from a document widget. Returns `null` if unable.
  */
-function getTextSelection(
-  widget: Widget | null
-): SelectionWatcher.Selection | null {
+async function getTextSelection(
+  widget: Widget | null,
+  languages: IEditorLanguageRegistry
+): Promise<SelectionWatcher.Selection | null> {
   const editor = getEditor(widget);
   // widget type check is redundant but hints the type to TypeScript
   if (!editor || !(widget instanceof DocumentWidget)) {
@@ -207,13 +225,18 @@ function getTextSelection(
     [start, end] = [end, start];
   }
 
+  const language = (await languages.getLanguage(editor?.model.mimeType))?.name;
   return {
     ...selectionObj,
     start,
     end,
     text,
     numLines: text.split('\n').length,
     widgetId: widget.id,
+
+    ...(language && {
+      language
+    }),
     ...(cellId && {
       cellId
     })

packages/jupyterlab-chat-extension/package.json

@@ -52,6 +52,7 @@
     "@jupyter/ydoc": "^3.0.0",
     "@jupyterlab/application": "^4.4.0",
     "@jupyterlab/apputils": "^4.5.0",
+    "@jupyterlab/codemirror": "^4.4.0",
     "@jupyterlab/coreutils": "^6.4.0",
     "@jupyterlab/docregistry": "^4.4.0",
     "@jupyterlab/filebrowser": "^4.4.0",

packages/jupyterlab-chat-extension/src/index.ts

@@ -35,6 +35,7 @@ import {
   createToolbarFactory,
   showErrorMessage
 } from '@jupyterlab/apputils';
+import { IEditorLanguageRegistry } from '@jupyterlab/codemirror';
 import { PathExt } from '@jupyterlab/coreutils';
 import { DocumentRegistry } from '@jupyterlab/docregistry';
 import { IDefaultFileBrowser } from '@jupyterlab/filebrowser';
@@ -813,9 +814,14 @@ const selectionWatcher: JupyterFrontEndPlugin<ISelectionWatcher> = {
   description: 'The selection watcher plugin.',
   autoStart: true,
   provides: ISelectionWatcherToken,
-  activate: (app: JupyterFrontEnd): ISelectionWatcher => {
+  optional: [IEditorLanguageRegistry],
+  activate: (
+    app: JupyterFrontEnd,
+    languages: IEditorLanguageRegistry
+  ): ISelectionWatcher => {
     return new SelectionWatcher({
-      shell: app.shell
+      shell: app.shell,
+      languages
     });
   }
 };

ui-tests/tests/send-message.spec.ts

@@ -164,7 +164,7 @@ test.describe('#sendMessages', () => {
     );
   });
 
-  test('should send with cell content', async ({ page }) => {
+  test('should send with code cell content', async ({ page }) => {
     const cellContent = 'a = 1\nprint(f"a={a}")';
     const chatPanel = await openChat(page, FILENAME);
     const messages = chatPanel.locator('.jp-chat-messages-container');
@@ -189,14 +189,56 @@ test.describe('#sendMessages', () => {
     await expect(sendWithSelection).toBeEnabled();
     await expect(sendWithSelection).toContainText('Code from 1 active cell');
     await sendWithSelection.click();
+    await expect(messages!.locator('.jp-chat-message')).toHaveCount(1);
+
+    // It seems that the markdown renderer adds a new line, but the '\n' inserter when
+    // pressing Enter above is trimmed.
+    const rendered = messages.locator(
+      '.jp-chat-message .jp-chat-rendered-markdown'
+    );
+    await expect(rendered).toHaveText(`${MSG_CONTENT}\n${cellContent}\n`);
 
+    // Code should have python language class.
+    await expect(rendered.locator('code')).toHaveClass('language-python');
+  });
+
+  test('should send with markdown cell content', async ({ page }) => {
+    const cellContent = 'markdown content';
+    const chatPanel = await openChat(page, FILENAME);
+    const messages = chatPanel.locator('.jp-chat-messages-container');
+    const input = chatPanel
+      .locator('.jp-chat-input-container')
+      .getByRole('combobox');
+    const openerButton = chatPanel.locator(
+      '.jp-chat-input-container .jp-chat-send-include-opener'
+    );
+    const sendWithSelection = page.locator('.jp-chat-send-include');
+
+    const notebook = await page.notebook.createNew();
+    // write content in the first cell after changing it to markdown.
+    const cell = (await page.notebook.getCellLocator(0))!;
+    await page.notebook.setCellType(0, 'markdown');
+    await cell.getByRole('textbox').pressSequentially(cellContent);
+
+    await splitMainArea(page, notebook!);
+
+    await input.pressSequentially(MSG_CONTENT);
+    await openerButton.click();
+    await expect(sendWithSelection).toBeVisible();
+    await expect(sendWithSelection).toBeEnabled();
+    await expect(sendWithSelection).toContainText('Code from 1 active cell');
+    await sendWithSelection.click();
     await expect(messages!.locator('.jp-chat-message')).toHaveCount(1);
 
     // It seems that the markdown renderer adds a new line, but the '\n' inserter when
     // pressing Enter above is trimmed.
-    await expect(
-      messages.locator('.jp-chat-message .jp-chat-rendered-markdown')
-    ).toHaveText(`${MSG_CONTENT}\n${cellContent}\n`);
+    const rendered = messages.locator(
+      '.jp-chat-message .jp-chat-rendered-markdown'
+    );
+    await expect(rendered).toHaveText(`${MSG_CONTENT}\n${cellContent}\n`);
+
+    // Code should not have python language class since it come from a markdown cell.
+    await expect(rendered.locator('code')).toHaveClass('');
   });
 
   test('should send with text selection', async ({ page }) => {
@@ -243,8 +285,12 @@ test.describe('#sendMessages', () => {
 
     // It seems that the markdown renderer adds a new line, but the '\n' inserter when
     // pressing Enter above is trimmed.
-    await expect(
-      messages.locator('.jp-chat-message .jp-chat-rendered-markdown')
-    ).toHaveText(`${MSG_CONTENT}\nprint\n`);
+    const rendered = messages.locator(
+      '.jp-chat-message .jp-chat-rendered-markdown'
+    );
+    await expect(rendered).toHaveText(`${MSG_CONTENT}\nprint\n`);
+
+    // Code should have python or ipython language class.
+    await expect(rendered.locator('code')).toHaveClass(/language-[i]?python/);
   });
 });

yarn.lock

@@ -2269,6 +2269,8 @@ __metadata:
     "@jupyter/react-components": ^0.15.2
     "@jupyterlab/application": ^4.2.0
     "@jupyterlab/apputils": ^4.3.0
+    "@jupyterlab/codeeditor": ^4.2.0
+    "@jupyterlab/codemirror": ^4.2.0
     "@jupyterlab/docmanager": ^4.2.0
     "@jupyterlab/filebrowser": ^4.2.0
     "@jupyterlab/fileeditor": ^4.2.0
@@ -2534,7 +2536,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyterlab/codeeditor@npm:^4.2.3, @jupyterlab/codeeditor@npm:^4.4.3":
+"@jupyterlab/codeeditor@npm:^4.2.0, @jupyterlab/codeeditor@npm:^4.2.3, @jupyterlab/codeeditor@npm:^4.4.3":
   version: 4.4.3
   resolution: "@jupyterlab/codeeditor@npm:4.4.3"
   dependencies:
@@ -2558,7 +2560,7 @@ __metadata:
   languageName: node
   linkType: hard
 
-"@jupyterlab/codemirror@npm:^4.2.3, @jupyterlab/codemirror@npm:^4.4.3":
+"@jupyterlab/codemirror@npm:^4.2.0, @jupyterlab/codemirror@npm:^4.2.3, @jupyterlab/codemirror@npm:^4.4.0, @jupyterlab/codemirror@npm:^4.4.3":
   version: 4.4.3
   resolution: "@jupyterlab/codemirror@npm:4.4.3"
   dependencies:
@@ -9724,6 +9726,7 @@ __metadata:
     "@jupyterlab/application": ^4.4.0
     "@jupyterlab/apputils": ^4.5.0
     "@jupyterlab/builder": ^4.2.0
+    "@jupyterlab/codemirror": ^4.4.0
     "@jupyterlab/coreutils": ^6.4.0
     "@jupyterlab/docregistry": ^4.4.0
     "@jupyterlab/filebrowser": ^4.4.0