Skip to content

Include cell line ranges in notebook summary tool (for faster/improved reading of notebooks) #637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DonJayamanne merged 2 commits into from
Aug 18, 2025
Merged

Include cell line ranges in notebook summary tool (for faster/improved reading of notebooks) #637

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a7f84d9
Include cell line ranges in notebook summary tool
DonJayamanne Aug 18, 2025
6bd2d4e
Updates
DonJayamanne Aug 18, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions package.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -796,7 +796,7 @@
"name": "copilot_editNotebook",
"toolReferenceName": "editNotebook",
"displayName": "%copilot.tools.editNotebook.name%",
"modelDescription": "This is a tool for editing an existing Notebook file in the workspace. Generate the \"explanation\" property first.\nThe system is very smart and can understand how to apply your edits to the notebooks.\nWhen updating the content of an existing cell, ensure newCode preserves whitespace and indentation exactly and does NOT include any code markers such as (...existing code...).",
"modelDescription": "This is a tool for editing an existing Notebook file in the workspace. Generate the \"explanation\" property first.\nThe system is very smart and can understand how to apply your edits to the notebooks.\nWhen updating the content of an existing cell, ensure newCode preserves whitespace and indentation exactly and does NOT include any code markers such as (...existing code...). Do NOT use this to read the notebook file or the contents of cell(s). To read contents of a specific cell use the copilot_getNotebookSummary followed by the read_file tool.",
"tags": [
"enable_other_tool_copilot_getNotebookSummary"
],
Expand Down Expand Up @@ -886,7 +886,7 @@
"name": "copilot_getNotebookSummary",
"toolReferenceName": "getNotebookSummary",
"displayName": "Get the structure of a notebook",
"modelDescription": "This is a tool returns the list of the Notebook cells along with the id, cell types, language, execution information and output mime types for each cell. This is useful to get Cell Ids when executing a notebook or determine what cells have been executed and what order, or what cells have outputs. Requery this tool if the contents of the notebook change.",
"modelDescription": "This is a tool returns the list of the Notebook cells along with the id, cell types, line ranges, language, execution information and output mime types for each cell. This is useful to get Cell Ids when executing a notebook or determine what cells have been executed and what order, or what cells have outputs. If required to read contents of a cell use this to determine the line range of a cells, and then use read_file tool to read a specific line range. Requery this tool if the contents of the notebook change.",
"tags": [],
"inputSchema": {
"type": "object",
Expand Down
2 changes: 1 addition & 1 deletion src/extension/prompts/node/agent/summarizedConversationHistory.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,7 @@ class WorkingNotebookSummary extends PromptElement<NotebookSummaryProps> {
return (
<UserMessage>
This is the current state of the notebook that you have been working on:<br />
<NotebookSummary notebook={this.props.notebook} />
<NotebookSummary notebook={this.props.notebook} includeCellLines={false} altDoc={undefined} />
</UserMessage>
);
}
Expand Down
27 changes: 22 additions & 5 deletions src/extension/tools/node/notebookSummaryTool.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,11 +13,13 @@ import { IPromptPathRepresentationService } from '../../../platform/prompts/comm
import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService';
import { findNotebook } from '../../../util/common/notebooks';
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
import { LanguageModelPromptTsxPart, LanguageModelToolResult, MarkdownString, NotebookCellKind } from '../../../vscodeTypes';
import { LanguageModelPromptTsxPart, LanguageModelToolResult, MarkdownString, NotebookCellKind, Position } from '../../../vscodeTypes';
import { IBuildPromptContext } from '../../prompt/common/intents';
import { renderPromptElementJSON } from '../../prompts/node/base/promptRenderer';
import { NotebookVariables } from '../../prompts/node/panel/notebookVariables';
import { ToolName } from '../common/toolNames';
import { ICopilotTool, ToolRegistry } from '../common/toolsRegistry';
import { renderPromptElementJSON } from '../../prompts/node/base/promptRenderer';
import { AlternativeNotebookDocument } from '../../../platform/notebook/common/alternativeNotebookDocument';


export interface INotebookSummaryToolParams {
Expand All @@ -26,6 +28,7 @@ export interface INotebookSummaryToolParams {

export class NotebookSummaryTool implements ICopilotTool<INotebookSummaryToolParams> {
public static toolName = ToolName.GetNotebookSummary;
private promptContext?: IBuildPromptContext;

constructor(
@IPromptPathRepresentationService protected readonly promptPathRepresentationService: IPromptPathRepresentationService,
Expand All @@ -52,13 +55,14 @@ export class NotebookSummaryTool implements ICopilotTool<INotebookSummaryToolPar

this.notebookStructureTracker.trackNotebook(notebook);
this.notebookStructureTracker.clearState(notebook);

const format = this.alternativeNotebookContent.getFormat(this.promptContext?.request?.model);
const altDoc = this.alternativeNotebookContent.create(format).getAlternativeDocument(notebook);
return new LanguageModelToolResult([
new LanguageModelPromptTsxPart(
await renderPromptElementJSON(
this.instantiationService,
NotebookSummary,
{ notebook },
{ notebook, altDoc, includeCellLines: true },
// If we are not called with tokenization options, have _some_ fake tokenizer
// otherwise we end up returning the entire document
options.tokenizationOptions ?? {
Expand All @@ -71,6 +75,11 @@ export class NotebookSummaryTool implements ICopilotTool<INotebookSummaryToolPar
]);
}

async resolveInput(input: INotebookSummaryToolParams, promptContext: IBuildPromptContext): Promise<INotebookSummaryToolParams> {
this.promptContext = promptContext;
return input;
}

prepareInvocation(options: vscode.LanguageModelToolInvocationPrepareOptions<INotebookSummaryToolParams>, token: vscode.CancellationToken): vscode.ProviderResult<vscode.PreparedToolInvocation> {
return {
invocationMessage: new MarkdownString(l10n.t`Retrieving Notebook summary.`)
Expand All @@ -84,6 +93,8 @@ ToolRegistry.registerTool(NotebookSummaryTool);

type NotebookStatePromptProps = PromptElementProps<{
notebook: vscode.NotebookDocument;
altDoc: AlternativeNotebookDocument | undefined;
includeCellLines: boolean;
}>;

export class NotebookSummary extends PromptElement<NotebookStatePromptProps> {
Expand All @@ -107,7 +118,8 @@ export class NotebookSummary extends PromptElement<NotebookStatePromptProps> {

private getSummary() {
const hasAnyCellBeenExecuted = this.props.notebook.getCells().some(cell => cell.executionSummary?.executionOrder !== undefined && cell.executionSummary?.timing);

const altDoc = this.props.altDoc;
const includeCellLines = this.props.includeCellLines && !!altDoc;
return (
<>
Below is a summary of the notebook {this.promptPathRepresentationService.getFilePath(this.props.notebook.uri)}:<br />
Expand All @@ -120,6 +132,10 @@ export class NotebookSummary extends PromptElement<NotebookStatePromptProps> {
const executionOrder = cell.executionSummary?.executionOrder;
const cellId = getCellId(cell);
let executionSummary = '';

const altCellStartLine = includeCellLines ? altDoc.fromCellPosition(cell, new Position(0, 0)).line + 1 : -1;
const altCellEndLine = includeCellLines ? altDoc.fromCellPosition(cell, new Position(cell.document.lineCount - 1, 0)).line + 1 : -1;
const cellLines = `From ${altCellStartLine} to ${altCellEndLine}`;
// If there's no timing, then means the notebook wasn't executed in current session.
// Timing information is generally not stored in notebooks.
if (executionOrder === undefined || !cell.executionSummary?.timing) {
Expand All @@ -138,6 +154,7 @@ export class NotebookSummary extends PromptElement<NotebookStatePromptProps> {
return (
<>{cellNumber}. Cell Id = {cellId}<br />
{indent}Cell Type = {cellType}{language}<br />
{includeCellLines && <>{indent}Cell Lines = {cellLines}<br /></>}
{indent}{executionSummary}<br />
{outputs}
</>
Expand Down