Merge pull request #7232 from continuedev/jacob/con-3563

sestinj · web-flow · commit e8dd689aeb20 · 2025-08-19T11:28:16.000-07:00
feat: update openai-adapter to support Inception models
diff --git a/core/llm/autodetect.ts b/core/llm/autodetect.ts
@@ -177,6 +177,7 @@ function isProviderHandlesTemplatingOrNoTemplateTypeRequired(
 // NOTE: When updating this list,
 // update core/nextEdit/templating/NextEditPromptEngine.ts as well.
 const MODEL_SUPPORTS_NEXT_EDIT: string[] = [
+  NEXT_EDIT_MODELS.MERCURY_CODER,
   NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT,
   NEXT_EDIT_MODELS.INSTINCT,
 ];
diff --git a/core/llm/constants.ts b/core/llm/constants.ts
@@ -20,6 +20,7 @@ export enum LLMConfigurationStatuses {
 }
 
 export enum NEXT_EDIT_MODELS {
+  MERCURY_CODER = "mercury-coder",
   MERCURY_CODER_NEXTEDIT = "mercury-coder-nextedit",
   INSTINCT = "instinct",
 }
diff --git a/core/llm/llms/Inception.ts b/core/llm/llms/Inception.ts
@@ -1,6 +1,8 @@
 import { streamSse } from "@continuedev/fetch";
-import { CompletionOptions, LLMOptions } from "../../index.js";
+import { ChatMessage, CompletionOptions, LLMOptions } from "../../index.js";
 
+import { ChatCompletionCreateParams } from "@continuedev/openai-adapters";
+import { UNIQUE_TOKEN } from "../../nextEdit/constants.js";
 import OpenAI from "./OpenAI.js";
 
 /**
@@ -32,6 +34,26 @@ class Inception extends OpenAI {
     return true;
   }
 
+  // It seems like this should be inherited automatically from the parent OpenAI class, but it sometimes doesn't.
+  // protected useOpenAIAdapterFor: (LlmApiRequestType | "*")[] = [
+  //   "chat",
+  //   "embed",
+  //   "list",
+  //   "rerank",
+  //   "streamChat",
+  //   "streamFim",
+  // ];
+
+  protected modifyChatBody(
+    body: ChatCompletionCreateParams,
+  ): ChatCompletionCreateParams {
+    const hasNextEditCapability = this.capabilities?.nextEdit ?? false;
+
+    // Add the nextEdit parameter for Inception-specific routing.
+    (body as any).nextEdit = hasNextEditCapability;
+    return body;
+  }
+
   async *_streamFim(
     prefix: string,
     suffix: string,
@@ -67,6 +89,79 @@ class Inception extends OpenAI {
       yield chunk.choices[0].text;
     }
   }
+
+  protected async *_streamChat(
+    messages: ChatMessage[],
+    signal: AbortSignal,
+    options: CompletionOptions,
+  ): AsyncGenerator<ChatMessage> {
+    if (this.isNextEdit(messages)) {
+      messages = this.removeNextEditToken(messages);
+
+      // Use edit/completions endpoint.
+      const endpoint = new URL("edit/completions", this.apiBase);
+
+      const resp = await this.fetch(endpoint, {
+        method: "POST",
+        body: JSON.stringify({
+          model: options.model,
+          messages: messages,
+          max_tokens: options.maxTokens,
+          temperature: options.temperature,
+          top_p: options.topP,
+          frequency_penalty: options.frequencyPenalty,
+          presence_penalty: options.presencePenalty,
+          stop: options.stop,
+          stream: true,
+        }),
+        headers: {
+          "Content-Type": "application/json",
+          Accept: "application/json",
+          Authorization: `Bearer ${this.apiKey}`,
+        },
+        signal,
+      });
+
+      for await (const chunk of streamSse(resp)) {
+        if (chunk.choices?.[0]?.delta?.content) {
+          yield {
+            role: "assistant",
+            content: chunk.choices[0].delta.content,
+          };
+        }
+      }
+    } else {
+      // Use regular chat/completions endpoint - call parent OpenAI implementation.
+      yield* super._streamChat(messages, signal, options);
+    }
+  }
+
+  private isNextEdit(messages: ChatMessage[]): boolean {
+    // Check if any message contains the unique next edit token.
+    return messages.some(
+      (message) =>
+        typeof message.content === "string" &&
+        message.content.endsWith(UNIQUE_TOKEN),
+    );
+  }
+
+  private removeNextEditToken(messages: ChatMessage[]): ChatMessage[] {
+    const lastMessage = messages[messages.length - 1];
+
+    if (
+      typeof lastMessage?.content === "string" &&
+      lastMessage.content.endsWith(UNIQUE_TOKEN)
+    ) {
+      const cleanedMessages = [...messages];
+      cleanedMessages[cleanedMessages.length - 1] = {
+        ...lastMessage,
+        content: lastMessage.content.slice(0, -UNIQUE_TOKEN.length),
+      };
+      return cleanedMessages;
+    }
+
+    return messages;
+  }
 }
 
 export default Inception;
diff --git a/core/nextEdit/NextEditProvider.ts b/core/nextEdit/NextEditProvider.ts
@@ -31,13 +31,14 @@ import { HelperVars } from "../autocomplete/util/HelperVars.js";
 import { AutocompleteInput } from "../autocomplete/util/types.js";
 import { myersDiff } from "../diff/myers.js";
 import { modelSupportsNextEdit } from "../llm/autodetect.js";
+import { NEXT_EDIT_MODELS } from "../llm/constants.js";
 import { countTokens } from "../llm/countTokens.js";
 import { localPathOrUriToPath } from "../util/pathToUri.js";
 import {
   INSTINCT_SYSTEM_PROMPT,
-  MERCURY_CODE_TO_EDIT_OPEN,
   MERCURY_SYSTEM_PROMPT,
   MODEL_WINDOW_SIZES,
+  UNIQUE_TOKEN,
 } from "./constants.js";
 import { createDiff, DiffFormatType } from "./context/diffFormatting.js";
 import {
@@ -392,9 +393,13 @@ export class NextEditProvider {
       this.ide.getWorkspaceDirs(),
     ]);
 
-    const modelName = helper.modelName.includes("mercury-coder-nextedit")
-      ? "mercury-coder-nextedit"
-      : "instinct";
+    const modelName = helper.modelName.includes(
+      NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT,
+    )
+      ? NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT
+      : helper.modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER)
+        ? NEXT_EDIT_MODELS.MERCURY_CODER
+        : NEXT_EDIT_MODELS.INSTINCT;
 
     const { editableRegionStartLine, editableRegionEndLine } =
       opts?.usingFullFileDiff
@@ -551,7 +556,10 @@ export class NextEditProvider {
     const modelName = helper.modelName;
     let ctx: any;
 
-    if (modelName.includes("mercury-coder-nextedit")) {
+    if (
+      modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER) ||
+      modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT)
+    ) {
       ctx = {
         recentlyViewedCodeSnippets:
           snippetPayload.recentlyVisitedRangesSnippets.map((snip) => ({
@@ -564,7 +572,7 @@ export class NextEditProvider {
         editDiffHistory: this.diffContext,
         currentFilePath: helper.filepath,
       };
-    } else if (modelName.includes("instinct")) {
+    } else if (modelName.includes(NEXT_EDIT_MODELS.INSTINCT)) {
       // Calculate the window around the cursor position (25 lines above and below).
       const windowStart = Math.max(0, helper.pos.line - 25);
       const windowEnd = Math.min(
@@ -608,9 +616,11 @@ export class NextEditProvider {
 
     const systemPrompt: Prompt = {
       role: "system",
-      content: modelName.includes("mercury-coder-nextedit")
-        ? MERCURY_SYSTEM_PROMPT
-        : INSTINCT_SYSTEM_PROMPT,
+      content:
+        modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER) ||
+        modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT)
+          ? MERCURY_SYSTEM_PROMPT
+          : INSTINCT_SYSTEM_PROMPT,
     };
 
     return [systemPrompt, promptMetadata.prompt];
@@ -677,25 +687,37 @@ export class NextEditProvider {
     const llm = await this._prepareLlm();
     if (!llm) return undefined;
 
-    const msg: ChatMessage = await llm.chat(prompts, token);
-    console.log("message");
-    console.log(msg.content);
+    // Check for mercury models, and inject unique token for next edit.
+    // Capture the unique tokens in llm Inception.ts and apis Inception.ts
+    const modelName = helper.modelName;
+    if (
+      modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER) ||
+      modelName.includes(NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT)
+    ) {
+      // Inject the unique token to the last prompt's content.
+      const lastPrompt = prompts[prompts.length - 1];
+      if (lastPrompt && typeof lastPrompt.content === "string") {
+        lastPrompt.content += UNIQUE_TOKEN;
+      }
+    }
+
+    const msg: ChatMessage = await llm.chat([prompts[1]], token, {
+      stream: false,
+    });
 
     if (typeof msg.content !== "string") {
       return undefined;
     }
 
-    // const nextCompletion = msg.content.split(
-    //   `${MERCURY_CODE_TO_EDIT_OPEN}\n`,
-    // )[1]
-    //   ? replaceEscapedCharacters(
-    //       msg.content.split(`${MERCURY_CODE_TO_EDIT_OPEN}\n`)[1],
-    //     ).replace(/\n$/, "")
-    //   : replaceEscapedCharacters(msg.content);
-    const nextCompletion =
-      msg.content
-        .split(`${MERCURY_CODE_TO_EDIT_OPEN}\n`)[1]
-        .replace(/\n$/, "") ?? msg.content;
+    // NOTE: Keep this in case there is a change in spec.
+    // const nextCompletion =
+    //   msg.content
+    //     .split(`${MERCURY_CODE_TO_EDIT_OPEN}\n`)[1]
+    //     .replace(/\n$/, "") ?? msg.content;
+    const nextCompletion = msg.content.slice(
+      msg.content.indexOf("\`\`\`\n") + "\'\'\'\n".length,
+      msg.content.lastIndexOf("\n\n\`\`\`"),
+    );
 
     if (opts?.usingFullFileDiff === false || !opts?.usingFullFileDiff) {
       return await this._handlePartialFileDiff(
diff --git a/core/nextEdit/constants.ts b/core/nextEdit/constants.ts
@@ -8,13 +8,19 @@ export const MODEL_WINDOW_SIZES: Record<
   NEXT_EDIT_MODELS,
   { topMargin: number; bottomMargin: number }
 > = {
+  "mercury-coder": {
+    topMargin: 0,
+    bottomMargin: 5,
+  }, // mercury coder uses full file diff, so this should be unnecessary
   "mercury-coder-nextedit": {
     topMargin: 0,
     bottomMargin: 5,
   }, // mercury coder nextedit uses full file diff, so this should be unnecessary
   instinct: { topMargin: 1, bottomMargin: 5 },
 };
 
+export const UNIQUE_TOKEN = "<|!@#IS_NEXT_EDIT!@#|>";
+
 // Model 1-specific tokens.
 export const INSTINCT_USER_CURSOR_IS_HERE_TOKEN = "<|user_cursor_is_here|>";
 export const INSTINCT_EDITABLE_REGION_START_TOKEN = "<|editable_region_start|>";
diff --git a/core/nextEdit/templating/NextEditPromptEngine.ts b/core/nextEdit/templating/NextEditPromptEngine.ts
@@ -39,8 +39,11 @@ import {
 type TemplateRenderer = (vars: TemplateVars) => string;
 
 const NEXT_EDIT_MODEL_TEMPLATES: Record<NEXT_EDIT_MODELS, NextEditTemplate> = {
+  "mercury-coder": {
+    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`,
+  },
   "mercury-coder-nextedit": {
-    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.`,
+    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`,
   },
   instinct: {
     template: `${INSTINCT_USER_PROMPT_PREFIX}\n\n### Context:\n{{{contextSnippets}}}\n\n### User Edits:\n\n{{{editDiffHistory}}}\n\n### User Excerpts:\n{{{currentFilePath}}}\n\n{{{currentFileContent}}}\`\`\`\n### Response:`,
@@ -90,7 +93,36 @@ export async function renderPrompt(
   let tv: TemplateVars;
 
   switch (modelName) {
-    case "mercury-coder-nextedit": {
+    case NEXT_EDIT_MODELS.MERCURY_CODER: {
+      userEdits = ctx.editDiffHistory;
+
+      // editedCodeWithTokens = insertTokens(
+      //   helper.fileContents.split("\n"),
+      //   helper.pos,
+      //   MERCURY_CURSOR,
+      // );
+
+      const mercuryCtx: MercuryTemplateVars = {
+        recentlyViewedCodeSnippets: recentlyViewedCodeSnippetsBlock(
+          ctx.recentlyViewedCodeSnippets,
+        ),
+        currentFileContent: currentFileContentBlock(
+          ctx.currentFileContent,
+          ctx.editableRegionStartLine,
+          ctx.editableRegionEndLine,
+          helper.pos,
+        ),
+        editDiffHistory: editHistoryBlock(ctx.editDiffHistory),
+        currentFilePath: ctx.currentFilePath,
+      };
+
+      tv = mercuryCtx;
+
+      editedCodeWithTokens = mercuryCtx.currentFileContent;
+
+      break;
+    }
+    case NEXT_EDIT_MODELS.MERCURY_CODER_NEXTEDIT: {
       userEdits = ctx.editDiffHistory;
 
       // editedCodeWithTokens = insertTokens(
@@ -119,7 +151,7 @@ export async function renderPrompt(
 
       break;
     }
-    case "instinct": {
+    case NEXT_EDIT_MODELS.INSTINCT: {
       userEdits = ctx.editDiffHistory;
 
       // editedCodeWithTokens = insertTokens(
diff --git a/packages/openai-adapters/src/apis/Inception.ts b/packages/openai-adapters/src/apis/Inception.ts

Original file line number	Diff line number	Diff line change
`@@ -20,6 +20,7 @@ export enum LLMConfigurationStatuses {`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`export enum NEXT_EDIT_MODELS {`
	`23`	`+ MERCURY_CODER = "mercury-coder",`
`23`	`24`	`MERCURY_CODER_NEXTEDIT = "mercury-coder-nextedit",`
`24`	`25`	`INSTINCT = "instinct",`
`25`	`26`	`}`