Merge two extensions and renderer.

package.json

@@ -96,6 +96,7 @@
     },
     "jupyterlab": {
         "extension": true,
+        "mimeExtension": "./lib/renderer.js",
         "outputDir": "jupyterlab_apod/labextension"
     },
     "eslintIgnore": [

src/chapyter_extension.ts

@@ -0,0 +1,336 @@
+import {
+    JupyterFrontEnd,
+    JupyterFrontEndPlugin
+  } from '@jupyterlab/application';
+  import {
+    INotebookTracker,
+    NotebookActions,
+    Notebook
+  } from '@jupyterlab/notebook';
+  import { CodeCell, Cell, isCodeCellModel } from '@jupyterlab/cells';
+
+  const CHAPYTER_CHAT_CELL = 'jp-chapyter-chat';
+  const CHAPYTER_CHAT_CELL_EXECUTING = 'jp-chapyter-chat-executing';
+  const CHAPYTER_ASSISTANCE_CELL = 'jp-chapyter-assistance';
+
+  type ChapyterCellMetadata = {
+    linkedCellId?: string;
+    cellType: 'generated' | 'original';
+  };
+
+  /**
+   * Check if the cell is not generated by Chapyter
+   */
+  function isCellNotGenerated(cell: Cell): boolean {
+    if (isCodeCellModel(cell.model)) {
+      let metadata =
+        (cell.model.getMetadata('ChapyterCell') as ChapyterCellMetadata) || null;
+      if (metadata && metadata.cellType === 'generated') {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Iterate through the notebook and find the cell with the given ID
+   */
+  function findCellById(notebook: Notebook, id: string): Cell | null {
+    for (let i = 0; i < notebook.widgets.length; i++) {
+      let cell = notebook.widgets[i];
+      if (cell.model.id === id) {
+        return cell;
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Iterate through the notebook and find the code cell that starts with the
+   * given template string. In our case, the template string is simply a manual
+   * template that's inserted by Chapyter.
+   */
+  function findCellByTemplateString(
+    notebook: Notebook,
+    executionId: string | number | null
+  ): CodeCell | null {
+    if (executionId) {
+      const searchTempalte = `# Assistant Code for Cell [${executionId}]:`;
+
+      for (let i = 0; i < notebook.widgets.length; i++) {
+        let cell = notebook.widgets[i];
+        if (cell.model.type === 'code') {
+          let codeCell = cell as CodeCell;
+          let codeCellText = codeCell.model.sharedModel.getSource();
+          if (codeCellText.split('\n')[0].startsWith(searchTempalte)) {
+            return cell as CodeCell;
+          }
+        }
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Find the index of the cell with the given ID
+   */
+  function findCellIndexById(notebook: Notebook, id: string): number {
+    for (let i = 0; i < notebook.widgets.length; i++) {
+      let cell = notebook.widgets[i];
+      if (cell.model.id === id) {
+        return i;
+      }
+    }
+    return -1;
+  }
+
+  /**
+   * Select the target cell based on its id by moving the cursor
+   * (using NotebookActions.selectAbove or NotebookActions.selectAbove).
+   */
+  function selectCellById(notebook: Notebook, id: string): void {
+    let activeCellIndex = notebook.activeCellIndex;
+    let targetCellIndex = findCellIndexById(notebook, id);
+    if (targetCellIndex !== -1) {
+      if (activeCellIndex !== targetCellIndex) {
+        if (activeCellIndex < targetCellIndex) {
+          for (let i = activeCellIndex; i < targetCellIndex; i++) {
+            if (!notebook.widgets[i].inputHidden) {
+              NotebookActions.selectBelow(notebook);
+            }
+          }
+        } else {
+          for (let i = activeCellIndex; i > targetCellIndex; i--) {
+            if (!notebook.widgets[i].inputHidden) {
+              NotebookActions.selectAbove(notebook);
+            }
+          }
+        }
+      }
+    }
+  }
+
+  /**
+   * Check if the code cell is a Chapyter magic cell
+   * i.e., the cell starts with %chat or %%chat
+   */
+  function isCellChapyterMagicCell(
+    cell: CodeCell,
+    strict: boolean = false
+  ): boolean {
+    let codeCellText = cell.model.sharedModel.getSource();
+    if (codeCellText.startsWith('%chat') || codeCellText.startsWith('%%chat')) {
+      if (!codeCellText.startsWith('%%chatonly') || !strict) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Check if a cell is a Chapyter magic cell in safe mode
+   * indicated by the -s or --safe flag
+   */
+  function isCellChapyterMagicCellSafeMode(cell: CodeCell): boolean {
+    let codeCellText = cell.model.sharedModel.getSource();
+    let firstLine = codeCellText.split('\n')[0];
+    return firstLine.includes('-s') || firstLine.includes('--safe');
+  }
+
+  /**
+   * Delete the cell from the notebook
+   */
+  function deleteCell(notebook: Notebook, cell: Cell): void {
+    const model = notebook.model!;
+    const sharedModel = model.sharedModel;
+    const toDelete: number[] = [];
+
+    notebook.widgets.forEach((child, index) => {
+      if (child === cell) {
+        const deletable = child.model.getMetadata('deletable') !== false;
+
+        if (deletable) {
+          toDelete.push(index);
+          notebook.model?.deletedCells.push(child.model.id);
+        }
+      }
+    });
+
+    if (toDelete.length > 0) {
+      // Delete the cells as one undo event.
+      sharedModel.transact(() => {
+        // Delete cells in reverse order to maintain the correct indices.
+        toDelete.reverse().forEach(index => {
+          sharedModel.deleteCell(index);
+        });
+      });
+      // Select the *first* interior cell not deleted or the cell
+      // *after* the last selected cell.
+      // Note: The activeCellIndex is clamped to the available cells,
+      // so if the last cell is deleted the previous cell will be activated.
+      // The *first* index is the index of the last cell in the initial
+      // toDelete list due to the `reverse` operation above.
+      notebook.activeCellIndex = toDelete[0] - toDelete.length + 1;
+    }
+
+    // Deselect any remaining, undeletable cells. Do this even if we don't
+    // delete anything so that users are aware *something* happened.
+    notebook.deselectAll();
+  }
+
+  /**
+   * Initialization data for the @shannon-shen/chapyter extension.
+   */
+  const plugin: JupyterFrontEndPlugin<void> = {
+    id: '@shannon-shen/chapyter:plugin',
+    description: 'A Natural Language-Based Python Program Interpreter',
+    autoStart: true,
+    requires: [INotebookTracker],
+    // optional: [ISettingRegistry],
+    activate: (app: JupyterFrontEnd, tracker: INotebookTracker) => {
+      NotebookActions.executed.connect((sender, args) => {
+        if (args.success) {
+          // It must be true that the cell is a code cell (otherwise it would not have been executed)
+          let chatCell = args.cell as CodeCell;
+
+          // We only want to automatically generate a new cell if the code cell starts with a magic command (e.g. %chat)
+          if (
+            isCellChapyterMagicCell(chatCell, true) &&
+            isCellNotGenerated(chatCell)
+          ) {
+            // this is the original code cell that was executed
+            if (chatCell.model.getMetadata('ChapyterCell') === undefined) {
+              chatCell.model.setMetadata('ChapyterCell', {
+                cellType: 'original'
+              });
+            }
+            let inSafeMode = isCellChapyterMagicCellSafeMode(chatCell);
+
+            // because it is successfully executed
+            let notebook = tracker.currentWidget;
+            if (notebook) {
+              let assistanceCell = findCellByTemplateString(
+                notebook.content,
+                chatCell.model.executionCount
+              );
+
+              if (assistanceCell) {
+                assistanceCell.model.setMetadata('ChapyterCell', {
+                  cellType: 'generated',
+                  linkedCellId: chatCell.model.id // the original cell ID
+                });
+
+                console.log(inSafeMode)
+                if (!inSafeMode) {
+                  selectCellById(notebook.content, assistanceCell.model.id);
+                  NotebookActions.run(notebook.content, notebook.sessionContext);
+                  assistanceCell.inputHidden = true;
+                }
+
+                // The removal of existing linked cells is handled in the executionScheduled event
+
+                /**
+                 * We want to run the next check for avoiding duplicate cells.
+                 * Imagine when we are redistributing the notebook: we have already run the
+                 * chapter cell with the magic command, and the jupyter notebook generates
+                 * a new cell below the executed cell. Then another person opens the notebook
+                 * and executes the same chapyter cell. We want to delete the original generated
+                 * cell and only keep the newly generated cell.
+                 *
+                 * The logic is important: if it's on the same machine, then the caching mechanism
+                 * in guidance will produce us the same code and the user won't feel any difference.
+                 * However if it's on a different machine, then the generated code will become
+                 * different and the user will see a different result.
+                 *
+                 * We also need to execute this check after the previous cell is executed. Consider
+                 * the corner case when the (previous) generated cell is the last cell inside a juptyer
+                 * notebook. If we execute the check before the previous cell is executed, then jupyter
+                 * will move up (instead of moving down) the active cell and it will confuse the logic
+                 * for executing the next cell.
+                 */
+
+                selectCellById(notebook.content, assistanceCell.model.id);
+                if (!inSafeMode) {
+                  NotebookActions.selectBelow(notebook.content);
+                }
+
+                // set the proper linked cell ID
+                chatCell.model.setMetadata('ChapyterCell', {
+                  cellType: 'original',
+                  linkedCellId: assistanceCell.model.id
+                });
+
+                chatCell.addClass(CHAPYTER_CHAT_CELL);
+                chatCell.removeClass(CHAPYTER_CHAT_CELL_EXECUTING);
+                assistanceCell.addClass(CHAPYTER_ASSISTANCE_CELL);
+              }
+            }
+          }
+        }
+      });
+
+      NotebookActions.executionScheduled.connect((sender, args) => {
+        // It must be true that the cell is a code cell (otherwise it would not have been executed)
+        let chatCell = args.cell as CodeCell;
+
+        // We want to automatically remove existing generated cells if we are running the chapyter cell
+        if (isCellChapyterMagicCell(chatCell) && isCellNotGenerated(chatCell)) {
+          chatCell.toggleClass(CHAPYTER_CHAT_CELL_EXECUTING);
+          let linkedCellId =
+            chatCell.model.getMetadata('ChapyterCell')?.linkedCellId;
+
+          let notebook = tracker.currentWidget;
+          if (notebook) {
+            if (linkedCellId) {
+              let linkedCell = findCellById(notebook.content, linkedCellId);
+              if (linkedCell) {
+                deleteCell(notebook.content, linkedCell);
+
+                /**
+                 * Make sure we select the right cell after the deletion:
+                 * Because we will use the selectBelow function when executing the generated
+                 * code cell, we want to make sure we are selecting the current codeCell in this
+                 * executionScheduled event.
+                 */
+                selectCellById(notebook.content, chatCell.model.id);
+              }
+            }
+          }
+        }
+      });
+
+      tracker.widgetAdded.connect((sender, notebookPanel) => {
+        notebookPanel.context.ready.then(() => {
+          notebookPanel.content.widgets.forEach(cell => {
+            switch (cell.model.type) {
+              case 'code': {
+                /**
+                 * The logic: 
+                 * When we load a notebook, we want to check if a code cell is a chapyter cell.
+                 * 1. if it is generated, then we want to add the class CHAPYTER_ASSISTANCE_CELL
+                 * 2. if it is original, 
+                 *  a. if the linked cell exists, then we want to add the class CHAPYTER_CHAT_CELL
+                 *  b. if the linked cell does not exist, then we want to add the class CHAPYTER_CHAT_CELL_EXECUTING
+                 */
+                if (cell.model.getMetadata('ChapyterCell')) {
+                  if (cell.model.getMetadata('ChapyterCell')?.cellType === 'original') {
+                    if (findCellById(notebookPanel.content, cell.model.getMetadata('ChapyterCell')?.linkedCellId)) {
+                      cell.addClass(CHAPYTER_CHAT_CELL);
+                    } else {
+                      cell.addClass(CHAPYTER_CHAT_CELL_EXECUTING);
+                    }
+                  } else if (cell.model.getMetadata('ChapyterCell')?.cellType === 'generated') {
+                    cell.addClass(CHAPYTER_ASSISTANCE_CELL);
+                  } else {
+                    console.log(cell.model.getMetadata('ChapyterCell'));
+                  }
+                }
+              }
+            }
+          })
+        })
+      });
+    }
+  };
+
+  export default plugin;