improve context collection

Author: Yahiewi

jupyter_ai_tutor/handlers.py

@@ -10,16 +10,28 @@ class ExplainHandler(APIHandler):
     @tornado.web.authenticated
     async def post(self):
         body = self.get_json_body()
-        if not body or "body" not in body:
-            raise tornado.web.HTTPError(400, "Missing 'body' field in request")
+        if not body:
+            raise tornado.web.HTTPError(400, "Missing request body")
 
-        message_body = body["body"]
-        description = body.get("description", "")
-        if description:
-            message_body = (
-                f"<exercise_description>\n{description}\n</exercise_description>\n\n"
-                f"{message_body}"
-            )
+        student_context = body.get("student_context", "")
+        student_answer = body.get("student_answer", "")
+
+        if student_context or student_answer:
+            message_body = ""
+            if student_context:
+                message_body += f"<student_context>\n{student_context}\n</student_context>\n\n"
+            if student_answer:
+                message_body += f"<student_answer>\n{student_answer}\n</student_answer>\n\n"
+        else:
+            if "body" not in body:
+                raise tornado.web.HTTPError(400, "Missing 'body' field in request")
+            message_body = body["body"]
+            description = body.get("description", "")
+            if description:
+                message_body = (
+                    f"<exercise_description>\n{description}\n</exercise_description>\n\n"
+                    f"{message_body}"
+                )
 
         config_manager = self.settings.get("jupyternaut.config_manager")
         if not config_manager:

src/api.ts

@@ -1,16 +1,24 @@
 import { URLExt } from '@jupyterlab/coreutils';
 import { ServerConnection } from '@jupyterlab/services';
 
+/**
+ * Options for streaming the explanation request.
+ */
+export interface IStreamExplanationOptions {
+  body: string;
+  description?: string;
+  signal?: AbortSignal;
+  studentContext?: string;
+  studentAnswer?: string;
+}
+
 /**
  * Streams the tutor explanation for the given message body via SSE.
  * Yields text chunks as they arrive from the backend.
- * @param body - The user message (code + question)
- * @param description - Optional exercise description from preceding markdown cells
+ * @param options - The request options containing body and structured contexts
  */
 export async function* streamExplanation(
-  body: string,
-  description?: string,
-  signal?: AbortSignal
+  options: IStreamExplanationOptions
 ): AsyncGenerator<string, void, undefined> {
   const settings = ServerConnection.makeSettings();
   const url = URLExt.join(settings.baseUrl, 'api/jupyter-ai-tutor/explain');
@@ -19,9 +27,14 @@ export async function* streamExplanation(
     url,
     {
       method: 'POST',
-      body: JSON.stringify({ body, description }),
+      body: JSON.stringify({
+        body: options.body,
+        description: options.description,
+        student_context: options.studentContext,
+        student_answer: options.studentAnswer
+      }),
       headers: { 'Content-Type': 'application/json' },
-      signal
+      signal: options.signal
     },
     settings
   );

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,45 +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 markdownCellsForAttachment = [];
+
+          for (const cCell of contextCells) {
+            const cSource = cCell.model.sharedModel.source.trim();
+            if (!cSource) {
+              continue;
+            }
+
+            if (cCell.model.type === 'markdown') {
+              contextStr += `${cSource}\n\n`;
+              markdownCellsForAttachment.push({
+                id: cCell.model.id,
+                input_type: 'markdown' as const
+              });
+            } 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 (markdownCellsForAttachment.length > 0) {
             attachment = {
               type: 'notebook',
               value: notebookPath,
-              cells: markdownCells.map(c => ({
-                id: c.id,
-                input_type: 'markdown' as const
-              }))
+              cells: markdownCellsForAttachment
             };
           }
         }
 
+        // 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}`;
+        }
+
         const question = errorSection
           ? 'Can you explain this code and the error it produced?'
           : 'Can you explain this code?';
@@ -254,7 +287,8 @@ const plugin: JupyterFrontEndPlugin<void> = {
 
         await tutorModel.sendMessageToAI({
           body,
-          description,
+          studentContext,
+          studentAnswer,
           attachments: attachment ? [attachment] : undefined
         });
       },

src/model.ts

@@ -15,6 +15,8 @@ import { AI_AVATAR } from './icons';
 interface ITutorNewMessage extends INewMessage {
   attachments?: IAttachment[];
   description?: string;
+  studentContext?: string;
+  studentAnswer?: string;
 }
 
 export const TUTOR_USER: IUser = {
@@ -87,11 +89,13 @@ export class TutorChatModel extends AbstractChatModel {
 
     try {
       let accumulated = '';
-      for await (const chunk of streamExplanation(
-        message.body,
-        message.description,
-        this._abortController.signal
-      )) {
+      for await (const chunk of streamExplanation({
+        body: message.body,
+        description: message.description,
+        signal: this._abortController.signal,
+        studentContext: message.studentContext,
+        studentAnswer: message.studentAnswer
+      })) {
         accumulated += chunk;
         streamingMsg.update({ body: accumulated });
       }