@vscode/chat-lib 0.0.1

package/src/_internal/platform/notebook/common/alternativeContentProvider.xml.ts

@@ -0,0 +1,205 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import type { CancellationToken, NotebookCell, NotebookDocument, Uri } from 'vscode';
+import { getLanguage } from '../../../util/common/languages';
+import { isUri } from '../../../util/common/types';
+import { EndOfLine, NotebookCellKind, Position } from '../../../util/common/test/shims/vscodeTypesShim';
+import { BaseAlternativeNotebookContentProvider } from './alternativeContentProvider';
+import { AlternativeNotebookDocument } from './alternativeNotebookDocument';
+import { EOL, getCellIdMap, getDefaultLanguage, LineOfCellText, LineOfText, summarize, SummaryCell } from './helpers';
+import { findLast } from '../../../util/vs/base/common/arraysFind';
+
+const StartDelimter = `<VSCode.Cell `;
+const StartEmptyCellDelimter = `<VSCode.Cell>`;
+const EndDelimter = `</VSCode.Cell>`;
+
+function generatePartialStartDelimiterWithId(id: string) {
+	return `${StartDelimter}id="${id}" `;
+}
+
+function generateCellMarker(cell: SummaryCell) {
+	return `${generatePartialStartDelimiterWithId(cell.id)}language="${cell.language}">`;
+}
+
+export function isXmlContent(text: string): boolean {
+	return text.includes(StartDelimter) || text.includes(EndDelimter) || text.includes(StartEmptyCellDelimter);
+}
+
+
+class AlternativeXmlDocument extends AlternativeNotebookDocument {
+	constructor(text: string, private readonly cellOffsetMap: { offset: number; cell: NotebookCell }[], notebook: NotebookDocument) {
+		super(text, notebook);
+	}
+
+	override fromCellPosition(cell: NotebookCell, position: Position): Position {
+		const cellSummary = summarize(cell);
+		const cellMarker = generateCellMarker(cellSummary);
+
+		const eolLength = cell.document.eol === EndOfLine.LF ? 1 : 2;
+
+		const alternativeContentText = this.getText();
+		const offsetInCell = cell.document.offsetAt(position);
+		const offset = alternativeContentText.indexOf(cellMarker) + cellMarker.length + eolLength + offsetInCell;
+		return this.positionAt(offset);
+	}
+
+	override toCellPosition(position: Position): { cell: NotebookCell; position: Position } | undefined {
+		const offset = this.offsetAt(position);
+		const cell = findLast(this.cellOffsetMap, (cell) => cell.offset <= offset);
+		if (!cell) {
+			return undefined;
+		}
+		const cellPosition = cell.cell.document.positionAt(offset - cell.offset);
+		return { cell: cell.cell, position: cellPosition };
+	}
+}
+
+export class AlternativeXmlNotebookContentProvider extends BaseAlternativeNotebookContentProvider {
+	constructor() {
+		super('xml');
+	}
+	public stripCellMarkers(text: string): string {
+		const lines = text.split(EOL);
+		if (lines.length && (lines[0].startsWith(StartDelimter) || lines[0].startsWith(StartEmptyCellDelimter))) {
+			lines.shift();
+		}
+		if (lines.length && lines[lines.length - 1].trim().endsWith(EndDelimter)) {
+			lines[lines.length - 1] = lines[lines.length - 1].substring(0, lines[lines.length - 1].lastIndexOf(EndDelimter));
+		}
+		return lines.join(EOL);
+	}
+
+	public override getSummaryOfStructure(notebook: NotebookDocument, cellsToInclude: NotebookCell[], existingCodeMarker: string): string {
+		const lines: string[] = [];
+		const existingCodeMarkerWithComment = `// ${existingCodeMarker}`;
+		notebook.getCells().forEach((cell) => {
+			if (cellsToInclude.includes(cell)) {
+				const cellSummary = summarize(cell);
+				lines.push(generateCellMarker(cellSummary));
+				if (cellSummary.source.length && cellSummary.source[0].trim().length) {
+					lines.push(cellSummary.source[0]);
+					lines.push(existingCodeMarkerWithComment);
+				} else if (cellSummary.source.length && cellSummary.source.some(line => line.trim().length)) {
+					cellSummary.source = [existingCodeMarkerWithComment, cellSummary.source.filter(line => line.trim().length)[0], existingCodeMarkerWithComment];
+				} else {
+					lines.push(existingCodeMarkerWithComment);
+				}
+				lines.push(EndDelimter);
+			} else if (!lines.length || lines[lines.length - 1] !== existingCodeMarkerWithComment) {
+				lines.push(existingCodeMarkerWithComment);
+			}
+		});
+		return lines.join(EOL);
+	}
+
+	public async *parseAlternateContent(notebookOrUri: NotebookDocument | Uri, inputStream: AsyncIterable<LineOfText>, token: CancellationToken): AsyncIterable<LineOfCellText> {
+		const isNotebook = !isUri(notebookOrUri);
+		const cellIdMap = isNotebook ? getCellIdMap(notebookOrUri) : new Map<string, NotebookCell>();
+
+
+		let index = -1;
+		let endDelimiterSeen = false;
+		const cellIdsSeen = new Set<string>();
+		let previousLineEndedWithEndCellMarker = false;
+		let previousLine: LineOfCellText | undefined = undefined;
+		const defaultLanguage = isNotebook ? getLanguage(getDefaultLanguage(notebookOrUri)).languageId : undefined;
+		for await (const lineOfText of inputStream) {
+			if (token.isCancellationRequested) {
+				break;
+			}
+			const line = lineOfText.value;
+			if ((line.startsWith(StartDelimter) || line.startsWith(StartEmptyCellDelimter)) && (index < 0 || (endDelimiterSeen || (previousLineEndedWithEndCellMarker && previousLine)))) {
+				if (!endDelimiterSeen && previousLineEndedWithEndCellMarker && previousLine) {
+					// Last line didn't finish, emit that, but strip the end delimiter.
+					previousLine.line = previousLine.line.substring(0, previousLine.line.lastIndexOf(EndDelimter));
+					yield previousLine;
+					yield { type: 'end', index: previousLine.index };
+				}
+				previousLineEndedWithEndCellMarker = false;
+				previousLine = undefined;
+
+				index += 1;
+				endDelimiterSeen = false;
+				const lineOfCellText: LineOfCellText = { type: 'start', index, uri: undefined, language: undefined, kind: NotebookCellKind.Code };
+				const cellParts = extractCellParts(line, defaultLanguage);
+				// LLM returns duplicate cell with the same id.
+				// We need tests for this.
+				// this is a work around to treat subsequent cells as new cells.
+				if (cellParts.id && cellIdMap.get(cellParts.id)?.document.languageId === cellParts.language) {
+					if (cellIdsSeen.has(cellParts.id)) {
+						cellParts.id = '';
+					} else {
+						cellIdsSeen.add(cellParts.id);
+					}
+				} else {
+					// Possible duplicate cell with the same id but different language.
+					// In such cases, treat them as new cells.
+					cellParts.id = '';
+				}
+				const cell = cellIdMap.get(cellParts.id)?.document.languageId === cellParts.language ? cellIdMap.get(cellParts.id) : undefined;
+				lineOfCellText.id = cellParts.id;
+				lineOfCellText.language = cellParts.language;
+				lineOfCellText.uri = cell?.document.uri;
+				lineOfCellText.kind = cell?.kind || (lineOfCellText.language === 'markdown' ? NotebookCellKind.Markup : NotebookCellKind.Code);
+				yield lineOfCellText;
+			} else if (line.startsWith(EndDelimter)) {
+				if (previousLineEndedWithEndCellMarker && previousLine) {
+					// The last line somehow ends with the cell marker (must have been added by the user),
+					// yield the previous line.
+					yield previousLine;
+				}
+
+				endDelimiterSeen = true;
+				previousLineEndedWithEndCellMarker = false;
+				previousLine = undefined;
+				yield { type: 'end', index };
+			} else if (index >= 0) {
+				if (previousLineEndedWithEndCellMarker && previousLine) {
+					// Some how we have two subsequent lines that end with the cell marker,
+					// Weird, shoudl not happen, if it does, yield the previous line.
+					yield previousLine;
+					previousLine = undefined;
+				}
+				previousLineEndedWithEndCellMarker = line.endsWith(EndDelimter);
+				if (previousLineEndedWithEndCellMarker) {
+					previousLine = { type: 'line', index, line };
+				} else {
+					yield { type: 'line', index, line };
+				}
+			}
+		}
+	}
+
+	public override getAlternativeDocument(notebook: NotebookDocument, excludeMarkdownCells?: boolean): AlternativeNotebookDocument {
+		const cells = notebook.getCells().filter(cell => excludeMarkdownCells ? cell.kind !== NotebookCellKind.Markup : true).map(cell => summarize(cell));
+
+		const cellContent = cells.map(cell => {
+			const cellMarker = generateCellMarker(cell);
+			const prefix = `${cellMarker}${EOL}`;
+			return { content: `${prefix}${cell.source.join(EOL)}${EOL}${EndDelimter}`, prefix, cell: notebook.cellAt(cell.index) };
+		});
+		const content = cellContent.map(cell => cell.content).join(EOL);
+		const cellOffsetMap = cellContent.map(cellContent => ({ offset: content.indexOf(cellContent.content) + cellContent.prefix.length, cell: cellContent.cell }));
+
+		return new AlternativeXmlDocument(content, cellOffsetMap, notebook);
+	}
+
+}
+
+
+function extractCellParts(line: string, defaultLanguage: string | undefined): { id: string; language: string } {
+	const idMatch = line.match(/id="([^"]+)"/);
+	const languageMatch = line.match(/language="([^"]+)"/);
+	if (!languageMatch) {
+		if (isXmlContent(line) && typeof defaultLanguage === 'string') {
+			// If we have a cell marker but no language, we assume the default language.
+			return { id: idMatch ? idMatch[1].trim() : '', language: defaultLanguage };
+		}
+		throw new Error(`Invalid cell part in ${line}`);
+	}
+
+	// New cells will not have an id.
+	return { id: idMatch ? idMatch[1].trim() : '', language: languageMatch[1].trim() };
+}

package/src/_internal/platform/notebook/common/alternativeNotebookDocument.ts

@@ -0,0 +1,155 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { NotebookCell, NotebookDocument, TextLine } from 'vscode';
+import { DEFAULT_WORD_REGEXP, getWordAtText } from '../../../util/vs/editor/common/core/wordHelper';
+import { Position, Range } from '../../../util/common/test/shims/vscodeTypesShim';
+import { PositionOffsetTransformer } from '../../editing/common/positionOffsetTransformer';
+import { SnapshotDocumentLine } from '../../editing/common/textDocumentSnapshot';
+
+
+export abstract class AlternativeNotebookDocument {
+	private _transformer: PositionOffsetTransformer | null = null;
+	private get transformer(): PositionOffsetTransformer {
+		if (!this._transformer) {
+			this._transformer = new PositionOffsetTransformer(this._text);
+		}
+		return this._transformer;
+	}
+
+	getText(range?: Range): string {
+		return range ? this._getTextInRange(range) : this._text;
+	}
+
+	private _getTextInRange(_range: Range): string {
+		const range = this.validateRange(_range);
+
+		if (range.isEmpty) {
+			return '';
+		}
+
+		const offsetRange = this.transformer.toOffsetRange(range);
+		return this._text.substring(offsetRange.start, offsetRange.endExclusive);
+	}
+
+	constructor(protected readonly _text: string, protected readonly notebook: NotebookDocument) {
+
+	}
+
+	protected positionToOffset(position: Position): number {
+		position = this.validatePosition(position);
+		return this.transformer.getOffset(position);
+	}
+
+	/**
+	 * Translates a position in the notebook document to the corresponding alternative position.
+	 */
+	abstract fromCellPosition(cell: NotebookCell, position: Position): Position;
+
+	/**
+	 * Translates a position in the alternative document to the corresponding cell index and position in the notebook document.
+	 */
+	abstract toCellPosition(position: Position): { cell: NotebookCell; position: Position } | undefined;
+
+	getWordRangeAtPosition(_position: Position): Range | undefined {
+		const position = this.validatePosition(_position);
+
+		const wordAtText = getWordAtText(
+			position.character + 1,
+			DEFAULT_WORD_REGEXP,
+			this.lines[position.line],
+			0
+		);
+
+		if (wordAtText) {
+			return new Range(position.line, wordAtText.startColumn - 1, position.line, wordAtText.endColumn - 1);
+		}
+		return undefined;
+	}
+
+
+	private _lines: string[] | null = null;
+
+	get lines(): string[] {
+		if (!this._lines) {
+			this._lines = this._text.split(/\r\n|\r|\n/g);
+		}
+		return this._lines;
+	}
+
+	get lineCount(): number {
+		return this.lines.length;
+	}
+
+	lineAt(line: number): TextLine;
+	lineAt(position: Position): TextLine;
+	lineAt(lineOrPosition: number | Position): TextLine {
+		let line: number | undefined;
+		if (lineOrPosition instanceof Position) {
+			line = lineOrPosition.line;
+		} else if (typeof lineOrPosition === 'number') {
+			line = lineOrPosition;
+		} else {
+			throw new Error(`Invalid argument`);
+		}
+		if (line < 0 || line >= this.lines.length) {
+			throw new Error('Illegal value for `line`');
+		}
+
+		return new SnapshotDocumentLine(line, this.lines[line], line === this.lines.length - 1);
+	}
+	offsetAt(position: Position): number {
+		return this.transformer.getOffset(position);
+	}
+
+	positionAt(offset: number): Position {
+		offset = Math.floor(offset);
+		offset = Math.max(0, offset);
+
+		return this.transformer.getPosition(offset);
+	}
+	validateRange(range: Range): Range {
+		const start = this.validatePosition(range.start);
+		const end = this.validatePosition(range.end);
+
+		if (start === range.start && end === range.end) {
+			return range;
+		}
+		return new Range(start.line, start.character, end.line, end.character);
+	}
+
+	validatePosition(position: Position): Position {
+		if (this._text.length === 0) {
+			return position.with(0, 0);
+		}
+
+		let { line, character } = position;
+		let hasChanged = false;
+
+		if (line < 0) {
+			line = 0;
+			character = 0;
+			hasChanged = true;
+		} else if (line >= this.lines.length) {
+			line = this.lines.length - 1;
+			character = this.lines[line].length;
+			hasChanged = true;
+		} else {
+			const maxCharacter = this.lines[line].length;
+			if (character < 0) {
+				character = 0;
+				hasChanged = true;
+			} else if (character > maxCharacter) {
+				character = maxCharacter;
+				hasChanged = true;
+			}
+		}
+
+		if (!hasChanged) {
+			return position;
+		}
+		return new Position(line, character);
+	}
+}

package/src/_internal/platform/notebook/common/helpers.ts

@@ -0,0 +1,212 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { ChatRequest, NotebookCell, NotebookDocument, TextDocument, Uri } from 'vscode';
+import { isLocation, isUri } from '../../../util/common/types';
+import { StringSHA1 } from '../../../util/vs/base/common/hash';
+import { removeAnsiEscapeCodes } from '../../../util/vs/base/common/strings';
+import { isUriComponents, URI } from '../../../util/vs/base/common/uri';
+import { NotebookCellData, NotebookCellKind } from '../../../util/common/test/shims/vscodeTypesShim';
+import { INotebookService } from './notebookService';
+
+
+export class LineOfText {
+	readonly __lineOfTextBrand: void = undefined;
+	public readonly value: string;
+	constructor(
+		value: string
+	) {
+		this.value = value.replace(/\r$/, '');
+	}
+}
+
+/** End of Line for alternative Notebook contnt is always \n */
+export const EOL = '\n';
+export type LineOfCellText = {
+	type: 'start';
+	/**
+	 * The cell index of the cell that this line belongs to.
+	*/
+	index: number;
+	id?: string;
+	/**
+	 * The Uri of the cell that this line belongs to.
+	 * Undefined if this is a cell that doesn't belong in the actual notebook.
+	*/
+	uri?: Uri;
+	/**
+	 * Language of the cell.
+	*/
+	language?: string;
+	/**
+	 * The type of cell.
+	*/
+	kind: NotebookCellKind;
+} | {
+	type: 'line';
+	/**
+	 * A line of text from a cell. Does not include the newline character.
+	 */
+	line: string;
+	/**
+	 * The cell index of the cell that this line belongs to.
+	 */
+	index: number;
+} | {
+	type: 'end';
+	/**
+	 *
+	 * The cell index of the cell that this line belongs to.
+	 */
+	index: number;
+};
+
+export type SummaryCell = {
+	cell_type: 'code' | 'markdown';
+	language: string;
+	id: string;
+	source: string[];
+	index: number;
+};
+
+export function summarize(cell: NotebookCell): SummaryCell {
+	const cellType = cell.kind === NotebookCellKind.Code ? 'code' : 'markdown';
+	const id = getCellId(cell);
+	const source = getCellCode(cell.document);
+	return { cell_type: cellType, id, language: cell.document.languageId, source, index: cell.index };
+}
+
+export function notebookCellToCellData(cell: NotebookCell): NotebookCellData {
+	const cellData = new NotebookCellData(cell.kind, cell.document.getText(), cell.document.languageId);
+	cellData.metadata = cell.metadata;
+	cellData.executionSummary = cell.executionSummary;
+	if (cell.outputs.length) {
+		cellData.outputs = [...cell.outputs];
+	}
+	return cellData;
+}
+
+export function getCellIdMap(notebook: NotebookDocument): Map<string, NotebookCell> {
+	const cellIdMap = new Map<string, NotebookCell>();
+	notebook.getCells().forEach(cell => {
+		cellIdMap.set(getCellId(cell), cell);
+	});
+	return cellIdMap;
+}
+
+const cellIdCache = new WeakMap<NotebookCell, string>();
+
+/** The length of the hash portion of cell IDs */
+const CELL_ID_HASH_LENGTH = 8;
+
+/** Use a unique enough cell id prefix so that we can easily identify cell ids*/
+const CELL_ID_PREFIX = '#VSC-';
+
+/** RegExp to match all Cell Ids */
+export const CellIdPatternRe = new RegExp(`(\\s+|^|\\b|\\W)(#VSC-[a-f0-9]{${CELL_ID_HASH_LENGTH}})\\b`, 'gi');
+
+/**
+ * Sometimes the model may return a cellId that is not in the expected format.
+ * This function attempts to convert such cellIds to the expected format.
+ */
+export function normalizeCellId(cellId: string): string {
+	if (cellId.startsWith(CELL_ID_PREFIX)) {
+		return cellId;
+	}
+	if (cellId.startsWith('VSC-')) {
+		return `#${cellId}`;
+	}
+	if (cellId.startsWith('#V-') && cellId.length === (CELL_ID_HASH_LENGTH + 3)) {
+		return `${CELL_ID_PREFIX}${cellId.substring(3)}`;
+	}
+	if (cellId.toLowerCase().startsWith('vscode-') && cellId.length === (CELL_ID_HASH_LENGTH + 7)) {
+		return `${CELL_ID_PREFIX}${cellId.substring(7)}`;
+	}
+	if (cellId.startsWith('-')) {
+		return `#VSC${cellId}`;
+	}
+	// Possible case where the cellId is just a hash without the prefix
+	return cellId.length === CELL_ID_HASH_LENGTH ? `${CELL_ID_PREFIX}${cellId}` : cellId;
+}
+
+const notebookIdCache = new WeakMap<NotebookDocument, string>();
+export function getNotebookId(notebook: NotebookDocument): string {
+	let id = notebookIdCache.get(notebook);
+	if (id) {
+		return id;
+	}
+	const hash = new StringSHA1();
+	hash.update(notebook.uri.toString());
+	id = hash.digest();
+	notebookIdCache.set(notebook, id);
+	return id;
+}
+
+/**
+ * Given a Notebook cell returns a unique identifier for the cell.
+ * The identifier is based on the cell's URI and is cached for performance.
+ * This is useful for tracking cells across sessions or for referencing cells in a consistent manner.
+ * The cell Id will have a specicial prefix as well do as to easily identify it as a cell Id.
+ */
+export function getCellId(cell: NotebookCell): string {
+	let oldId = cellIdCache.get(cell);
+	if (oldId) {
+		return oldId;
+	}
+	const hash = new StringSHA1();
+	hash.update(cell.document.uri.toString());
+	oldId = `${CELL_ID_PREFIX}${hash.digest().substring(0, CELL_ID_HASH_LENGTH)}`;
+	cellIdCache.set(cell, oldId);
+	return oldId;
+}
+
+function getCellCode(document: TextDocument): string[] {
+	if (document.lineCount === 0) {
+		return [];
+	}
+	return new Array(document.lineCount).fill('').map((_, i) => document.lineAt(i).text);
+}
+
+export function getDefaultLanguage(notebook: NotebookDocument): string | undefined {
+	const codeCell = notebook.getCells().find(cell => cell.kind === NotebookCellKind.Code);
+	if (codeCell) {
+		return codeCell.document.languageId;
+	}
+	// Fallback for Jupyter Notebooks that do not have a code cell.
+	if (notebook.notebookType === 'jupyter-notebook') {
+		return notebook.metadata?.language_info?.name || notebook.metadata?.kernelspec?.language || 'python';
+	}
+}
+
+
+const notebookTermsToLookFor = ['jupyter', 'notebook', 'cell.', 'cells.', ' cell ', 'cells', 'notebook cell'];
+export function requestHasNotebookRefs(request: ChatRequest, notebookService: INotebookService, options?: { checkPromptAsWell: boolean }): boolean {
+	const prompt = (request.prompt || '').toLowerCase();
+	if (options?.checkPromptAsWell && notebookTermsToLookFor.some(term => prompt.includes(term))) {
+		return true;
+	}
+	return request.references.some(ref => {
+		if (isLocation(ref.value)) {
+			return notebookService.hasSupportedNotebooks(ref.value.uri);
+		}
+		if (isUriComponents(ref.value)) {
+			return notebookService.hasSupportedNotebooks(URI.revive(ref.value));
+		}
+		if (isUri(ref.value)) {
+			return notebookService.hasSupportedNotebooks(ref.value);
+		}
+		return false;
+	});
+}
+
+export function parseAndCleanStack(jsonString: string): string {
+	try {
+		// Parse the JSON string
+		const parsed = JSON.parse(jsonString) as Partial<Error>;
+		return removeAnsiEscapeCodes(parsed?.stack || parsed.message || '') || parsed.message || parsed.name || jsonString;
+	} catch {
+		return jsonString; // Return the original string if parsing fails
+	}
+}

package/src/_internal/platform/notebook/common/notebookService.ts

@@ -0,0 +1,43 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import type { NotebookCell, Uri } from 'vscode';
+import { createServiceIdentifier } from '../../../util/common/services';
+
+
+export interface Variable {
+	name: string;
+	value: string;
+	type?: string;
+	summary?: string;
+}
+
+export interface VariablesResult {
+	variable: Variable;
+	hasNamedChildren: boolean;
+	indexedChildrenCount: number;
+}
+
+export interface PipPackage {
+	name: string;
+	version: string;
+}
+
+export const INotebookService = createServiceIdentifier<INotebookService>('INotebookService');
+
+export interface INotebookService {
+	readonly _serviceBrand: undefined;
+	getVariables(notebook: Uri): Promise<VariablesResult[]>;
+	getPipPackages(notebook: Uri): Promise<PipPackage[]>;
+	getCellExecutions(notebook: Uri): NotebookCell[];
+	runCells(notebook: Uri, range: { start: number; end: number }, autoReveal: boolean): Promise<void>;
+	trackAgentUsage(): void;
+	setFollowState(state: boolean): void;
+	getFollowState(): boolean;
+	ensureKernelSelected(notebook: Uri): Promise<void>;
+	hasSupportedNotebooks(uri: Uri): boolean;
+	// testing utility
+	setVariables(notebook: Uri, variables: VariablesResult[]): void;
+}