From 146e6ad14e6ee2cbf6ae9dd1e6f673e0a8899f41 Mon Sep 17 00:00:00 2001 From: Nakul Date: Mon, 24 Nov 2025 19:39:44 +0530 Subject: [PATCH 1/6] Getting model from widget first than fallback to documentManager. --- src/chat-model.ts | 112 +++++++++++++++++++++++++++++++--------------- 1 file changed, 75 insertions(+), 37 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index 6cad61e..d5fd9e1 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -623,23 +623,39 @@ ${toolsList} } try { - const model = await this.input.documentManager?.services.contents.get( - attachment.value - ); - if (!model || model.type !== 'notebook') { - return null; - } + // Try reading from live notebook if open + const widget = this.input.documentManager?.findWidget(attachment.value); + let cellData: any[] | null = null; + let kernelLang = 'text'; + + const shared = widget?.context?.model?.sharedModel as any; + + if (shared?.cells) { + cellData = shared.cells.map((c: any) => c.toJSON()); + + kernelLang = + shared?.metadata?.language_info?.name || + shared?.metadata?.kernelspec?.language || + 'text'; + } else { + // Fallback: reading from disk + const model = await this.input.documentManager?.services.contents.get( + attachment.value + ); + if (!model || model.type !== 'notebook') { + return null; + } + cellData = model.content.cells; - const kernelLang = - model.content?.metadata?.language_info?.name || - model.content?.metadata?.kernelspec?.language || - 'text'; + kernelLang = + model.content?.metadata?.language_info?.name || + model.content?.metadata?.kernelspec?.language || + 'text'; + } const selectedCells = attachment.cells .map(cellInfo => { - const cell = model.content.cells.find( - (c: any) => c.id === cellInfo.id - ); + const cell = cellData!.find(c => c.id === cellInfo.id); if (!cell) { return null; } @@ -777,36 +793,58 @@ ${toolsList} } try { - const model = await this.input.documentManager?.services.contents.get( + // Try reading from an open widget first + const widget = this.input.documentManager?.findWidget(attachment.value); + + if (widget && widget.context && widget.context.model) { + const model = widget.context.model; + + if ((model as any).sharedModel?.getSource) { + return (model as any).sharedModel.getSource(); + } + + if ((model as any).sharedModel?.getCells) { + const sharedModel = (model as any).sharedModel; + const cells = sharedModel.getCells().map((cell: any) => ({ + ...cell, + outputs: [], + execution_count: null + })); + + return JSON.stringify({ + cells, + metadata: sharedModel.metadata || {}, + nbformat: 4, + nbformat_minor: 5 + }); + } + } + + // If not open, load from disk + const diskModel = await this.input.documentManager?.services.contents.get( attachment.value ); - if (!model?.content) { + + if (!diskModel?.content) { return null; } - if (model.type === 'file') { + + if (diskModel.type === 'file') { // Regular file content - return model.content; - } else if (model.type === 'notebook') { - // Clear outputs from notebook cells before sending to LLM - // TODO: make this configurable? - const cells = model.content.cells.map((cell: any) => { - const cleanCell = { ...cell }; - if (cleanCell.outputs) { - cleanCell.outputs = []; - } - if (cleanCell.execution_count) { - cleanCell.execution_count = null; - } - return cleanCell; - }); - - const notebookModel = { - cells, - metadata: (model as any).metadata || {}, - nbformat: (model as any).nbformat || 4, - nbformat_minor: (model as any).nbformat_minor || 4 + return diskModel.content; + } + + if (diskModel.type === 'notebook') { + const cleaned = { + ...diskModel, + cells: diskModel.content.cells.map((cell: any) => ({ + ...cell, + outputs: [], + execution_count: null + })) }; - return JSON.stringify(notebookModel); + + return JSON.stringify(cleaned); } return null; } catch (error) { From 552143442521e3f2417feaa0904a8f0bf4d95a91 Mon Sep 17 00:00:00 2001 From: Nakul Date: Sat, 29 Nov 2025 16:19:12 +0530 Subject: [PATCH 2/6] prevent llm from recieving object --- src/chat-model.ts | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index d5fd9e1..da90e10 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -800,7 +800,10 @@ ${toolsList} const model = widget.context.model; if ((model as any).sharedModel?.getSource) { - return (model as any).sharedModel.getSource(); + const source = (model as any).sharedModel.getSource(); + return typeof source === 'string' + ? source + : JSON.stringify(source, null, 2); } if ((model as any).sharedModel?.getCells) { @@ -811,12 +814,16 @@ ${toolsList} execution_count: null })); - return JSON.stringify({ - cells, - metadata: sharedModel.metadata || {}, - nbformat: 4, - nbformat_minor: 5 - }); + return JSON.stringify( + { + cells, + metadata: sharedModel.metadata || {}, + nbformat: 4, + nbformat_minor: 5 + }, + null, + 2 + ); } } From 0c28684f32b0161e27e69d28f888198fdb96cb35 Mon Sep 17 00:00:00 2001 From: Nakul Date: Sat, 29 Nov 2025 16:44:43 +0530 Subject: [PATCH 3/6] using nbformat types --- src/chat-model.ts | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index da90e10..157bdf0 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -808,16 +808,17 @@ ${toolsList} if ((model as any).sharedModel?.getCells) { const sharedModel = (model as any).sharedModel; - const cells = sharedModel.getCells().map((cell: any) => ({ + const cells = sharedModel.getCells().map((cell: nbformat.ICell) => ({ ...cell, - outputs: [], + outputs: [] as nbformat.IOutput[], execution_count: null })); return JSON.stringify( { cells, - metadata: sharedModel.metadata || {}, + metadata: + sharedModel.metadata || ({} as nbformat.INotebookMetadata), nbformat: 4, nbformat_minor: 5 }, @@ -844,9 +845,9 @@ ${toolsList} if (diskModel.type === 'notebook') { const cleaned = { ...diskModel, - cells: diskModel.content.cells.map((cell: any) => ({ + cells: diskModel.content.cells.map((cell: nbformat.ICell) => ({ ...cell, - outputs: [], + outputs: [] as nbformat.IOutput[], execution_count: null })) }; From 8a4b7042c90179d51b5b9141607d034eb0007e7c Mon Sep 17 00:00:00 2001 From: Nakul Date: Mon, 1 Dec 2025 20:29:54 +0530 Subject: [PATCH 4/6] removing any types --- src/chat-model.ts | 51 ++++++++++++++++++++++++++--------------------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index 157bdf0..56def44 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -12,6 +12,10 @@ import { PathExt } from '@jupyterlab/coreutils'; import { IDocumentManager } from '@jupyterlab/docmanager'; +import { IDocumentWidget } from '@jupyterlab/docregistry'; + +import { INotebookModel, Notebook } from '@jupyterlab/notebook'; + import { UUID } from '@lumino/coreutils'; import { ISignal, Signal } from '@lumino/signaling'; @@ -24,6 +28,8 @@ import { AISettingsModel } from './models/settings-model'; import { ITokenUsage } from './tokens'; +import { YNotebook } from '@jupyter/ydoc'; + import * as nbformat from '@jupyterlab/nbformat'; /** @@ -624,7 +630,9 @@ ${toolsList} try { // Try reading from live notebook if open - const widget = this.input.documentManager?.findWidget(attachment.value); + const widget = this.input.documentManager?.findWidget( + attachment.value + ) as IDocumentWidget | undefined; let cellData: any[] | null = null; let kernelLang = 'text'; @@ -794,37 +802,34 @@ ${toolsList} try { // Try reading from an open widget first - const widget = this.input.documentManager?.findWidget(attachment.value); + const widget = this.input.documentManager?.findWidget( + attachment.value + ) as IDocumentWidget | undefined; if (widget && widget.context && widget.context.model) { const model = widget.context.model; + const ymodel = model.sharedModel as YNotebook; - if ((model as any).sharedModel?.getSource) { - const source = (model as any).sharedModel.getSource(); + if (typeof ymodel.getSource === 'function') { + const source = ymodel.getSource(); return typeof source === 'string' ? source : JSON.stringify(source, null, 2); } - if ((model as any).sharedModel?.getCells) { - const sharedModel = (model as any).sharedModel; - const cells = sharedModel.getCells().map((cell: nbformat.ICell) => ({ - ...cell, - outputs: [] as nbformat.IOutput[], - execution_count: null - })); - - return JSON.stringify( - { - cells, - metadata: - sharedModel.metadata || ({} as nbformat.INotebookMetadata), - nbformat: 4, - nbformat_minor: 5 - }, - null, - 2 - ); + if (typeof ymodel.toJSON === 'function') { + const nb = ymodel.toJSON(); + + const cleaned = { + ...nb, + cells: nb.cells.map(cell => ({ + ...cell, + outputs: [], + execution_count: null + })) + }; + + return JSON.stringify(cleaned, null, 2); } } From b1124d0cc9bc279788a9f451b72bbea1e3f087d3 Mon Sep 17 00:00:00 2001 From: Nakul Date: Mon, 1 Dec 2025 21:56:39 +0530 Subject: [PATCH 5/6] using nbformat types instead of any --- src/chat-model.ts | 31 ++++++++++++++++++------------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index 56def44..1d0a3ad 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -633,18 +633,22 @@ ${toolsList} const widget = this.input.documentManager?.findWidget( attachment.value ) as IDocumentWidget | undefined; - let cellData: any[] | null = null; + let cellData: nbformat.ICell[] | null = null; let kernelLang = 'text'; - const shared = widget?.context?.model?.sharedModel as any; + const ymodel = widget?.context?.model?.sharedModel as YNotebook; - if (shared?.cells) { - cellData = shared.cells.map((c: any) => c.toJSON()); + if (ymodel) { + const nb = ymodel.toJSON(); - kernelLang = - shared?.metadata?.language_info?.name || - shared?.metadata?.kernelspec?.language || + cellData = nb.cells; + + const lang = + nb.metadata?.language_info?.name || + nb.metadata?.kernelspec?.language || 'text'; + + kernelLang = String(lang); } else { // Fallback: reading from disk const model = await this.input.documentManager?.services.contents.get( @@ -684,7 +688,7 @@ ${toolsList} 'text/plain' ]; - function extractDisplay(data: any): string { + function extractDisplay(data: nbformat.IMimeBundle): string { for (const mime of DISPLAY_PRIORITY) { if (!(mime in data)) { continue; @@ -697,13 +701,13 @@ ${toolsList} switch (mime) { case 'application/vnd.jupyter.widget-view+json': - return `Widget: ${(value as any).model_id ?? 'unknown model'}`; + return `Widget: ${(value as { model_id?: string }).model_id ?? 'unknown model'}`; case 'image/png': - return `![image](data:image/png;base64,${value.slice(0, 100)}...)`; + return `![image](data:image/png;base64,${String(value).slice(0, 100)}...)`; case 'image/jpeg': - return `![image](data:image/jpeg;base64,${value.slice(0, 100)}...)`; + return `![image](data:image/jpeg;base64,${String(value).slice(0, 100)}...)`; case 'image/svg+xml': return String(value).slice(0, 500) + '...\n[svg truncated]'; @@ -736,8 +740,9 @@ ${toolsList} let outputs = ''; if (cellType === 'code' && Array.isArray(cell.outputs)) { - outputs = cell.outputs - .map((output: nbformat.IOutput) => { + const outputsArray = cell.outputs as nbformat.IOutput[]; + outputs = outputsArray + .map(output => { if (output.output_type === 'stream') { return (output as nbformat.IStream).text; } else if (output.output_type === 'error') { From a41f6d03091bf005e920ba1a39b2b2f7be541c77 Mon Sep 17 00:00:00 2001 From: Nakul Date: Tue, 2 Dec 2025 18:18:03 +0530 Subject: [PATCH 6/6] implementing suggestions and removing dead code block --- src/chat-model.ts | 31 ++++++++----------------------- 1 file changed, 8 insertions(+), 23 deletions(-) diff --git a/src/chat-model.ts b/src/chat-model.ts index 1d0a3ad..3af3127 100644 --- a/src/chat-model.ts +++ b/src/chat-model.ts @@ -633,10 +633,10 @@ ${toolsList} const widget = this.input.documentManager?.findWidget( attachment.value ) as IDocumentWidget | undefined; - let cellData: nbformat.ICell[] | null = null; + let cellData: nbformat.ICell[]; let kernelLang = 'text'; - const ymodel = widget?.context?.model?.sharedModel as YNotebook; + const ymodel = widget?.context.model.sharedModel as YNotebook; if (ymodel) { const nb = ymodel.toJSON(); @@ -644,8 +644,8 @@ ${toolsList} cellData = nb.cells; const lang = - nb.metadata?.language_info?.name || - nb.metadata?.kernelspec?.language || + nb.metadata.language_info?.name || + nb.metadata.kernelspec?.language || 'text'; kernelLang = String(lang); @@ -657,17 +657,17 @@ ${toolsList} if (!model || model.type !== 'notebook') { return null; } - cellData = model.content.cells; + cellData = model.content.cells ?? []; kernelLang = - model.content?.metadata?.language_info?.name || - model.content?.metadata?.kernelspec?.language || + model.content.metadata.language_info?.name || + model.content.metadata.kernelspec?.language || 'text'; } const selectedCells = attachment.cells .map(cellInfo => { - const cell = cellData!.find(c => c.id === cellInfo.id); + const cell = cellData.find(c => c.id === cellInfo.id); if (!cell) { return null; } @@ -821,21 +821,6 @@ ${toolsList} ? source : JSON.stringify(source, null, 2); } - - if (typeof ymodel.toJSON === 'function') { - const nb = ymodel.toJSON(); - - const cleaned = { - ...nb, - cells: nb.cells.map(cell => ({ - ...cell, - outputs: [], - execution_count: null - })) - }; - - return JSON.stringify(cleaned, null, 2); - } } // If not open, load from disk