Skip to content

Commit e1d39bb

Browse files
RomneyDaRomneyDa
authored
Merge pull request #7029 from continuedev/jacob/con-3283
feat: set capabilities for next edit
2 parents 50a7b8c + 6f65295 commit e1d39bb

File tree

11 files changed

+207
-56
lines changed

11 files changed

+207
-56
lines changed

‎core/config/yaml/models.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ async function modelConfigToBaseLLM({
6868
capabilities: {
6969
tools: model.capabilities?.includes("tool_use"),
7070
uploadImage: model.capabilities?.includes("image_input"),
71+
nextEdit: model.capabilities?.includes("next_edit"),
7172
},
7273
autocompleteOptions: model.autocompleteOptions,
7374
};

‎core/index.d.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,6 +1141,7 @@ export interface BaseCompletionOptions {
11411141
export interface ModelCapability {
11421142
uploadImage?: boolean;
11431143
tools?: boolean;
1144+
nextEdit?: boolean;
11441145
}
11451146

11461147
export interface ModelDescription {

‎core/llm/autodetect.ts‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import { ChatMessage, ModelCapability, TemplateType } from "../index.js";
2+
import { NEXT_EDIT_MODELS } from "./constants.js";
23

34
import {
45
anthropicTemplateMessages,
@@ -173,6 +174,34 @@ function isProviderHandlesTemplatingOrNoTemplateTypeRequired(
173174
);
174175
}
175176

177+
// NOTE: When updating this list,
178+
// update core/nextEdit/templating/NextEditPromptEngine.ts as well.
179+
const MODEL_SUPPORTS_NEXT_EDIT: string[] = [
180+
NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT,
181+
NEXT_EDIT_MODELS.MODEL_1,
182+
];
183+
184+
function modelSupportsNextEdit(
185+
capabilities: ModelCapability | undefined,
186+
model: string,
187+
title: string | undefined,
188+
): boolean {
189+
if (capabilities?.nextEdit !== undefined) {
190+
return capabilities.nextEdit;
191+
}
192+
193+
const lower = model.toLowerCase();
194+
if (
195+
MODEL_SUPPORTS_NEXT_EDIT.some(
196+
(modelName) => lower.includes(modelName) || title?.includes(modelName),
197+
)
198+
) {
199+
return true;
200+
}
201+
202+
return false;
203+
}
204+
176205
function autodetectTemplateType(model: string): TemplateType | undefined {
177206
const lower = model.toLowerCase();
178207

@@ -391,4 +420,5 @@ export {
391420
autodetectTemplateType,
392421
llmCanGenerateInParallel,
393422
modelSupportsImages,
423+
modelSupportsNextEdit,
394424
};

‎core/llm/autodetect.vitest.ts‎

Lines changed: 136 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
1-
import { expect, test } from "vitest";
2-
import { autodetectTemplateType } from "./autodetect";
1+
import { describe, expect, it, test } from "vitest";
2+
import { autodetectTemplateType, modelSupportsNextEdit } from "./autodetect";
33

44
test("autodetectTemplateType returns 'codellama-70b' for CodeLlama 70B models", () => {
55
expect(autodetectTemplateType("codellama-70b")).toBe("codellama-70b");
@@ -201,3 +201,137 @@ test("autodetectTemplateType handles models with mixed keywords", () => {
201201
expect(autodetectTemplateType("gpt-llama")).toBe(undefined); // gpt comes first, returns undefined
202202
expect(autodetectTemplateType("claude-llama")).toBe("llama2"); // llama comes first, returns llama2
203203
});
204+
205+
describe("modelSupportsNextEdit", () => {
206+
describe("when capabilities.nextEdit is defined", () => {
207+
it("should return true when capabilities.nextEdit is true", () => {
208+
expect(
209+
modelSupportsNextEdit(
210+
{
211+
nextEdit: true,
212+
},
213+
"any-model",
214+
"Any Title",
215+
),
216+
).toBe(true);
217+
});
218+
219+
it("should return false when capabilities.nextEdit is false", () => {
220+
expect(
221+
modelSupportsNextEdit(
222+
{ nextEdit: false },
223+
"mercury-coder-nextedit",
224+
"Mercury Coder",
225+
),
226+
).toBe(false);
227+
});
228+
229+
it("should prioritize capabilities over model name matching", () => {
230+
// Even though model name matches, capabilities should take precedence.
231+
expect(
232+
modelSupportsNextEdit(
233+
{ nextEdit: false },
234+
"mercury-coder-nextedit",
235+
"Mercury Coder",
236+
),
237+
).toBe(false);
238+
});
239+
});
240+
241+
describe("when capabilities.nextEdit is undefined", () => {
242+
it("should return true for mercury-coder-nextedit model (case insensitive)", () => {
243+
expect(
244+
modelSupportsNextEdit(undefined, "Mercury-Coder-Nextedit", undefined),
245+
).toBe(true);
246+
});
247+
248+
it("should return true for model-1", () => {
249+
expect(modelSupportsNextEdit(undefined, "model-1", undefined)).toBe(true);
250+
});
251+
252+
it("should return true when model contains supported model name as substring", () => {
253+
expect(
254+
modelSupportsNextEdit(
255+
undefined,
256+
"provider/mercury-coder-nextedit-v2",
257+
undefined,
258+
),
259+
).toBe(true);
260+
});
261+
262+
it("should return true when title contains supported model name", () => {
263+
expect(
264+
modelSupportsNextEdit(
265+
undefined,
266+
"some-model",
267+
"This is mercury-coder-nextedit model",
268+
),
269+
).toBe(true);
270+
});
271+
272+
it("should return true when title contains model-1", () => {
273+
expect(
274+
modelSupportsNextEdit(undefined, "some-model", "model-1 deployment"),
275+
).toBe(true);
276+
});
277+
278+
it("should return true for unsupported models that have capabilities explicitly set to true", () => {
279+
expect(
280+
modelSupportsNextEdit(
281+
{
282+
nextEdit: true,
283+
},
284+
"gpt-4",
285+
"GPT-4 Model",
286+
),
287+
).toBe(true);
288+
});
289+
290+
it("should return false for unsupported models", () => {
291+
expect(modelSupportsNextEdit(undefined, "gpt-4", "GPT-4 Model")).toBe(
292+
false,
293+
);
294+
});
295+
296+
it("should return false when model and title are both undefined/null", () => {
297+
expect(modelSupportsNextEdit(undefined, "", undefined)).toBe(false);
298+
});
299+
300+
it("should return false when model and title do not contain supported names", () => {
301+
expect(
302+
modelSupportsNextEdit(undefined, "claude-3", "Claude 3 Sonnet"),
303+
).toBe(false);
304+
});
305+
});
306+
307+
describe("edge cases", () => {
308+
it("should handle empty strings", () => {
309+
expect(modelSupportsNextEdit(undefined, "", "")).toBe(false);
310+
});
311+
312+
it("should handle undefined title gracefully", () => {
313+
expect(
314+
modelSupportsNextEdit(undefined, "mercury-coder-nextedit", undefined),
315+
).toBe(true);
316+
});
317+
318+
it("should handle case sensitivity correctly", () => {
319+
expect(
320+
modelSupportsNextEdit(undefined, "MERCURY-CODER-NEXTEDIT", "MODEL-1"),
321+
).toBe(true);
322+
});
323+
324+
it("should handle capabilities with other properties", () => {
325+
expect(
326+
modelSupportsNextEdit(
327+
{
328+
nextEdit: true,
329+
uploadImage: false,
330+
},
331+
"unsupported-model",
332+
undefined,
333+
),
334+
).toBe(true);
335+
});
336+
});
337+
});

‎core/llm/constants.ts‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@ export enum LLMConfigurationStatuses {
1919
MISSING_ENV_SECRET = "missing-env-secret",
2020
}
2121

22+
export enum NEXT_EDIT_MODELS {
23+
MERCURY_CODER_NEXTEDIT = "mercury-coder-nextedit",
24+
MODEL_1 = "model-1",
25+
}
26+
2227
export {
2328
DEFAULT_ARGS,
2429
DEFAULT_CONTEXT_LENGTH,

‎core/nextEdit/NextEditProvider.ts‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import AutocompleteLruCache from "../autocomplete/util/AutocompleteLruCache.js";
3030
import { HelperVars } from "../autocomplete/util/HelperVars.js";
3131
import { AutocompleteInput } from "../autocomplete/util/types.js";
3232
import { myersDiff } from "../diff/myers.js";
33+
import { modelSupportsNextEdit } from "../llm/autodetect.js";
3334
import { countTokens } from "../llm/countTokens.js";
3435
import { localPathOrUriToPath } from "../util/pathToUri.js";
3536
import { replaceEscapedCharacters } from "../util/text.js";
@@ -341,6 +342,14 @@ export class NextEditProvider {
341342
return { token, startTime, helper: undefined };
342343
}
343344

345+
// In vscode, this check is done in extensions/vscode/src/extension/VsCodeExtension.ts.
346+
// For other editors, this check should be done in their respective config reloaders.
347+
// This is left for a final check.
348+
if (!modelSupportsNextEdit(llm.capabilities, llm.model, llm.title)) {
349+
console.error(`${llm.model} is not capable of next edit.`);
350+
return { token, startTime, helper: undefined };
351+
}
352+
344353
if (llm.promptTemplates?.autocomplete) {
345354
options.template = llm.promptTemplates.autocomplete as string;
346355
}

‎core/nextEdit/templating/NextEditPromptEngine.ts‎

Lines changed: 5 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -30,29 +30,22 @@ import {
3030
insertCursorToken,
3131
insertEditableRegionTokensWithStaticRange,
3232
} from "./utils";
33+
import { NEXT_EDIT_MODELS } from "../../llm/constants";
3334

3435
type TemplateRenderer = (vars: TemplateVars) => string;
3536

36-
export type NextEditModelName =
37-
| "mercury-coder-nextedit"
38-
| "model-1"
39-
| "this field is not used";
40-
41-
const NEXT_EDIT_MODEL_TEMPLATES: Record<NextEditModelName, NextEditTemplate> = {
37+
const NEXT_EDIT_MODEL_TEMPLATES: Record<NEXT_EDIT_MODELS, NextEditTemplate> = {
4238
"mercury-coder-nextedit": {
4339
template: `${MERCURY_RECENTLY_VIEWED_CODE_SNIPPETS_OPEN}\n{{{recentlyViewedCodeSnippets}}}\n${MERCURY_RECENTLY_VIEWED_CODE_SNIPPETS_CLOSE}\n\n${MERCURY_CURRENT_FILE_CONTENT_OPEN}\n{{{currentFileContent}}}\n${MERCURY_CURRENT_FILE_CONTENT_CLOSE}\n\n${MERCURY_EDIT_DIFF_HISTORY_OPEN}\n{{{editDiffHistory}}}\n${MERCURY_EDIT_DIFF_HISTORY_CLOSE}\n\nThe developer was working on a section of code within the tags \`<|code_to_edit|>\` in the file located at {{{currentFilePath}}}.\nUsing the given \`recently_viewed_code_snippets\`, \`current_file_content\`, \`edit_diff_history\`, and the cursor position marked as \`<|cursor|>\`, please continue the developer's work. Update the \`code_to_edit\` section by predicting and completing the changes they would have made next. Provide the revised code that was between the \`<|code_to_edit|>\` and \`<|/code_to_edit|>\` tags, including the tags themselves.`,
4440
},
4541
"model-1": {
4642
template:
4743
"### User Edits:\n\n{{{userEdits}}}\n\n### User Excerpts:\n\n```{{{languageShorthand}}}\n{{{userExcerpts}}}```",
4844
},
49-
"this field is not used": {
50-
template: "NEXT_EDIT",
51-
},
5245
};
5346

5447
function templateRendererOfModel(
55-
modelName: NextEditModelName,
48+
modelName: NEXT_EDIT_MODELS,
5649
): TemplateRenderer {
5750
let template = NEXT_EDIT_MODEL_TEMPLATES[modelName];
5851
if (!template) {
@@ -70,18 +63,7 @@ export async function renderPrompt(
7063
helper: HelperVars,
7164
ctx: any,
7265
): Promise<PromptMetadata> {
73-
let modelName = helper.modelName as NextEditModelName;
74-
75-
if (modelName === "this field is not used") {
76-
return {
77-
prompt: {
78-
role: "user",
79-
content: "NEXT_EDIT",
80-
},
81-
userEdits: "",
82-
userExcerpts: helper.fileContents,
83-
};
84-
}
66+
let modelName = helper.modelName as NEXT_EDIT_MODELS;
8567

8668
// Validate that the modelName is actually a supported model.
8769
if (!Object.keys(NEXT_EDIT_MODEL_TEMPLATES).includes(modelName)) {
@@ -91,7 +73,7 @@ export async function renderPrompt(
9173
);
9274

9375
if (matchingModel) {
94-
modelName = matchingModel as NextEditModelName;
76+
modelName = matchingModel as NEXT_EDIT_MODELS;
9577
} else {
9678
throw new Error(
9779
`${helper.modelName} is not yet supported for next edit.`,

‎core/nextEdit/templating/NextEditPromptEngine.vitest.ts‎

Lines changed: 5 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,11 @@ import {
1414
MODEL_1_USER_CURSOR_IS_HERE_TOKEN,
1515
} from "../constants";
1616
import {
17-
NextEditModelName,
1817
renderDefaultSystemPrompt,
1918
renderDefaultUserPrompt,
2019
renderPrompt,
2120
} from "./NextEditPromptEngine";
21+
import { NEXT_EDIT_MODELS } from "../../llm/constants";
2222

2323
describe("NextEditPromptEngine", () => {
2424
describe("renderPrompt", () => {
@@ -30,25 +30,19 @@ describe("NextEditPromptEngine", () => {
3030

3131
beforeEach(() => {
3232
mercuryHelper = {
33-
modelName: "inception/mercury-coder-nextedit" as NextEditModelName,
33+
modelName: "inception/mercury-coder-nextedit" as NEXT_EDIT_MODELS,
3434
fileContents: "function test() {\n const a = 1;\n return a;\n}",
3535
pos: { line: 1, character: 12 } as Position,
3636
lang: { name: "typescript" },
3737
} as HelperVars;
3838
model1Helper = {
39-
modelName: "continuedev/model-1" as NextEditModelName,
40-
fileContents: "function test() {\n const a = 1;\n return a;\n}",
41-
pos: { line: 1, character: 12 } as Position,
42-
lang: { name: "typescript" },
43-
} as HelperVars;
44-
testHelper = {
45-
modelName: "this field is not used" as NextEditModelName,
39+
modelName: "continuedev/model-1" as NEXT_EDIT_MODELS,
4640
fileContents: "function test() {\n const a = 1;\n return a;\n}",
4741
pos: { line: 1, character: 12 } as Position,
4842
lang: { name: "typescript" },
4943
} as HelperVars;
5044
unsupportedHelper = {
51-
modelName: "mistral/codestral" as NextEditModelName,
45+
modelName: "mistral/codestral" as NEXT_EDIT_MODELS,
5246
fileContents: "function test() {\n const a = 1;\n return a;\n}",
5347
pos: { line: 1, character: 12 } as Position,
5448
lang: { name: "typescript" },
@@ -99,16 +93,6 @@ describe("NextEditPromptEngine", () => {
9993
expect(result.userExcerpts).toContain(MODEL_1_USER_CURSOR_IS_HERE_TOKEN);
10094
});
10195

102-
it("should handle 'this field is not used' model name", async () => {
103-
const result = await renderPrompt(testHelper, ctx);
104-
105-
expect(result).toHaveProperty("prompt");
106-
expect(result.prompt.role).toBe("user");
107-
expect(result.prompt.content).toBe("NEXT_EDIT");
108-
expect(result.userEdits).toBe("");
109-
expect(result.userExcerpts).toBe(testHelper.fileContents);
110-
});
111-
11296
it("should throw error for unsupported model name", async () => {
11397
await expect(renderPrompt(unsupportedHelper, ctx)).rejects.toThrow(
11498
"mistral/codestral is not yet supported for next edit.",
@@ -209,7 +193,7 @@ describe("NextEditPromptEngine", () => {
209193
describe("insertTokens", () => {
210194
it("should correctly insert cursor and editable region tokens", () => {
211195
const mercuryHelper = {
212-
modelName: "inception/mercury-coder-nextedit" as NextEditModelName,
196+
modelName: "inception/mercury-coder-nextedit" as NEXT_EDIT_MODELS,
213197
fileContents: "function test() {\n const a = 1;\n return a;\n}",
214198
pos: { line: 1, character: 12 } as Position,
215199
lang: { name: "typescript" },

‎core/nextEdit/utils.ts‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,13 @@
1-
type NextEditModelName = "mercury-coder-nextedit";
1+
export function isNextEditTest(): boolean {
2+
const enabled = process.env.NEXT_EDIT_TEST_ENABLED;
23

3-
export function isModelCapableOfNextEdit(modelName: string): boolean {
4-
// In test mode, we can control whether next edit is enabled via environment variable.
5-
if (process.env.NEXT_EDIT_TEST_ENABLED === "false") {
4+
if (enabled === "false") {
65
return false;
76
}
87

9-
if (process.env.NEXT_EDIT_TEST_ENABLED === "true") {
8+
if (enabled === "true") {
109
return true;
1110
}
1211

13-
const supportedModels: NextEditModelName[] = ["mercury-coder-nextedit"];
14-
return supportedModels.some((supported) => modelName.includes(supported));
12+
return false;
1513
}

0 commit comments

Comments
 (0)