diff --git a/lib/shared/src/completions/types.ts b/lib/shared/src/completions/types.ts index 967099f4381c..335f549bd505 100644 --- a/lib/shared/src/completions/types.ts +++ b/lib/shared/src/completions/types.ts @@ -16,11 +16,10 @@ interface AutocompleteContextSnippetMetadata { retrieverMetadata?: AutocompleteContextSnippetMetadataFields } -export interface AutocompleteFileContextSnippet { +export interface AutocompleteBaseContextSnippet { + type: 'base' identifier: string uri: URI - startLine: number - endLine: number content: string /** * Metadata populated by the context retriever. @@ -30,12 +29,21 @@ export interface AutocompleteFileContextSnippet { metadata?: AutocompleteContextSnippetMetadata } -export interface AutocompleteSymbolContextSnippet extends AutocompleteFileContextSnippet { +export interface AutocompleteFileContextSnippet extends Omit { + type: 'file' + startLine: number + endLine: number +} + +export interface AutocompleteSymbolContextSnippet extends Omit { + type: 'symbol' symbol: string } + export type AutocompleteContextSnippet = | AutocompleteFileContextSnippet | AutocompleteSymbolContextSnippet + | AutocompleteBaseContextSnippet export interface DocumentContext extends DocumentDependentContext, LinesContext { position: vscode.Position diff --git a/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts b/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts index dc0575c4c18b..91f32a01e585 100644 --- a/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts +++ b/vscode/src/autoedits/prompt/default-prompt-strategy.test.ts @@ -17,6 +17,7 @@ describe('DefaultUserPromptStrategy', () => { identifier: string, filePath = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(filePath), diff --git a/vscode/src/autoedits/prompt/prompt-utils.test.ts b/vscode/src/autoedits/prompt/prompt-utils.test.ts index d3c78723042b..9088eacd7594 100644 --- a/vscode/src/autoedits/prompt/prompt-utils.test.ts +++ b/vscode/src/autoedits/prompt/prompt-utils.test.ts @@ -1,9 +1,13 @@ -import { type AutocompleteContextSnippet, ps, testFileUri } from '@sourcegraph/cody-shared' import dedent from 'dedent' -import { describe, expect, it } from 'vitest' +import { beforeEach, describe, expect, it, vi } from 'vitest' +import * as vscode from 'vscode' +import { NotebookCellKind } from 'vscode-languageserver-protocol' + +import { type AutocompleteContextSnippet, ps, testFileUri } from '@sourcegraph/cody-shared' + import { RetrieverIdentifier } from '../../completions/context/utils' import { getCurrentDocContext } from '../../completions/get-current-doc-context' -import { documentAndPosition } from '../../completions/test-helpers' +import { documentAndPosition, mockNotebookAndPosition } from '../../completions/test-helpers' import { getCompletionsPromptWithSystemPrompt, getContextItemsInTokenBudget, @@ -17,7 +21,18 @@ import { getRecentEditsPrompt, getRecentlyViewedSnippetsPrompt, joinPromptsWithNewlineSeperator, -} from '../prompt/prompt-utils' +} from './prompt-utils' + +// A helper to set up your global "activeNotebookEditor" mock. +function mockActiveNotebookEditor(notebook: vscode.NotebookDocument | undefined) { + vi.spyOn(vscode.window, 'activeNotebookEditor', 'get').mockReturnValue( + notebook + ? ({ + notebook, + } as vscode.NotebookEditor) + : undefined + ) +} describe('getContextPromptWithPath', () => { it('correct prompt with path', () => { @@ -52,12 +67,14 @@ describe('getCurrentFilePromptComponents', () => { const content = `${longPrefix}\ncursor█line\n${longSuffix}` const { document, position } = documentAndPosition(content) + const maxPrefixLength = 100 + const maxSuffixLength = 100 const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -115,11 +132,14 @@ describe('getCurrentFilePromptComponents', () => { const { document, position } = documentAndPosition(content) + const maxPrefixLength = 30 + const maxSuffixLength = 30 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 30, - maxSuffixLength: 30, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -130,6 +150,8 @@ describe('getCurrentFilePromptComponents', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFilePromptComponents(options) @@ -159,14 +181,145 @@ describe('getCurrentFilePromptComponents', () => { }) describe('getCurrentFileContext', () => { + beforeEach(() => { + vi.restoreAllMocks() + vi.clearAllMocks() + }) + + it('correctly handles the notebook document with only code cells', async () => { + const { notebookDoc, position } = mockNotebookAndPosition({ + uri: 'file://test.ipynb', + cells: [ + { + kind: NotebookCellKind.Code, + text: 'console.log("cell0 code")', + languageId: 'python', + }, + { + kind: NotebookCellKind.Code, + text: 'console.log("cell1█ code")', + languageId: 'python', + }, + { + kind: NotebookCellKind.Code, + text: 'console.log("cell2 code")', + languageId: 'python', + }, + ], + }) + mockActiveNotebookEditor(notebookDoc) + const document = notebookDoc.cellAt(1).document + + const maxPrefixLength = 100 + const maxSuffixLength = 100 + + const docContext = getCurrentDocContext({ + document, + position, + maxPrefixLength, + maxSuffixLength, + }) + + const options = { + docContext, + document, + position, + maxPrefixLinesInArea: 1, + maxSuffixLinesInArea: 1, + codeToRewritePrefixLines: 1, + codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, + } + + const result = getCurrentFileContext(options) + + // Verify the results + expect(result.codeToRewrite.toString()).toBe('console.log("cell1 code")') + expect(result.codeToRewritePrefix.toString()).toBe('console.log("cell1') + expect(result.codeToRewriteSuffix.toString()).toBe(' code")') + expect(result.prefixInArea.toString()).toBe('') + expect(result.suffixInArea.toString()).toBe('') + expect(result.prefixBeforeArea.toString()).toBe('console.log("cell0 code")\n') + expect(result.suffixAfterArea.toString()).toBe('\nconsole.log("cell2 code")') + }) + + it('correctly handles the notebook document and relevant context', async () => { + const { notebookDoc, position } = mockNotebookAndPosition({ + uri: 'file://test.ipynb', + cells: [ + { + kind: NotebookCellKind.Code, + text: 'console.log("cell0 code")', + languageId: 'python', + }, + { + kind: NotebookCellKind.Code, + text: 'console.log("cell1█ code")', + languageId: 'python', + }, + { + kind: NotebookCellKind.Markup, + text: '# This is a markdown cell', + languageId: 'markdown', + }, + { + kind: NotebookCellKind.Code, + text: 'console.log("cell2 code")', + languageId: 'python', + }, + ], + }) + mockActiveNotebookEditor(notebookDoc) + const document = notebookDoc.cellAt(1).document + + const maxPrefixLength = 100 + const maxSuffixLength = 100 + + const docContext = getCurrentDocContext({ + document, + position, + maxPrefixLength, + maxSuffixLength, + }) + + const options = { + docContext, + document, + position, + maxPrefixLinesInArea: 1, + maxSuffixLinesInArea: 1, + codeToRewritePrefixLines: 1, + codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, + } + + const result = getCurrentFileContext(options) + + // Verify the results + expect(result.codeToRewrite.toString()).toBe('console.log("cell1 code")') + expect(result.codeToRewritePrefix.toString()).toBe('console.log("cell1') + expect(result.codeToRewriteSuffix.toString()).toBe(' code")') + expect(result.prefixInArea.toString()).toBe('') + expect(result.suffixInArea.toString()).toBe('') + expect(result.prefixBeforeArea.toString()).toBe('console.log("cell0 code")\n') + expect(result.suffixAfterArea.toString()).toBe( + '\n# # This is a markdown cell\n\nconsole.log("cell2 code")' + ) + }) + it('correctly splits content into different areas based on cursor position', () => { const { document, position } = documentAndPosition('line1\nline2\nline3█line4\nline5\nline6') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -177,6 +330,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -196,11 +351,14 @@ describe('getCurrentFileContext', () => { it('handles cursor at start of line', () => { const { document, position } = documentAndPosition('line1\nline2\n█line3\nline4\nline5') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -211,6 +369,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -226,11 +386,14 @@ describe('getCurrentFileContext', () => { it('handles single line content', () => { const { document, position } = documentAndPosition('const foo = █bar') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -241,6 +404,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -257,11 +422,14 @@ describe('getCurrentFileContext', () => { it('handles cursor at start of file', () => { const { document, position } = documentAndPosition('█line1\nline2\nline3') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -272,6 +440,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -288,11 +458,14 @@ describe('getCurrentFileContext', () => { it('handles cursor at end of file', () => { const { document, position } = documentAndPosition('line1\nline2\nline3█') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -303,6 +476,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -321,11 +496,14 @@ describe('getCurrentFileContext', () => { 'line1\nline2\nline3\nline4\nline5\n█line6\nline7' ) + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -336,6 +514,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 3, // Increased prefix lines codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -357,11 +537,14 @@ describe('getCurrentFileContext', () => { const { document, position } = documentAndPosition(content) + const maxPrefixLength = 30 + const maxSuffixLength = 30 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 30, - maxSuffixLength: 30, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -372,6 +555,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 1, codeToRewritePrefixLines: 1, codeToRewriteSuffixLines: 1, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -394,11 +579,14 @@ describe('getCurrentFileContext', () => { const { document, position } = documentAndPosition(content) + const maxPrefixLength = 20 + const maxSuffixLength = 20 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 20, - maxSuffixLength: 20, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -409,6 +597,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 2, codeToRewritePrefixLines: 2, codeToRewriteSuffixLines: 2, + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -426,11 +616,14 @@ describe('getCurrentFileContext', () => { it('handles file shorter than requested ranges', () => { const { document, position } = documentAndPosition('line1\n█line2\nline3\n') + const maxPrefixLength = 100 + const maxSuffixLength = 100 + const docContext = getCurrentDocContext({ document, position, - maxPrefixLength: 100, - maxSuffixLength: 100, + maxPrefixLength, + maxSuffixLength, }) const options = { @@ -441,6 +634,8 @@ describe('getCurrentFileContext', () => { maxSuffixLinesInArea: 5, // Larger than file codeToRewritePrefixLines: 3, // Larger than file codeToRewriteSuffixLines: 3, // Larger than file + maxPrefixLength, + maxSuffixLength, } const result = getCurrentFileContext(options) @@ -459,6 +654,7 @@ describe('getCurrentFileContext', () => { describe('getContextItemsInTokenBudget', () => { const getContextItem = (content: string, identifier: string): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri('foo.ts'), @@ -532,6 +728,7 @@ describe('getLintErrorsPrompt', () => { identifier: string, fileName = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(fileName), @@ -610,6 +807,7 @@ describe('getRecentCopyPrompt', () => { identifier: string, fileName = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(fileName), @@ -668,6 +866,7 @@ describe('getRecentEditsPrompt', () => { identifier: string, fileName = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(fileName), @@ -750,6 +949,7 @@ describe('getRecentlyViewedSnippetsPrompt', () => { identifier: string, fileName = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(fileName), @@ -828,6 +1028,7 @@ describe('getJaccardSimilarityPrompt', () => { identifier: string, fileName = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(fileName), diff --git a/vscode/src/autoedits/prompt/prompt-utils.ts b/vscode/src/autoedits/prompt/prompt-utils.ts index 3958c4c151b4..318e913c5c00 100644 --- a/vscode/src/autoedits/prompt/prompt-utils.ts +++ b/vscode/src/autoedits/prompt/prompt-utils.ts @@ -7,6 +7,12 @@ import { } from '@sourcegraph/cody-shared' import { Uri } from 'vscode' import * as vscode from 'vscode' +import { getTextFromNotebookCells } from '../../completions/context/retrievers/recent-user-actions/notebook-utils' +import { + getActiveNotebookUri, + getCellIndexInActiveNotebookEditor, + getNotebookCells, +} from '../../completions/context/retrievers/recent-user-actions/notebook-utils' import { RetrieverIdentifier } from '../../completions/context/utils' import { autoeditsOutputChannelLogger } from '../output-channel-logger' import { clip, splitLinesKeepEnds } from '../utils' @@ -40,6 +46,7 @@ export interface CurrentFilePromptResponse { } interface CurrentFileContext { + filePath: PromptString codeToRewrite: PromptString codeToRewritePrefix: PromptString codeToRewriteSuffix: PromptString @@ -100,7 +107,7 @@ export function getCurrentFilePromptComponents( ) const filePrompt = getCurrentFileContextPromptWithPath( - PromptString.fromDisplayPath(options.document.uri), + currentFileContext.filePath, joinPromptsWithNewlineSeperator( constants.FILE_TAG_OPEN, fileWithMarker, @@ -171,19 +178,84 @@ export function getCurrentFileContext(options: CurrentFilePromptOptions): Curren prefixBeforeArea: new vscode.Range(positionAtLineStart(minLine), positionAtLineStart(areaStart)), suffixAfterArea: new vscode.Range(positionAtLineEnd(areaEnd), positionAtLineEnd(maxLine)), } + const { prefixBeforeArea, suffixAfterArea } = getUpdatedCurrentFilePrefixAndSuffixOutsideArea( + document, + ranges.prefixBeforeArea, + ranges.suffixAfterArea + ) + // Convert ranges to PromptStrings return { + filePath: getCurrentFilePath(document), codeToRewrite: PromptString.fromDocumentText(document, ranges.codeToRewrite), codeToRewritePrefix: PromptString.fromDocumentText(document, ranges.codeToRewritePrefix), codeToRewriteSuffix: PromptString.fromDocumentText(document, ranges.codeToRewriteSuffix), prefixInArea: PromptString.fromDocumentText(document, ranges.prefixInArea), suffixInArea: PromptString.fromDocumentText(document, ranges.suffixInArea), - prefixBeforeArea: PromptString.fromDocumentText(document, ranges.prefixBeforeArea), - suffixAfterArea: PromptString.fromDocumentText(document, ranges.suffixAfterArea), + prefixBeforeArea, + suffixAfterArea, range: ranges.codeToRewrite, } } +export function getCurrentFilePath(document: vscode.TextDocument): PromptString { + const uri = + document.uri.scheme === 'vscode-notebook-cell' + ? getActiveNotebookUri() ?? document.uri + : document.uri + return PromptString.fromDisplayPath(uri) +} + +export function getUpdatedCurrentFilePrefixAndSuffixOutsideArea( + document: vscode.TextDocument, + rangePrefixBeforeArea: vscode.Range, + rangeSuffixAfterArea: vscode.Range +): { + prefixBeforeArea: PromptString + suffixAfterArea: PromptString +} { + const { prefixBeforeAreaForNotebook, suffixAfterAreaForNotebook } = + getPrefixAndSuffixForAreaForNotebook(document) + + const prefixBeforeArea = ps`${prefixBeforeAreaForNotebook}${PromptString.fromDocumentText( + document, + rangePrefixBeforeArea + )}` + + const suffixAfterArea = ps`${PromptString.fromDocumentText( + document, + rangeSuffixAfterArea + )}${suffixAfterAreaForNotebook}` + + return { + prefixBeforeArea, + suffixAfterArea, + } +} + +export function getPrefixAndSuffixForAreaForNotebook(document: vscode.TextDocument): { + prefixBeforeAreaForNotebook: PromptString + suffixAfterAreaForNotebook: PromptString +} { + const currentCellIndex = getCellIndexInActiveNotebookEditor(document) + if (currentCellIndex === -1) { + return { + prefixBeforeAreaForNotebook: ps``, + suffixAfterAreaForNotebook: ps``, + } + } + const activeNotebook = vscode.window.activeNotebookEditor?.notebook! + const notebookCells = getNotebookCells(activeNotebook) + const cellsBeforeCurrentCell = notebookCells.slice(0, currentCellIndex) + const cellsAfterCurrentCell = notebookCells.slice(currentCellIndex + 1) + const beforeContent = getTextFromNotebookCells(activeNotebook, cellsBeforeCurrentCell) + const afterContent = getTextFromNotebookCells(activeNotebook, cellsAfterCurrentCell) + return { + prefixBeforeAreaForNotebook: ps`${beforeContent}\n`, + suffixAfterAreaForNotebook: ps`\n${afterContent}`, + } +} + export function getLintErrorsPrompt(contextItems: AutocompleteContextSnippet[]): PromptString { const lintErrors = getContextItemsForIdentifier( contextItems, diff --git a/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts b/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts index 24fca6d98e9f..d810ce0fade8 100644 --- a/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts +++ b/vscode/src/autoedits/prompt/short-term-diff-prompt-strategy.test.ts @@ -21,6 +21,7 @@ describe('ShortTermPromptStrategy', () => { identifier: string, filePath = 'foo.ts' ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(filePath), @@ -385,6 +386,7 @@ describe('ShortTermPromptStrategy', () => { filePath = 'foo.ts', identifier: string = RetrieverIdentifier.RecentViewPortRetriever ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(filePath), @@ -451,6 +453,7 @@ describe('ShortTermPromptStrategy', () => { filePath = 'foo.ts', identifier: string = RetrieverIdentifier.RecentEditsRetriever ): AutocompleteContextSnippet => ({ + type: 'file', content, identifier, uri: testFileUri(filePath), diff --git a/vscode/src/completions/analytics-logger.ts b/vscode/src/completions/analytics-logger.ts index c4b2ed7d341d..0d54d717978f 100644 --- a/vscode/src/completions/analytics-logger.ts +++ b/vscode/src/completions/analytics-logger.ts @@ -71,8 +71,8 @@ interface InlineCompletionItemRetrievedContext { identifier: string content: string filePath: string - startLine: number - endLine: number + startLine?: number + endLine?: number } interface InlineContextItemsParams extends GitIdentifiersForFile { @@ -767,19 +767,18 @@ function getInlineContextItemContext( suffix: docContext.completeSuffix.slice(0, MAX_PREFIX_SUFFIX_SIZE_BYTES), triggerLine: position.line, triggerCharacter: position.character, - context: inlineContextParams.context.map( - ({ identifier, content, startLine, endLine, uri, metadata }) => ({ - identifier, - content, - startLine, - endLine, - filePath: displayPathWithoutWorkspaceFolderPrefix(uri), - metadata, - }) - ), + context: inlineContextParams.context.map(snippet => ({ + identifier: snippet.identifier, + content: snippet.content, + filePath: displayPathWithoutWorkspaceFolderPrefix(snippet.uri), + metadata: snippet.metadata, + ...(snippet.type !== 'base' && { + startLine: snippet.startLine, + endLine: snippet.endLine, + }), + })), } } - function suggestionDocumentDiffTracker( interactionId: CompletionAnalyticsID, document: vscode.TextDocument, diff --git a/vscode/src/completions/context/completions-context-ranker.test.ts b/vscode/src/completions/context/completions-context-ranker.test.ts index ebd7ea0821d1..e93bf59d6893 100644 --- a/vscode/src/completions/context/completions-context-ranker.test.ts +++ b/vscode/src/completions/context/completions-context-ranker.test.ts @@ -7,6 +7,7 @@ import type { RetrievedContextResults } from './completions-context-ranker' describe('DefaultCompletionsContextRanker', () => { describe('getContextSnippetsAsPerTimeBasedStrategy', () => { const getContextSnippet = (time?: number): AutocompleteContextSnippet => ({ + type: 'file', identifier: 'test', uri: vscode.Uri.parse('file:///test'), startLine: 1, diff --git a/vscode/src/completions/context/completions-context-ranker.ts b/vscode/src/completions/context/completions-context-ranker.ts index 2b65af5a6c4a..592a1f6f173b 100644 --- a/vscode/src/completions/context/completions-context-ranker.ts +++ b/vscode/src/completions/context/completions-context-ranker.ts @@ -131,12 +131,9 @@ export class DefaultCompletionsContextRanker implements CompletionsContextRanker const fusedResults = fuseResults( results.map(r => r.snippets), result => { - // Ensure that context retrieved works when we do not have a startLine and - // endLine yet. - if (typeof result.startLine === 'undefined' || typeof result.endLine === 'undefined') { + if (result.type === 'base') { return [result.uri.toString()] } - const lineIds = [] for (let i = result.startLine; i <= result.endLine; i++) { lineIds.push(`${result.uri.toString()}:${i}`) diff --git a/vscode/src/completions/context/context-mixer.test.ts b/vscode/src/completions/context/context-mixer.test.ts index 481cf8a158ae..efea853009b3 100644 --- a/vscode/src/completions/context/context-mixer.test.ts +++ b/vscode/src/completions/context/context-mixer.test.ts @@ -96,6 +96,7 @@ describe('ContextMixer', () => { strategyFactory: createMockStrategy([ [ { + type: 'file', identifier: 'jaccard-similarity', uri: testFileUri('foo.ts'), content: 'function foo() {}', @@ -103,6 +104,7 @@ describe('ContextMixer', () => { endLine: 0, }, { + type: 'file', identifier: 'jaccard-similarity', uri: testFileUri('bar.ts'), content: 'function bar() {}', @@ -121,6 +123,7 @@ describe('ContextMixer', () => { identifier: 'jaccard-similarity', startLine: 0, endLine: 0, + type: 'file', }, { fileName: 'bar.ts', @@ -128,6 +131,7 @@ describe('ContextMixer', () => { identifier: 'jaccard-similarity', startLine: 0, endLine: 0, + type: 'file', }, ]) expect(logSummary).toEqual({ @@ -155,6 +159,7 @@ describe('ContextMixer', () => { strategyFactory: createMockStrategy([ [ { + type: 'file', identifier: 'retriever1', uri: testFileUri('foo.ts'), content: 'function foo1() {}', @@ -162,6 +167,7 @@ describe('ContextMixer', () => { endLine: 0, }, { + type: 'file', identifier: 'retriever1', uri: testFileUri('bar.ts'), content: 'function bar1() {}', @@ -172,6 +178,7 @@ describe('ContextMixer', () => { [ { + type: 'file', identifier: 'retriever2', uri: testFileUri('foo.ts'), content: 'function foo3() {}', @@ -179,6 +186,7 @@ describe('ContextMixer', () => { endLine: 10, }, { + type: 'file', identifier: 'retriever2', uri: testFileUri('foo.ts'), content: 'function foo1() {}\nfunction foo2() {}', @@ -186,6 +194,7 @@ describe('ContextMixer', () => { endLine: 1, }, { + type: 'file', identifier: 'retriever2', uri: testFileUri('bar.ts'), content: 'function bar1() {}\nfunction bar2() {}', @@ -209,6 +218,7 @@ describe('ContextMixer', () => { "fileName": "foo.ts", "identifier": "retriever1", "startLine": 0, + "type": "file", }, { "content": "function foo1() {} @@ -217,6 +227,7 @@ describe('ContextMixer', () => { "fileName": "foo.ts", "identifier": "retriever2", "startLine": 0, + "type": "file", }, { "content": "function bar1() {}", @@ -224,6 +235,7 @@ describe('ContextMixer', () => { "fileName": "bar.ts", "identifier": "retriever1", "startLine": 0, + "type": "file", }, { "content": "function bar1() {} @@ -232,6 +244,7 @@ describe('ContextMixer', () => { "fileName": "bar.ts", "identifier": "retriever2", "startLine": 0, + "type": "file", }, { "content": "function foo3() {}", @@ -239,6 +252,7 @@ describe('ContextMixer', () => { "fileName": "foo.ts", "identifier": "retriever2", "startLine": 10, + "type": "file", }, ] `) @@ -283,6 +297,7 @@ describe('ContextMixer', () => { strategyFactory: createMockStrategy([ [ { + type: 'file', identifier: 'retriever1', uri: testFileUri('foo.ts'), content: 'function foo1() {}', @@ -290,6 +305,7 @@ describe('ContextMixer', () => { endLine: 0, }, { + type: 'file', identifier: 'retriever1', uri: testFileUri('foo/bar.ts'), content: 'function bar1() {}', @@ -299,6 +315,7 @@ describe('ContextMixer', () => { ], [ { + type: 'file', identifier: 'retriever2', uri: testFileUri('test/foo.ts'), content: 'function foo3() {}', @@ -306,6 +323,7 @@ describe('ContextMixer', () => { endLine: 10, }, { + type: 'file', identifier: 'retriever2', uri: testFileUri('foo.ts'), content: 'function foo1() {}\nfunction foo2() {}', @@ -313,6 +331,7 @@ describe('ContextMixer', () => { endLine: 1, }, { + type: 'file', identifier: 'retriever2', uri: testFileUri('example/bar.ts'), content: 'function bar1() {}\nfunction bar2() {}', diff --git a/vscode/src/completions/context/retrievers/jaccard-similarity/jaccard-similarity-retriever.ts b/vscode/src/completions/context/retrievers/jaccard-similarity/jaccard-similarity-retriever.ts index 40d2264bb9d6..397a75e302b3 100644 --- a/vscode/src/completions/context/retrievers/jaccard-similarity/jaccard-similarity-retriever.ts +++ b/vscode/src/completions/context/retrievers/jaccard-similarity/jaccard-similarity-retriever.ts @@ -89,7 +89,7 @@ export class JaccardSimilarityRetriever extends CachedRetriever implements Conte continue } - matches.push({ ...match, uri }) + matches.push({ type: 'file', uri, ...match }) } } @@ -245,6 +245,7 @@ export class JaccardSimilarityRetriever extends CachedRetriever implements Conte } interface JaccardMatchWithFilename extends JaccardMatch { + type: 'file' uri: URI } diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts index 9bb60b41c0a7..5921b0b60828 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.test.ts @@ -426,9 +426,15 @@ describe('DiagnosticsRetriever', () => { diagnostics ) expect(snippets).toHaveLength(3) - expect(snippets[0].startLine).toBe(9) - expect(snippets[1].startLine).toBe(13) - expect(snippets[2].startLine).toBe(5) + const expectedStartLines: number[] = [9, 13, 5] + for (const [index, snippet] of snippets.entries()) { + expect(snippet.type).toBe('file') + expect(snippet).toHaveProperty('startLine') + expect(snippet).toHaveProperty('endLine') + if (snippet.type === 'file') { + expect(snippet.startLine).toBe(expectedStartLines[index]) + } + } }) }) diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts index cad338b9ad81..805de30ec4b6 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/diagnostics-retriever.ts @@ -4,6 +4,7 @@ import { XMLBuilder } from 'fast-xml-parser' import * as vscode from 'vscode' import type { ContextRetriever, ContextRetrieverOptions } from '../../../types' import { RetrieverIdentifier } from '../../utils' +import { getCellIndexInActiveNotebookEditor, getNotebookCells } from './notebook-utils' // XML builder instance for formatting diagnostic messages const XML_BUILDER = new XMLBuilder({ format: true }) @@ -40,12 +41,48 @@ export class DiagnosticsRetriever implements vscode.Disposable, ContextRetriever this.useCaretToIndicateErrorLocation = options.useCaretToIndicateErrorLocation ?? true } - public async retrieve({ + public retrieve({ document, position, }: ContextRetrieverOptions): Promise { + if (getCellIndexInActiveNotebookEditor(document) !== -1) { + // Handle the diagnostic error for the notebook + return this.getDiagnosticsForNotebook(position) + } + return this.getDiagnosticsForDocument(document, position) + } + + private async getDiagnosticsForNotebook( + position: vscode.Position + ): Promise { + const activeNotebook = vscode.window.activeNotebookEditor?.notebook + if (!activeNotebook) { + return [] + } + const notebookCells = getNotebookCells(activeNotebook) + const diagnostics = await Promise.all( + notebookCells.map(cell => { + const diagnostics = vscode.languages.getDiagnostics(cell.document.uri) + return this.getDiagnosticsPromptFromInformation(cell.document, position, diagnostics) + }) + ) + return diagnostics.flat().map(snippet => ({ + ...snippet, + uri: activeNotebook!.uri, + })) + } + + private async getDiagnosticsForDocument( + document: vscode.TextDocument, + position: vscode.Position + ): Promise { const diagnostics = vscode.languages.getDiagnostics(document.uri) - return this.getDiagnosticsPromptFromInformation(document, position, diagnostics) + const diagnosticsSnippets = await this.getDiagnosticsPromptFromInformation( + document, + position, + diagnostics + ) + return diagnosticsSnippets } public async getDiagnosticsPromptFromInformation( @@ -61,6 +98,7 @@ export class DiagnosticsRetriever implements vscode.Disposable, ContextRetriever ) return Promise.all( diagnosticInfos.map(async info => ({ + type: 'file', identifier: this.identifier, content: await this.getDiagnosticPromptMessage(info), uri: document.uri, diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/notebook-utils.ts b/vscode/src/completions/context/retrievers/recent-user-actions/notebook-utils.ts new file mode 100644 index 000000000000..8472c7e2d1c3 --- /dev/null +++ b/vscode/src/completions/context/retrievers/recent-user-actions/notebook-utils.ts @@ -0,0 +1,73 @@ +import { PromptString, ps } from '@sourcegraph/cody-shared' +import * as vscode from 'vscode' +import { getNewLineChar } from '../../../../completions/text-processing' +import { getLanguageConfig } from '../../../../tree-sitter/language' + +export function getTextFromNotebookCells( + notebook: vscode.NotebookDocument, + cells: vscode.NotebookCell[] +): PromptString { + const orderedCells = cells.sort((a, b) => a.index - b.index) + const cellText: PromptString[] = [] + const languageId = getNotebookLanguageId(notebook) + for (const cell of orderedCells) { + const text = PromptString.fromDocumentText(cell.document) + if (text.trim().length === 0) { + continue + } + if (cell.kind === vscode.NotebookCellKind.Code) { + cellText.push(text) + } else if (cell.kind === vscode.NotebookCellKind.Markup) { + // Add the markdown content as a comment + cellText.push(getCellMarkupContent(languageId, text)) + } + } + return PromptString.join(cellText, ps`\n\n`) +} + +export function getNotebookLanguageId(notebook: vscode.NotebookDocument): string { + const cells = notebook.getCells() + for (const cell of cells) { + if (cell.kind === vscode.NotebookCellKind.Code) { + return cell.document.languageId + } + } + return cells.length > 0 ? cells[0].document.languageId : '' +} + +export function getCellMarkupContent(languageId: string, text: PromptString): PromptString { + if (text.trim().length === 0) { + return ps`` + } + const languageConfig = getLanguageConfig(languageId) + const commentStart = languageConfig ? languageConfig.commentStart : ps`// ` + const newLineChar = getNewLineChar(text.toString()) + + const contentLines = text.split(newLineChar).map(line => ps`${commentStart}${line}`) + return PromptString.join(contentLines, ps`\n`) +} + +/** + * Returns the index of a notebook cell within the currently active notebook editor. + * Each cell in the notebook is treated as a seperate document, so this function + * can be used to find the index of a cell within the notebook. + * @param document The VS Code text document to find the index for + * @returns The zero-based index of the cell within the notebook cells, or -1 if not found or no active notebook + */ +export function getCellIndexInActiveNotebookEditor(document: vscode.TextDocument): number { + const activeNotebook = vscode.window.activeNotebookEditor?.notebook + if (!activeNotebook || document.uri.scheme !== 'vscode-notebook-cell') { + return -1 + } + const notebookCells = getNotebookCells(activeNotebook) + const currentCellIndex = notebookCells.findIndex(cell => cell.document === document) + return currentCellIndex +} + +export function getNotebookCells(notebook: vscode.NotebookDocument): vscode.NotebookCell[] { + return notebook.getCells().sort((a, b) => a.index - b.index) +} + +export function getActiveNotebookUri(): vscode.Uri | undefined { + return vscode.window.activeNotebookEditor?.notebook.uri +} diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts index dccd87354cae..c048203a4026 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.test.ts @@ -85,6 +85,7 @@ describe('RecentCopyRetriever', () => { expect(snippets).toHaveLength(1) expect(snippets[0]).toEqual({ + type: 'file', content: mockClipboardContent, uri: testDocument.uri, startLine: selection.start.line, diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts index da6c08911601..aec864aac6e9 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-copy.ts @@ -42,6 +42,7 @@ export class RecentCopyRetriever implements vscode.Disposable, ContextRetriever const selectionItem = this.getSelectionItemIfExist(clipboardContent) if (selectionItem) { const autocompleteItem: AutocompleteContextSnippet = { + type: 'file', identifier: this.identifier, content: selectionItem.content, uri: selectionItem.uri, diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts index 02858df60545..22648d989b9c 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-retriever.ts @@ -55,7 +55,8 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever const retrievalTriggerTime = Date.now() for (const diff of diffs) { const content = diff.diff.toString() - const autocompleteSnippet = { + const autocompleteSnippet: AutocompleteContextSnippet = { + type: 'base', uri: diff.uri, identifier: this.identifier, content, @@ -63,12 +64,9 @@ export class RecentEditsRetriever implements vscode.Disposable, ContextRetriever timeSinceActionMs: retrievalTriggerTime - diff.latestChangeTimestamp, retrieverMetadata: diff.diffStrategyMetadata, }, - } satisfies Omit + } autocompleteContextSnippets.push(autocompleteSnippet) } - // remove the startLine and endLine from the response similar to how we did - // it for BFG. - // @ts-ignore return autocompleteContextSnippets } diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-tracker.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-tracker.ts index 2d03f304a039..b3ef35a40eca 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-tracker.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-edits-tracker.ts @@ -88,16 +88,15 @@ export class RecentEditsTracker implements vscode.Disposable { } private trackDocument(document: vscode.TextDocument): void { - if (document.uri.scheme !== 'file') { - return - } - const trackedDocument: TrackedDocument = { - content: document.getText(), - languageId: document.languageId, - uri: document.uri, - changes: [], + if (document.uri.scheme === 'file' || document.uri.scheme === 'vscode-notebook-cell') { + const trackedDocument: TrackedDocument = { + content: document.getText(), + languageId: document.languageId, + uri: document.uri, + changes: [], + } + this.trackedDocuments.set(document.uri.toString(), trackedDocument) } - this.trackedDocuments.set(document.uri.toString(), trackedDocument) } private reconcileOutdatedChanges(): void { diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts index 5109fc69d7ee..e73d52d3c107 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.test.ts @@ -91,6 +91,12 @@ describe('RecentViewPortRetriever', () => { onDidChangeActiveTextEditor() { return { dispose: () => {} } }, + onDidChangeNotebookEditorVisibleRanges() { + return { dispose: () => {} } + }, + onDidChangeActiveNotebookEditor() { + return { dispose: () => {} } + }, activeTextEditor: undefined, } ) @@ -135,8 +141,6 @@ describe('RecentViewPortRetriever', () => { expect(snippets).toHaveLength(1) expect(snippets[0]).toMatchObject({ uri: doc.uri, - startLine: 1, - endLine: 2, identifier: retriever.identifier, }) expect(snippets[0].content).toMatchInlineSnapshot(dedent` @@ -156,8 +160,6 @@ describe('RecentViewPortRetriever', () => { const snippets = await retriever.retrieve(getContextRetrieverOptionsFromDoc(doc2)) expect(snippets).toHaveLength(1) - expect(snippets[0].startLine).toBe(1) - expect(snippets[0].endLine).toBe(2) }) it('should handle empty visible ranges', async () => { diff --git a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts index 6f657b7b1cf5..904797246010 100644 --- a/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts +++ b/vscode/src/completions/context/retrievers/recent-user-actions/recent-view-port.ts @@ -4,10 +4,22 @@ import { LRUCache } from 'lru-cache' import * as vscode from 'vscode' import type { ContextRetriever, ContextRetrieverOptions } from '../../../types' import { RetrieverIdentifier, type ShouldUseContextParams, shouldBeUsedAsContext } from '../../utils' +import { + getActiveNotebookUri, + getCellIndexInActiveNotebookEditor, + getNotebookLanguageId, + getTextFromNotebookCells, +} from './notebook-utils' + +interface TrackViewPortLines { + startLine: number + endLine: number +} interface TrackedViewPort { uri: vscode.Uri - visibleRange: vscode.Range + content: string + lines?: TrackViewPortLines languageId: string lastAccessTimestamp: number } @@ -29,7 +41,11 @@ export class RecentViewPortRetriever implements vscode.Disposable, ContextRetrie options: RecentViewPortRetrieverOptions, readonly window: Pick< typeof vscode.window, - 'onDidChangeTextEditorVisibleRanges' | 'onDidChangeActiveTextEditor' | 'activeTextEditor' + | 'onDidChangeTextEditorVisibleRanges' + | 'onDidChangeActiveTextEditor' + | 'activeTextEditor' + | 'onDidChangeNotebookEditorVisibleRanges' + | 'onDidChangeActiveNotebookEditor' > = vscode.window ) { this.maxTrackedViewPorts = options.maxTrackedViewPorts @@ -42,7 +58,11 @@ export class RecentViewPortRetriever implements vscode.Disposable, ContextRetrie window.onDidChangeTextEditorVisibleRanges( debounce(this.onDidChangeTextEditorVisibleRanges.bind(this), 300) ), - window.onDidChangeActiveTextEditor(this.onDidChangeActiveTextEditor.bind(this)) + window.onDidChangeNotebookEditorVisibleRanges( + debounce(this.onDidChangeNotebookEditorVisibleRanges.bind(this), 300) + ), + window.onDidChangeActiveTextEditor(this.onDidChangeActiveTextEditor.bind(this)), + window.onDidChangeActiveNotebookEditor(this.onDidChangeActiveNotebookEditor.bind(this)) ) } @@ -50,27 +70,32 @@ export class RecentViewPortRetriever implements vscode.Disposable, ContextRetrie const sortedViewPorts = this.getValidViewPorts(document) const snippetPromises = sortedViewPorts.map(async viewPort => { - const document = await vscode.workspace.openTextDocument(viewPort.uri) - const content = document.getText(viewPort.visibleRange) - - const snippet: AutocompleteContextSnippet = { + const baseSnippetData = { uri: viewPort.uri, - content, - startLine: viewPort.visibleRange.start.line, - endLine: viewPort.visibleRange.end.line, + content: viewPort.content, identifier: this.identifier, metadata: { timeSinceActionMs: Date.now() - viewPort.lastAccessTimestamp, }, } + const snippet: AutocompleteContextSnippet = + viewPort.lines !== undefined + ? { + type: 'file', + ...viewPort.lines, + ...baseSnippetData, + } + : { + type: 'base', + ...baseSnippetData, + } return snippet }) - const viewPortSnippets = await Promise.all(snippetPromises) - return viewPortSnippets + return Promise.all(snippetPromises) } private getValidViewPorts(document: vscode.TextDocument): TrackedViewPort[] { - const currentFileUri = document.uri.toString() + const currentFileUri = this.getCurrentDocumentUri(document).toString() const currentLanguageId = document.languageId const viewPorts = Array.from(this.viewportsByDocumentUri.entries()) .map(([_, value]) => value) @@ -91,6 +116,13 @@ export class RecentViewPortRetriever implements vscode.Disposable, ContextRetrie return sortedViewPorts } + private getCurrentDocumentUri(document: vscode.TextDocument): vscode.Uri { + if (getCellIndexInActiveNotebookEditor(document) !== -1) { + return getActiveNotebookUri() ?? document.uri + } + return document.uri + } + public isSupportedForLanguageId(): boolean { return true } @@ -99,43 +131,91 @@ export class RecentViewPortRetriever implements vscode.Disposable, ContextRetrie if (this.activeTextEditor) { // Update the previous editor which was active before this one // Most of the property would remain same, but lastAccessTimestamp would be updated on the update - this.updateTrackedViewPort( - this.activeTextEditor.document, - this.activeTextEditor.visibleRanges[0], - this.activeTextEditor.document.languageId - ) + this.updateTrackedViewPort({ + uri: this.activeTextEditor.document.uri, + content: this.activeTextEditor.document.getText(this.activeTextEditor.visibleRanges[0]), + languageId: this.activeTextEditor.document.languageId, + startLine: this.activeTextEditor.visibleRanges.at(-1)?.start.line, + endLine: this.activeTextEditor.visibleRanges.at(-1)?.end.line, + }) } - this.activeTextEditor = editor - if (!editor?.visibleRanges?.[0]) { + if (!editor) { return } - this.updateTrackedViewPort(editor.document, editor.visibleRanges[0], editor.document.languageId) + this.updateTextEditor(editor, editor.visibleRanges) + } + + private onDidChangeActiveNotebookEditor(editor: vscode.NotebookEditor | undefined): void { + if (!editor?.notebook) { + return + } + this.updateNotebookEditor(editor, editor.visibleRanges) } private onDidChangeTextEditorVisibleRanges(event: vscode.TextEditorVisibleRangesChangeEvent): void { - const { textEditor, visibleRanges } = event + this.updateTextEditor(event.textEditor, event.visibleRanges) + } + + private onDidChangeNotebookEditorVisibleRanges( + event: vscode.NotebookEditorVisibleRangesChangeEvent + ) { + this.updateNotebookEditor(event.notebookEditor, event.visibleRanges) + } + + private updateTextEditor(editor: vscode.TextEditor, visibleRanges: readonly vscode.Range[]): void { if (visibleRanges.length === 0) { return } - this.updateTrackedViewPort(textEditor.document, visibleRanges[0], textEditor.document.languageId) + this.updateTrackedViewPort({ + uri: editor.document.uri, + content: editor.document.getText(visibleRanges?.at(-1)), + languageId: editor.document.languageId, + startLine: visibleRanges?.at(-1)?.start.line, + endLine: visibleRanges?.at(-1)?.end.line, + }) } - private updateTrackedViewPort( - document: vscode.TextDocument, - visibleRange: vscode.Range, - languageId: string + private updateNotebookEditor( + notebookEditor: vscode.NotebookEditor, + visibleRanges: readonly vscode.NotebookRange[] ): void { - if (document.uri.scheme !== 'file') { + if (!notebookEditor.notebook || visibleRanges.length === 0) { + return + } + const visibleCells = notebookEditor.notebook.getCells(visibleRanges?.at(-1)) + const content = getTextFromNotebookCells(notebookEditor.notebook, visibleCells).toString() + + this.updateTrackedViewPort({ + uri: notebookEditor.notebook.uri, + content, + languageId: getNotebookLanguageId(notebookEditor.notebook), + }) + } + + private updateTrackedViewPort(params: { + uri: vscode.Uri + content: string + languageId: string + startLine?: number + endLine?: number + }): void { + if (params.uri.scheme !== 'file') { return } const now = Date.now() - const key = document.uri.toString() + const key = params.uri.toString() this.viewportsByDocumentUri.set(key, { - uri: document.uri, - visibleRange, - languageId, + uri: params.uri, + content: params.content, + languageId: params.languageId, lastAccessTimestamp: now, + ...(params.startLine !== undefined || params.endLine !== undefined + ? { + startLine: params.startLine, + endLine: params.endLine, + } + : undefined), }) } diff --git a/vscode/src/completions/context/retrievers/tsc/tsc-retriever.ts b/vscode/src/completions/context/retrievers/tsc/tsc-retriever.ts index 1ded851fe00e..9e9743e9a0f3 100644 --- a/vscode/src/completions/context/retrievers/tsc/tsc-retriever.ts +++ b/vscode/src/completions/context/retrievers/tsc/tsc-retriever.ts @@ -536,6 +536,7 @@ class SymbolCollector { // Skip module declarations because they can be too large. // We still format them to queue the referenced types. const snippet: AutocompleteContextSnippet = { + type: 'symbol', identifier: RetrieverIdentifier.TscRetriever, symbol: sym.name, content, diff --git a/vscode/src/completions/model-helpers/__tests__/test-data.ts b/vscode/src/completions/model-helpers/__tests__/test-data.ts index 3afec237642d..0ef7b1519304 100644 --- a/vscode/src/completions/model-helpers/__tests__/test-data.ts +++ b/vscode/src/completions/model-helpers/__tests__/test-data.ts @@ -1,4 +1,4 @@ -import { testFileUri } from '@sourcegraph/cody-shared' +import { type AutocompleteContextSnippet, testFileUri } from '@sourcegraph/cody-shared' import { paramsWithInlinedCompletion } from '../../get-inline-completions-tests/helpers' @@ -18,8 +18,9 @@ export const completionParams = paramsWithInlinedCompletion( { documentUri: testFileUri('codebase/test.ts') } )! -export const contextSnippets = [ +export const contextSnippets: AutocompleteContextSnippet[] = [ { + type: 'file', identifier: 'jaccard-similarity', uri: testFileUri('codebase/context1.ts'), content: 'function contextSnippetOne() {}', @@ -27,6 +28,7 @@ export const contextSnippets = [ endLine: 2, }, { + type: 'file', identifier: 'jaccard-similarity', uri: testFileUri('codebase/context2.ts'), content: 'const contextSnippet2 = {}', @@ -34,6 +36,7 @@ export const contextSnippets = [ endLine: 2, }, { + type: 'symbol', identifier: 'jaccard-similarity', uri: testFileUri('codebase/context3.ts'), content: 'interface ContextParams {}', diff --git a/vscode/src/completions/test-helpers.ts b/vscode/src/completions/test-helpers.ts index 10c3cb6cbc22..d2b3011bed3a 100644 --- a/vscode/src/completions/test-helpers.ts +++ b/vscode/src/completions/test-helpers.ts @@ -1,6 +1,6 @@ import fs from 'node:fs' import dedent from 'dedent' -import type { Position as VSCodePosition, TextDocument as VSCodeTextDocument } from 'vscode' +import type * as vscode from 'vscode' import { TextDocument } from 'vscode-languageserver-textdocument' import { type CompletionResponse, testFileUri } from '@sourcegraph/cody-shared' @@ -39,11 +39,11 @@ export function document( text: string, languageId = 'typescript', uriString = testFileUri('test.ts').toString() -): VSCodeTextDocument { +): vscode.TextDocument { return wrapVSCodeTextDocument(TextDocument.create(uriString, languageId, 0, text)) } -export function documentFromFilePath(filePath: string, languageId = 'typescript'): VSCodeTextDocument { +export function documentFromFilePath(filePath: string, languageId = 'typescript'): vscode.TextDocument { return wrapVSCodeTextDocument( TextDocument.create( Uri.file(filePath).toString(), @@ -58,7 +58,7 @@ export function documentAndPosition( textWithCursor: string, languageId?: string, uriString?: string -): { document: VSCodeTextDocument; position: VSCodePosition } { +): { document: vscode.TextDocument; position: vscode.Position } { const { prefix, suffix, cursorIndex } = prefixAndSuffix(textWithCursor) const doc = document(prefix + suffix, languageId, uriString) @@ -80,3 +80,193 @@ export function prefixAndSuffix(textWithCursor: string): { return { prefix, suffix, cursorIndex } } + +//------------------------------------------- +// vscode.NotebookDocument mocks +//------------------------------------------- + +export function mockNotebookAndPosition({ + uri, + cells, + notebookType, +}: { + /** + * Notebook URI as a string, e.g. 'file://test.ipynb'. + */ + uri: string + /** + * Cells to create. + * - `kind`: NotebookCellKind.Code or NotebookCellKind.Markup + * - `text`: Cell content + * - `languageId`: The language ID for the cell document + */ + cells: { + kind: vscode.NotebookCellKind + text: string + languageId?: string + }[] + notebookType?: string +}): { notebookDoc: vscode.NotebookDocument; position: vscode.Position } { + const mockUri = Uri.parse(uri) + + const notebookDoc = new MockNotebookDocument({ + uri: mockUri, + notebookType: notebookType ?? 'mock-notebook', + cells: [], + }) + + const positionsWithCursorInNotebook: vscode.Position[] = [] + + const mockCells: vscode.NotebookCell[] = cells.map((cellData, index) => { + let cellText = cellData.text + let notebookCursorIndex = -1 + + if (cellData.text.includes(CURSOR_MARKER)) { + const { prefix, suffix, cursorIndex } = prefixAndSuffix(cellData.text) + cellText = prefix + suffix + notebookCursorIndex = cursorIndex + } + const cellTextDoc = wrapVSCodeTextDocument( + TextDocument.create( + `vscode-notebook-cell://mock/${index}`, + cellData.languageId ?? 'plaintext', + 0, // version + cellText + ) + ) + if (notebookCursorIndex !== -1) { + positionsWithCursorInNotebook.push(cellTextDoc.positionAt(notebookCursorIndex)) + } + + return new MockNotebookCell({ + index, + notebook: notebookDoc, + kind: cellData.kind, + document: cellTextDoc, + }) + }) + + if (positionsWithCursorInNotebook.length !== 1) { + throw new Error( + `Only one cell is currently supported to have a cursor position in notebook, received: ${positionsWithCursorInNotebook.length}` + ) + } + // Patch the internal cells array + // We do not create cells first to ensure they have a + // parent notebook reference. + ;(notebookDoc as any).cellsInternal = mockCells + notebookDoc.cellCount = mockCells.length + + return { notebookDoc, position: positionsWithCursorInNotebook[0] } +} + +class NotebookCellOutput implements vscode.NotebookCellOutput { + public readonly items: vscode.NotebookCellOutputItem[] + public readonly metadata: { [key: string]: any } + + constructor(items: vscode.NotebookCellOutputItem[], metadata: { [key: string]: any } = {}) { + this.items = items + this.metadata = metadata + } +} + +class MockNotebookCell implements vscode.NotebookCell { + public readonly index: number + public readonly notebook: vscode.NotebookDocument + public readonly kind: vscode.NotebookCellKind + public readonly document: vscode.TextDocument + public readonly metadata: { readonly [key: string]: any } + public readonly outputs: readonly NotebookCellOutput[] + public readonly executionSummary: vscode.NotebookCellExecutionSummary | undefined + + constructor({ + index, + notebook, + kind, + document, + metadata = {}, + outputs = [], + executionSummary, + }: { + index: number + notebook: vscode.NotebookDocument + kind: vscode.NotebookCellKind + document: vscode.TextDocument + metadata?: { readonly [key: string]: any } + outputs?: NotebookCellOutput[] + executionSummary?: vscode.NotebookCellExecutionSummary + }) { + this.index = index + this.notebook = notebook + this.kind = kind + this.document = document + this.metadata = metadata + this.outputs = outputs + this.executionSummary = executionSummary + } +} + +class MockNotebookDocument implements vscode.NotebookDocument { + public readonly uri: Uri + public readonly notebookType: string + public version: number + public isDirty: boolean + public isUntitled: boolean + public isClosed: boolean + public readonly metadata: { [key: string]: any } + public cellCount: number + private cellsInternal: vscode.NotebookCell[] + + constructor({ + uri, + notebookType = 'mock-notebook', + version = 1, + isDirty = false, + isUntitled = false, + isClosed = false, + metadata = {}, + cells = [], + }: { + uri: Uri + notebookType?: string + version?: number + isDirty?: boolean + isUntitled?: boolean + isClosed?: boolean + metadata?: { [key: string]: any } + cells?: vscode.NotebookCell[] + }) { + this.uri = uri + this.notebookType = notebookType + this.version = version + this.isDirty = isDirty + this.isUntitled = isUntitled + this.isClosed = isClosed + this.metadata = metadata + this.cellsInternal = cells + this.cellCount = cells.length + } + + public cellAt(index: number): vscode.NotebookCell { + const cell = this.cellsInternal[index] + if (!cell) { + throw new Error(`No cell found at index ${index}`) + } + return cell + } + + public getCells(range?: vscode.NotebookRange): vscode.NotebookCell[] { + if (!range) { + return this.cellsInternal + } + const start = Math.max(0, range.start) + const end = Math.min(this.cellsInternal.length, range.end) + return this.cellsInternal.slice(start, end) + } + + public async save(): Promise { + // Simulate saving + this.isDirty = false + return true + } +} diff --git a/vscode/src/graph/lsp/symbol-context-snippets.ts b/vscode/src/graph/lsp/symbol-context-snippets.ts index 9c3dfb86357f..ae058f313cc2 100644 --- a/vscode/src/graph/lsp/symbol-context-snippets.ts +++ b/vscode/src/graph/lsp/symbol-context-snippets.ts @@ -127,6 +127,7 @@ async function getSnippetForLocationGetter( } const symbolContextSnippet = { + type: 'symbol', identifier: RetrieverIdentifier.LspLightRetriever, key: `${definitionUri}::${definitionRange.start.line}:${definitionRange.start.character}`, uri: definitionUri, diff --git a/vscode/src/testutils/mocks.ts b/vscode/src/testutils/mocks.ts index 6dab60d39596..d22b053a60c4 100644 --- a/vscode/src/testutils/mocks.ts +++ b/vscode/src/testutils/mocks.ts @@ -229,6 +229,10 @@ export class CodeActionKind { constructor(public readonly value: string) {} } +export enum NotebookCellKind { + Markup = 1, + Code = 2, +} // biome-ignore lint/complexity/noStaticOnlyClass: mock export class QuickInputButtons { public static readonly Back: vscode_types.QuickInputButton = { @@ -793,6 +797,7 @@ export const vsCodeMocks = { showErrorMessage(message: string) { console.error(message) }, + activeNotebookEditor: undefined, activeTextEditor: { document: { uri: { scheme: 'not-cody' } }, options: { tabSize: 4 }, @@ -867,6 +872,7 @@ export const vsCodeMocks = { FoldingRange, FoldingRangeKind, CodeActionKind, + NotebookCellKind, DiagnosticSeverity, ViewColumn, TextDocumentChangeReason,