Improve context collection

README.md

@@ -13,8 +13,6 @@ Once Jupyterlab started, the model must be configured via the menu `Settings>Jup
 
 See Jupyternaut [documentation](https://jupyter-ai.readthedocs.io/en/v3/users/jupyternaut/index.html#model-selection) to set it up.
 
-
-
 ## Requirements
 
 - JupyterLab >= 4.0.0

jupyter_ai_tutor/TUTOR.md

@@ -18,12 +18,17 @@ understand code — never to write code for them.
 - Suggest what to search for or which documentation to read.
 - Break a complex problem into smaller steps and ask the student to tackle one at a time.
 
-## Exercise Context
-
-When the student's message contains an `<exercise_description>` block, that content
-comes from the markdown cells immediately preceding the code cell in the notebook.
-Use it to understand what the student is expected to accomplish, and tailor your
-guidance to that goal. The block is not visible to the student.
+## Request formatting
+
+- The cell to explain is in a `<source>` block.
+- When the message contains an `<context>` block, that content comes from the markdown
+  cells immediately preceding the code cell in the notebook.
+  Use it to understand what the student is expected to accomplish, and tailor your
+  guidance to that goal. The block is not visible to the student.
+- A `<reference_solution>` block can also be embedded. You should use it to help guiding
+  the student, without exposing its content.
+- A `<evaluation_criteria>` block may also be embedded. You should use this to help the
+  student in accordance with the teacher's expectations.
 
 ## Tone
 

src/index.ts

@@ -184,6 +184,11 @@ const plugin: JupyterFrontEndPlugin<void> = {
         const codeModel = cell.model as ICodeCellModel;
         const outputs = codeModel.outputs;
         let errorSection = '';
+        let jsonError: {
+          ename: string;
+          evalue: string;
+          traceback: string[];
+        } | null = null;
 
         for (let i = 0; i < outputs.length; i++) {
           const output = outputs.get(i);
@@ -193,6 +198,7 @@ const plugin: JupyterFrontEndPlugin<void> = {
               evalue: string;
               traceback: string[];
             };
+            jsonError = json;
             const traceback = json.traceback
               .map(line => line.replace(ANSI_ESCAPE, ''))
               .join('\n');
@@ -203,44 +209,72 @@ const plugin: JupyterFrontEndPlugin<void> = {
           }
         }
 
-        // Collect preceding markdown cells as exercise description context.
+        // Collect preceding cells context.
         const notebook = notebookTracker?.currentWidget?.content;
         const notebookPath = notebookTracker?.currentWidget?.context.path ?? '';
-        let description: string | undefined;
+        let studentContext = '';
         let attachment: INotebookAttachment | undefined;
 
         if (notebook) {
           const activeCellIndex = notebook.activeCellIndex;
-          const markdownCells: { id: string; source: string }[] = [];
+          let lastMdIdx = -1;
 
+          // Find the index of the most recent markdown cell above the active cell
           for (let i = activeCellIndex - 1; i >= 0; i--) {
             const precedingCell = notebook.widgets[i];
-            if (precedingCell.model.type === 'code') {
+            if (precedingCell.model.type === 'markdown') {
+              lastMdIdx = i;
               break;
             }
-            if (precedingCell.model.type === 'markdown') {
-              const mdSource = precedingCell.model.sharedModel.source.trim();
-              if (mdSource) {
-                markdownCells.unshift({
-                  id: precedingCell.model.id,
-                  source: mdSource
-                });
-              }
+          }
+
+          // Gather all cells from that markdown cell up to activeCellIndex - 1
+          const startIdx = lastMdIdx !== -1 ? lastMdIdx : 0;
+
+          const contextCells = [];
+          for (let i = startIdx; i < activeCellIndex; i++) {
+            contextCells.push(notebook.widgets[i]);
+          }
+
+          let contextStr = '';
+          const cellsForAttachment = [];
+
+          for (const cCell of contextCells) {
+            const cSource = cCell.model.sharedModel.source.trim();
+            if (!cSource) {
+              continue;
+            }
+
+            cellsForAttachment.push({
+              id: cCell.model.id,
+              input_type: cCell.model.type as 'raw' | 'markdown' | 'code'
+            });
+
+            if (cCell.model.type === 'markdown') {
+              contextStr += `${cSource}\n\n`;
+            } else if (cCell.model.type === 'code') {
+              contextStr += `Preceding Code:\n\`\`\`${language}\n${cSource}\n\`\`\`\n\n`;
             }
           }
 
-          if (markdownCells.length > 0) {
-            description = markdownCells.map(c => c.source).join('\n\n');
+          studentContext = contextStr.trim();
+
+          if (cellsForAttachment.length > 0) {
             attachment = {
               type: 'notebook',
               value: notebookPath,
-              cells: markdownCells.map(c => ({
-                id: c.id,
-                input_type: 'markdown' as const
-              }))
+              cells: cellsForAttachment
             };
           }
         }
+        // Format student answer
+        let studentAnswer = source;
+        if (jsonError) {
+          const traceback = jsonError.traceback
+            .map(line => line.replace(ANSI_ESCAPE, ''))
+            .join('\n');
+          studentAnswer += `\n\nError:\n${jsonError.ename}: ${jsonError.evalue}\n${traceback}`;
+        }
 
         // Retrieve and decode reference_solution from metadata
         const rawSolution = cell.model.getMetadata('reference_solution');
@@ -252,29 +286,33 @@ const plugin: JupyterFrontEndPlugin<void> = {
         const evaluationCriteria = formatEvaluationCriteria(rawCriteria);
 
         const question = errorSection
-          ? 'Can you explain this code and the error it produced?'
-          : 'Can you explain this code?';
+          ? 'Explain code and error'
+          : 'Explain code?';
         const bodyContent = `${question}\n\n\`\`\`${language}\n${source}\n\`\`\`${errorSection}\n`;
 
         let formattedBody = '';
-        if (description) {
-          formattedBody += `<exercise_description>\n${description}\n</exercise_description>\n\n`;
+        if (studentContext) {
+          formattedBody += `<context>\n${studentContext}\n</context>\n\n`;
         }
-        formattedBody += bodyContent;
+        formattedBody += `<source>\n${studentAnswer}\n</source>`;
+
         if (referenceSolution) {
           formattedBody += `\n\n<reference_solution>\n${referenceSolution}\n</reference_solution>`;
         }
         if (evaluationCriteria) {
           formattedBody += `\n\n<evaluation_criteria>\n${evaluationCriteria}\n</evaluation_criteria>`;
         }
 
+        formattedBody += '\n';
+
         if (!chatWidget.isAttached) {
           app.shell.add(chatWidget, 'right');
         }
         app.shell.activateById(chatWidget.id);
 
         await tutorModel.sendMessageToAI({
-          body: formattedBody,
+          body: bodyContent,
+          formattedBody: formattedBody,
           notebookPath,
           attachments: attachment ? [attachment] : undefined
         });

src/model.ts

@@ -15,6 +15,7 @@ import { AI_AVATAR } from './icons';
 interface ITutorNewMessage extends INewMessage {
   attachments?: IAttachment[];
   notebookPath?: string;
+  formattedBody?: string;
 }
 
 export const TUTOR_USER: IUser = {
@@ -88,7 +89,7 @@ export class TutorChatModel extends AbstractChatModel {
     try {
       let accumulated = '';
       for await (const chunk of streamExplanation(
-        message.body,
+        message.formattedBody ?? message.body,
         message.notebookPath,
         this._abortController.signal
       )) {