Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 84 additions & 15 deletions src/extension/prompts/node/inline/inlineChat2Prompt.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -4,11 +4,12 @@
*--------------------------------------------------------------------------------------------*/

import { AssistantMessage, PromptElement, PromptElementProps, PromptReference, PromptSizing, SystemMessage, ToolCall, ToolMessage, useKeepWith, UserMessage } from '@vscode/prompt-tsx';
import type { ExtendedLanguageModelToolResult } from 'vscode';
import { ChatResponsePart } from '@vscode/prompt-tsx/dist/base/vscodeTypes';
import type { CancellationToken, ExtendedLanguageModelToolResult, Position, Progress } from 'vscode';
import { TextDocumentSnapshot } from '../../../../platform/editing/common/textDocumentSnapshot';
import { CacheType } from '../../../../platform/endpoint/common/endpointTypes';
import { IPromptPathRepresentationService } from '../../../../platform/prompts/common/promptPathRepresentationService';
import { ChatRequest, ChatRequestEditorData } from '../../../../vscodeTypes';
import { ChatRequest, ChatRequestEditorData, Range } from '../../../../vscodeTypes';
import { ChatVariablesCollection } from '../../../prompt/common/chatVariablesCollection';
import { IToolCall } from '../../../prompt/common/intents';
import { CopilotIdentityRules } from '../base/copilotIdentity';
Expand Down Expand Up @@ -41,15 +42,7 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {

const snapshotAtRequest = this.props.snapshotAtRequest;

// the full lines of the selection
// TODO@jrieken
// * if the selection is empty and if the line with the selection is empty we could hint to add code and
// generally with empty selections we could allow the model to be a bit more creative
// * use the true selected text (now we extend to full lines)
const selectedLines = snapshotAtRequest.getText(this.props.data.selection.with({
start: this.props.data.selection.start.with({ character: 0 }),
end: this.props.data.selection.end.with({ character: Number.MAX_SAFE_INTEGER }),
}));
const selection = this.props.data.selection;

const variables = new ChatVariablesCollection(this.props.request.references);
const filepath = this._promptPathRepresentationService.getFilePath(snapshotAtRequest.uri);
Expand All @@ -74,9 +67,10 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {
<Tag name='file'>
<CodeBlock includeFilepath={false} languageId={snapshotAtRequest.languageId} uri={snapshotAtRequest.uri} references={[new PromptReference(snapshotAtRequest.uri, undefined, undefined)]} code={snapshotAtRequest.getText()} />
</Tag>
<Tag name='file-selection'>
<CodeBlock includeFilepath={false} languageId={snapshotAtRequest.languageId} uri={snapshotAtRequest.uri} references={[new PromptReference(snapshotAtRequest.uri, undefined, undefined)]} code={selectedLines} />
</Tag>
{selection.isEmpty
? <FileContextElement snapshot={snapshotAtRequest} position={selection.start} />
: <FileSelectionElement snapshot={snapshotAtRequest} selection={selection} />
}
<ChatVariables flexGrow={3} priority={898} chatVariables={variables} useFixCookbook={true} />
<Tag name='reminder'>
If there is a user selection, focus on it, and try to make changes to the selected code and its context.<br />
Expand All @@ -98,7 +92,82 @@ export class InlineChat2Prompt extends PromptElement<InlineChat2PromptProps> {
}
}

export type EditAttemptsElementProps = PromptElementProps<{

export type FileContextElementProps = PromptElementProps<{
snapshot: TextDocumentSnapshot;
position: Position;
}>;

export class FileContextElement extends PromptElement<FileContextElementProps> {

override render(state: void, sizing: PromptSizing, _progress?: Progress<ChatResponsePart>, _token?: CancellationToken) {

let startLine = this.props.position.line;
let endLine = this.props.position.line;
let n = 0;
let seenNonEmpty = false;
while (startLine > 0) {
seenNonEmpty = seenNonEmpty || !this.props.snapshot.lineAt(startLine).isEmptyOrWhitespace;
startLine--;
n++;
if (n >= 3 && seenNonEmpty) {
break;
}
}
n = 0;
seenNonEmpty = false;
while (endLine < this.props.snapshot.lineCount - 1) {
seenNonEmpty = seenNonEmpty || !this.props.snapshot.lineAt(endLine).isEmptyOrWhitespace;
endLine++;
n++;
if (n >= 3 && seenNonEmpty) {
break;
}
}

const textBefore = this.props.snapshot.getText(new Range(this.props.position.with({ line: startLine, character: 0 }), this.props.position));
const textAfter = this.props.snapshot.getText(new Range(this.props.position, this.props.position.with({ line: endLine, character: Number.MAX_SAFE_INTEGER })));

const code = `${textBefore}$CURSOR$${textAfter}`;

return <>
<Tag name='file-cursor-context'>
<CodeBlock includeFilepath={false} languageId={this.props.snapshot.languageId} uri={this.props.snapshot.uri} references={[new PromptReference(this.props.snapshot.uri, undefined, undefined)]} code={code} />
</Tag>
</>;
}
}


export type FileSelectionElementProps = PromptElementProps<{
snapshot: TextDocumentSnapshot;
selection: Range;
}>;

export class FileSelectionElement extends PromptElement<FileSelectionElementProps> {

override render(state: void, sizing: PromptSizing, progress?: Progress<ChatResponsePart>, token?: CancellationToken) {


// the full lines of the selection
// TODO@jrieken
// * use the true selected text (now we extend to full lines)

const selectedLines = this.props.snapshot.getText(this.props.selection.with({
start: this.props.selection.start.with({ character: 0 }),
end: this.props.selection.end.with({ character: Number.MAX_SAFE_INTEGER }),
}));

return <>
<Tag name='file-selection'>
<CodeBlock includeFilepath={false} languageId={this.props.snapshot.languageId} uri={this.props.snapshot.uri} references={[new PromptReference(this.props.snapshot.uri, undefined, undefined)]} code={selectedLines} />
</Tag>
</>;
}
}


type EditAttemptsElementProps = PromptElementProps<{
editAttempts: [IToolCall, ExtendedLanguageModelToolResult][];
data: ChatRequestEditorData;
documentVersionAtRequest: number;
Expand Down
243 changes: 243 additions & 0 deletions src/extension/prompts/node/inline/test/inlineChat2Prompt.spec.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,243 @@
/*---------------------------------------------------------------------------------------------
* Copyright (c) Microsoft Corporation. All rights reserved.
* Licensed under the MIT License. See License.txt in the project root for license information.
*--------------------------------------------------------------------------------------------*/

import { expect, suite, test } from 'vitest';
import { TextDocumentSnapshot } from '../../../../../platform/editing/common/textDocumentSnapshot';
import { createTextDocumentData } from '../../../../../util/common/test/shims/textDocument';
import { URI } from '../../../../../util/vs/base/common/uri';
import { Position, Range } from '../../../../../vscodeTypes';
import { FileContextElement, FileSelectionElement } from '../inlineChat2Prompt';

function createSnapshot(content: string, languageId: string = 'typescript'): TextDocumentSnapshot {
const uri = URI.file('/workspace/file.ts');
const docData = createTextDocumentData(uri, content, languageId);
return TextDocumentSnapshot.create(docData.document);
}

suite('FileContextElement', () => {

test('cursor at the beginning of the file', async () => {
const content = `line 1
line 2
line 3
line 4
line 5`;
const snapshot = createSnapshot(content);
const position = new Position(0, 0);

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('$CURSOR$');
expect(output).toContain('line 1');
expect(output).toContain('line 2');
expect(output).toContain('line 3');
});

test('cursor in the middle of a file', async () => {
const content = `line 1
line 2
line 3
line 4
line 5
line 6
line 7`;
const snapshot = createSnapshot(content);
const position = new Position(3, 2); // after "li" in "line 4"

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('$CURSOR$');
// Should include lines before and after cursor
expect(output).toContain('line 2');
expect(output).toContain('line 3');
// Cursor position (3, 2) splits "line 4" into "li" + "$CURSOR$" + "ne 4"
expect(output).toContain('li$CURSOR$ne 4');
expect(output).toContain('line 5');
expect(output).toContain('line 6');
});

test('cursor at the end of file', async () => {
const content = `line 1
line 2
line 3
line 4
line 5`;
const snapshot = createSnapshot(content);
const position = new Position(4, 6); // end of "line 5"

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('$CURSOR$');
expect(output).toContain('line 3');
expect(output).toContain('line 4');
expect(output).toContain('line 5');
});

test('cursor with empty lines - includes extra lines until non-empty', async () => {
const content = `

line 3
line 4

`;
const snapshot = createSnapshot(content);
const position = new Position(2, 0); // start of "line 3"

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('$CURSOR$');
expect(output).toContain('line 3');
expect(output).toContain('line 4');
});

test('single line file', async () => {
const content = `only one line`;
const snapshot = createSnapshot(content);
const position = new Position(0, 5); // middle of line

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('only $CURSOR$one line');
});

test('cursor position splits text correctly', async () => {
const content = `hello world`;
const snapshot = createSnapshot(content);
const position = new Position(0, 6); // after "hello "

const element = new FileContextElement({ snapshot, position });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('hello $CURSOR$world');
});
});

suite('FileSelectionElement', () => {

test('single line selection', async () => {
const content = `line 1
line 2
line 3
line 4
line 5`;
const snapshot = createSnapshot(content);
const selection = new Range(1, 0, 1, 6); // "line 2"

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('line 2');
expect(output).not.toContain('line 1');
expect(output).not.toContain('line 3');
});

test('multi-line selection', async () => {
const content = `line 1
line 2
line 3
line 4
line 5`;
const snapshot = createSnapshot(content);
const selection = new Range(1, 0, 3, 6); // "line 2" through "line 4"

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('line 2');
expect(output).toContain('line 3');
expect(output).toContain('line 4');
expect(output).not.toContain('line 1');
expect(output).not.toContain('line 5');
});

test('partial line selection extends to full lines', async () => {
const content = `line 1
line 2
line 3`;
const snapshot = createSnapshot(content);
// Select from middle of line 2 to middle of line 2 (partial)
const selection = new Range(1, 2, 1, 4);

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
// Should include the full line, not just "ne"
expect(output).toContain('line 2');
});

test('selection at start of file', async () => {
const content = `line 1
line 2
line 3`;
const snapshot = createSnapshot(content);
const selection = new Range(0, 0, 0, 6);

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('line 1');
expect(output).not.toContain('line 2');
});

test('selection at end of file', async () => {
const content = `line 1
line 2
line 3`;
const snapshot = createSnapshot(content);
const selection = new Range(2, 0, 2, 6);

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('line 3');
expect(output).not.toContain('line 2');
});

test('selection spanning partial lines extends to full lines', async () => {
const content = `first line here
second line here
third line here`;
const snapshot = createSnapshot(content);
// Select from middle of "first" to middle of "second"
const selection = new Range(0, 6, 1, 7);

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
// Should include full lines
expect(output).toContain('first line here');
expect(output).toContain('second line here');
expect(output).not.toContain('third line here');
});

test('preserves language id for code block', async () => {
const content = `const x = 1;`;
const snapshot = createSnapshot(content, 'javascript');
const selection = new Range(0, 0, 0, 12);

const element = new FileSelectionElement({ snapshot, selection });
const rendered = await element.render(undefined, { tokenBudget: 1000, countTokens: () => Promise.resolve(0), endpoint: {} as any });

const output = typeof rendered === 'string' ? rendered : JSON.stringify(rendered) ?? '';
expect(output).toContain('javascript');
});
});
Loading