Skip to content

Commit 0f95b20

Browse files
DonJayamanneDonJayamanne
authored
Move notebook tests and add translation of alt offset to cell position (#291)
* Move notebook tests and add translation of alt offset to cell position * Updates * Add tests * Updates with more tests
1 parent ec0adc8 commit 0f95b20

File tree

132 files changed

+177
-129
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

132 files changed

+177
-129
lines changed

‎src/extension/prompts/node/agent/agentPrompt.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -485,8 +485,8 @@ class CurrentEditorContext extends PromptElement<CurrentEditorContextProps> {
485485
const startCell = cellsInRange[0];
486486
const endCell = cellsInRange[cellsInRange.length - 1];
487487
const lastLine = endCell.document.lineAt(endCell.document.lineCount - 1);
488-
const startPosition = altDocument.fromCellPosition(startCell.index, new Position(0, 0));
489-
const endPosition = altDocument.fromCellPosition(endCell.index, new Position(endCell.document.lineCount - 1, lastLine.text.length));
488+
const startPosition = altDocument.fromCellPosition(startCell, new Position(0, 0));
489+
const endPosition = altDocument.fromCellPosition(endCell, new Position(endCell.document.lineCount - 1, lastLine.text.length));
490490
const selection = new Range(startPosition, endPosition);
491491
selectionText = selection ? ` The current selection is from line ${selection.start.line + 1} to line ${selection.end.line + 1}.` : '';
492492
}

‎src/extension/prompts/node/panel/chatVariables.tsx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -244,8 +244,8 @@ class DiagnosticVariable extends PromptElement<IDiagnosticVariableProps> {
244244
}
245245

246246
const altDocument = this.alternativeNotebookContent.create(this.alternativeNotebookContent.getFormat(this.endpoint)).getAlternativeDocument(notebook);
247-
const start = altDocument.fromCellPosition(cell.index, range.start);
248-
const end = altDocument.fromCellPosition(cell.index, range.end);
247+
const start = altDocument.fromCellPosition(cell, range.start);
248+
const end = altDocument.fromCellPosition(cell, range.end);
249249
const newRange = new Range(start, end);
250250
return [notebook.uri, newRange];
251251
}

‎src/extension/prompts/node/panel/currentEditor.tsx

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -118,9 +118,10 @@ export class CurrentEditor extends PromptElement<CurrentEditorPromptProps> {
118118
return undefined;
119119
}
120120
const altDocument = this._alternativeNotebookContentService.create(format).getAlternativeDocument(notebook);
121+
const cell = notebook.cellAt(cellIndex);
121122
const ranges = editor.visibleRanges.map(range => {
122-
const start = altDocument.fromCellPosition(cellIndex, range.start);
123-
const end = altDocument.fromCellPosition(cellIndex, range.end);
123+
const start = altDocument.fromCellPosition(cell, range.start);
124+
const end = altDocument.fromCellPosition(cell, range.end);
124125
return new Range(start, end);
125126
});
126127

@@ -164,8 +165,8 @@ export class CurrentEditor extends PromptElement<CurrentEditorPromptProps> {
164165
const ranges = notebookRanges.map(range => {
165166
const cell = editor.notebook.cellAt(range.start);
166167
const lastLine = cell.document.lineAt(cell.document.lineCount - 1);
167-
const start = altDocument.fromCellPosition(range.start, new Position(0, 0));
168-
const end = altDocument.fromCellPosition(range.start, lastLine.range.end);
168+
const start = altDocument.fromCellPosition(cell, new Position(0, 0));
169+
const end = altDocument.fromCellPosition(cell, lastLine.range.end);
169170
return new Range(start, end);
170171
});
171172

‎src/extension/prompts/node/panel/fileVariable.tsx

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -136,9 +136,8 @@ export class FileVariable extends PromptElement<FileVariableProps, unknown> {
136136
range = cellRange;
137137
}
138138
const altDocument = this.alternativeNotebookContent.create(this.alternativeNotebookContent.getFormat(this.promptEndpoint)).getAlternativeDocument(notebook);
139-
const cellIndex = notebook.getCells().indexOf(cell);
140139
//Translate the range to alternative content.
141-
range = new Range(altDocument.fromCellPosition(cellIndex, range.start), altDocument.fromCellPosition(cellIndex, range.end));
140+
range = new Range(altDocument.fromCellPosition(cell, range.start), altDocument.fromCellPosition(cell, range.end));
142141
} else {
143142
range = undefined;
144143
}

‎src/extension/tools/node/editNotebookTool.tsx

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import { ILogService } from '../../../platform/log/common/logService';
1414
import { IChatEndpoint } from '../../../platform/networking/common/networking';
1515
import { IAlternativeNotebookContentService } from '../../../platform/notebook/common/alternativeContent';
1616
import { BaseAlternativeNotebookContentProvider } from '../../../platform/notebook/common/alternativeContentProvider';
17-
import { getCellId, getDefaultLanguage } from '../../../platform/notebook/common/helpers';
17+
import { getCellId, getCellIdMap, getDefaultLanguage } from '../../../platform/notebook/common/helpers';
1818
import { IPromptPathRepresentationService } from '../../../platform/prompts/common/promptPathRepresentationService';
1919
import { ITelemetryService } from '../../../platform/telemetry/common/telemetry';
2020
import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService';
@@ -133,14 +133,14 @@ export class EditNotebookTool implements ICopilotTool<IEditNotebookToolParams> {
133133
try {
134134
// First validate all of the args begore applying any changes.
135135
this.fixInput(options.input, notebook, provider);
136-
this.validateInput(options.input, notebook, provider);
136+
this.validateInput(options.input, notebook);
137137
stream.notebookEdit(notebookUri, []);
138138
let previousCellIdUsedForInsertion = '';
139139
const explanation = options.input.explanation;
140140
const { editType, language, newCode } = options.input;
141141
const cellCode = Array.isArray(newCode) ? newCode.join(EOL) : newCode;
142142
let cellId = options.input.cellId || '';
143-
143+
const cellMap = getCellIdMap(notebook);
144144
if (editType === 'insert') {
145145
counters.insert++;
146146
// Model can send two subsequent inserts, and only first insert might contain the cellid.
@@ -168,7 +168,7 @@ export class EditNotebookTool implements ICopilotTool<IEditNotebookToolParams> {
168168
cellsCellIndex = cells.length;
169169
notebookCellIndex = cells.filter(item => item.type !== 'delete').length;
170170
} else {
171-
const cell = cellId ? provider.getCell(notebook, cellId) : undefined;
171+
const cell = cellId ? cellMap.get(cellId) : undefined;
172172
if (cellId && !cell) {
173173
throw new ErrorWithTelemetrySafeReason(`Invalid cell id: ${cellId}, notebook may have been modified, try reading the file again`, 'invalid_cell_id_insert_after');
174174
}
@@ -198,7 +198,7 @@ export class EditNotebookTool implements ICopilotTool<IEditNotebookToolParams> {
198198
stream.notebookEdit(notebookUri, NotebookEdit.insertCells(notebookCellIndex, [cell]));
199199
} else {
200200
previousCellIdUsedForInsertion = '';
201-
const cell = cellId ? provider.getCell(notebook, cellId) : undefined;
201+
const cell = cellId ? cellMap.get(cellId) : undefined;
202202
if (!cell) {
203203
throw new ErrorWithTelemetrySafeReason(`Invalid cell id: ${cellId}, notebook may have been modified, try reading the file again`, 'invalid_cell_id_empty');
204204
}
@@ -338,11 +338,12 @@ export class EditNotebookTool implements ICopilotTool<IEditNotebookToolParams> {
338338
};
339339
}
340340

341-
private validateInput(input: IEditNotebookToolParams, notebook: vscode.NotebookDocument, provider: BaseAlternativeNotebookContentProvider) {
341+
private validateInput(input: IEditNotebookToolParams, notebook: vscode.NotebookDocument) {
342342
const { editType, cellId, newCode } = input;
343343
// Possible we'll get cellId as a number such as -1 when inserting a cell at the top.
344344
const id = ((typeof (cellId as any) === 'number' ? `${cellId}` : cellId) || '').trim();
345-
const cell = (id && id !== 'top' && id !== 'bottom') ? provider.getCell(notebook, id) : undefined;
345+
const cellMap = getCellIdMap(notebook);
346+
const cell = (id && id !== 'top' && id !== 'bottom') ? cellMap.get(id) : undefined;
346347
if (id && id !== 'top' && id !== 'bottom' && !cell) {
347348
throw new ErrorWithTelemetrySafeReason(`None of the edits were applied as cell id: ${id} is invalid. Notebook may have been modified, try reading the file again`, 'invalidCellId');
348349
}

‎src/extension/tools/node/getNotebookCellOutputTool.tsx

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import { ToolName } from '../common/toolNames';
1818
import { ICopilotTool, ToolRegistry } from '../common/toolsRegistry';
1919
import { RunNotebookCellOutput } from './runNotebookCellTool';
2020
import { ITelemetryService } from '../../../platform/telemetry/common/telemetry';
21+
import { getCellIdMap } from '../../../platform/notebook/common/helpers';
2122

2223
export class GetNotebookCellOutputTool implements ICopilotTool<IGetNotebookCellOutputToolParams> {
2324
public static toolName = ToolName.ReadCellOutput;
@@ -52,9 +53,7 @@ export class GetNotebookCellOutputTool implements ICopilotTool<IGetNotebookCellO
5253
throw ex;
5354
}
5455

55-
const altDocProvider = this.alternativeNotebookContent.create(this.alternativeNotebookContent.getFormat(this._promptContext?.request?.model));
56-
57-
const cell = altDocProvider.getCell(notebook, cellId);
56+
const cell = getCellIdMap(notebook).get(cellId);
5857
if (!cell) {
5958
sendOutcomeTelemetry(this.telemetryService, this.endpointProvider, options, 'cellNotFound');
6059
throw new Error(`Cell not found, use the ${ToolName.ReadFile} file tool to get the latest content of the notebook file.`);

‎src/extension/tools/node/runNotebookCellTool.tsx

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import type * as vscode from 'vscode';
1010
import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider';
1111
import { IExtensionsService } from '../../../platform/extensions/common/extensionsService';
1212
import { IAlternativeNotebookContentService } from '../../../platform/notebook/common/alternativeContent';
13-
import { parseAndCleanStack } from '../../../platform/notebook/common/helpers';
13+
import { getCellIdMap, parseAndCleanStack } from '../../../platform/notebook/common/helpers';
1414
import { INotebookService } from '../../../platform/notebook/common/notebookService';
1515
import { IPromptPathRepresentationService } from '../../../platform/prompts/common/promptPathRepresentationService';
1616
import { ITelemetryService } from '../../../platform/telemetry/common/telemetry';
@@ -235,8 +235,7 @@ export class RunNotebookCellTool implements ICopilotTool<IRunNotebookCellToolPar
235235
throw new Error(`Notebook ${resolvedUri} not found.`);
236236
}
237237

238-
const altDocProvider = this.alternativeNotebookContent.create(this.alternativeNotebookContent.getFormat(this._promptContext?.request?.model));
239-
const cell = altDocProvider.getCell(notebook, cellId);
238+
const cell = getCellIdMap(notebook).get(cellId);
240239
if (!cell) {
241240
throw new Error(`Cell ${cellId} not found in the notebook, use the ${ToolName.ReadFile} file tool to get the latest content of the notebook file`);
242241
}

‎src/platform/inlineEdits/common/notebook.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/

‎src/platform/notebook/common/alternativeContent.ts

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,13 @@ export class AlternativeNotebookContentService implements IAlternativeNotebookCo
7777

7878
export function getAltNotebookRange(range: Range, cellUri: Uri, notebook: NotebookDocument, format: AlternativeContentFormat) {
7979
// If we have a range for cell, then translate that from notebook cell range to alternative range.
80-
const cellIndex = findCell(cellUri, notebook)?.index;
81-
if (cellIndex === undefined || cellIndex === -1) {
80+
const cell = findCell(cellUri, notebook);
81+
if (!cell) {
8282
return undefined;
8383
}
8484
const doc = getAlternativeNotebookDocumentProvider(format).getAlternativeDocument(notebook);
8585
return new Range(
86-
doc.fromCellPosition(cellIndex, range.start),
87-
doc.fromCellPosition(cellIndex, range.end),
86+
doc.fromCellPosition(cell, range.start),
87+
doc.fromCellPosition(cell, range.end),
8888
);
8989
}

‎src/platform/notebook/common/alternativeContentProvider.json.ts

Lines changed: 22 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,7 @@ export function isJsonContent(text: string): boolean {
2525
}
2626

2727
class AlternativeJsonDocument extends AlternativeNotebookDocument {
28-
override fromCellPosition(cellIndex: number, position: Position): Position {
29-
const cell = this.notebook.cellAt(cellIndex);
28+
override fromCellPosition(cell: NotebookCell, position: Position): Position {
3029
const cellId = getCellId(cell);
3130

3231
const alternativeContentText = this.getText();
@@ -47,6 +46,9 @@ class AlternativeJsonDocument extends AlternativeNotebookDocument {
4746
// -1 to exclude to trailing `"`
4847
return new Position(linePositionInAltContent, characterPositionInAltContent.length);
4948
}
49+
override toCellPosition(position: Position): { cell: NotebookCell; position: Position } | undefined {
50+
throw new Error('Method not implemented.');
51+
}
5052
}
5153

5254
export class AlternativeJsonNotebookContentProvider extends BaseAlternativeNotebookContentProvider {
@@ -62,8 +64,24 @@ export class AlternativeJsonNotebookContentProvider extends BaseAlternativeNoteb
6264
return this.parseAlternateContentImpl(notebookOrUri, inputStream, token);
6365
}
6466

65-
public override getAlternativeDocument(notebook: NotebookDocument): AlternativeNotebookDocument {
66-
const text = this.getAlternativeContent(notebook);
67+
public override getAlternativeDocument(notebook: NotebookDocument, excludeMarkdownCells?: boolean): AlternativeNotebookDocument {
68+
const cells = notebook.getCells().filter(cell => excludeMarkdownCells ? cell.kind !== NotebookCellKind.Markup : true).map(cell => {
69+
const summary = summarize(cell);
70+
const source = getCellCode(cell.document);
71+
72+
return {
73+
cell_type: summary.cell_type,
74+
id: summary.id,
75+
metadata: {
76+
language: summary.language
77+
},
78+
source,
79+
} satisfies SummaryCell;
80+
});
81+
82+
const json: Notebook = { cells };
83+
const text = JSON.stringify(json, undefined, IndentSize);
84+
6785
return new AlternativeJsonDocument(text, notebook);
6886
}
6987

@@ -214,25 +232,6 @@ export class AlternativeJsonNotebookContentProvider extends BaseAlternativeNoteb
214232
}
215233
});
216234
}
217-
218-
protected getAlternativeContent(notebook: NotebookDocument): string {
219-
const cells = notebook.getCells().map(cell => {
220-
const summary = summarize(cell);
221-
const source = getCellCode(cell.document);
222-
223-
return {
224-
cell_type: summary.cell_type,
225-
id: summary.id,
226-
metadata: {
227-
language: summary.language
228-
},
229-
source,
230-
} satisfies SummaryCell;
231-
});
232-
233-
const json: Notebook = { cells };
234-
return JSON.stringify(json, undefined, IndentSize);
235-
}
236235
}
237236

238237
function getCellCode(document: TextDocument): string[] {

0 commit comments

Comments
 (0)