feat(autocomplete): add FIM support with Mistral Codestral provider

package.json

@@ -150,6 +150,7 @@
     "@ai-sdk/openai-compatible": "^1.0.7",
     "@ai-sdk/provider": "^2.0.0",
     "ai": "^5.0.14",
+    "ai-sdk-mistral-fim": "^0.0.1",
     "ai-sdk-ollama": "^0.5.0",
     "uuid": "^11.1.0"
   }

src/autocomplete/defaultHoleFiller.ts

@@ -0,0 +1,114 @@
+import { HoleFiller, PromptArgs, AutoCompleteContext } from "./holeFiller";
+
+// Source: continue/core/autocomplete/templating/AutocompleteTemplate.ts (holeFillerTemplate)
+export class DefaultHoleFiller implements HoleFiller {
+  systemPrompt(): string {
+    // From https://github.com/VictorTaelin/AI-scripts
+    return `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'.
+		Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.
+		All completions MUST be truthful, accurate, well-written and correct.
+## EXAMPLE QUERY:
+
+<QUERY>
+function sum_evens(lim) {
+  var sum = 0;
+  for (var i = 0; i < lim; ++i) {
+    {{FILL_HERE}}
+  }
+  return sum;
+}
+</QUERY>
+
+TASK: Fill the {{FILL_HERE}} hole.
+
+## CORRECT COMPLETION
+
+<COMPLETION>if (i % 2 === 0) {
+      sum += i;
+    }</COMPLETION>
+
+## EXAMPLE QUERY:
+
+<QUERY>
+def sum_list(lst):
+  total = 0
+  for x in lst:
+  {{FILL_HERE}}
+  return total
+
+print sum_list([1, 2, 3])
+</QUERY>
+
+## CORRECT COMPLETION:
+
+<COMPLETION>  total += x</COMPLETION>
+
+## EXAMPLE QUERY:
+
+<QUERY>
+// data Tree a = Node (Tree a) (Tree a) | Leaf a
+
+// sum :: Tree Int -> Int
+// sum (Node lft rgt) = sum lft + sum rgt
+// sum (Leaf val)     = val
+
+// convert to TypeScript:
+{{FILL_HERE}}
+</QUERY>
+
+## CORRECT COMPLETION:
+
+<COMPLETION>type Tree<T>
+  = {$:"Node", lft: Tree<T>, rgt: Tree<T>}
+  | {$:"Leaf", val: T};
+
+function sum(tree: Tree<number>): number {
+  switch (tree.$) {
+    case "Node":
+      return sum(tree.lft) + sum(tree.rgt);
+    case "Leaf":
+      return tree.val;
+  }
+}</COMPLETION>
+
+## EXAMPLE QUERY:
+
+The 5th {{FILL_HERE}} is Jupiter.
+
+## CORRECT COMPLETION:
+
+<COMPLETION>planet from the Sun</COMPLETION>
+
+## EXAMPLE QUERY:
+
+function hypothenuse(a, b) {
+  return Math.sqrt({{FILL_HERE}}b ** 2);
+}
+
+## CORRECT COMPLETION:
+
+<COMPLETION>a ** 2 + </COMPLETION>
+`;
+  }
+
+  userPrompt(ctx: AutoCompleteContext): string {
+    let context = '';
+    if (ctx.filename !== '') {
+      context += `// Filename: "${ctx.filename}" \n`;
+    }
+    if (ctx.language !== '') {
+      context += `// Programming language: "${ctx.language}" \n`;
+    }
+    return `${context}<QUERY>\n${ctx.textBeforeCursor}{{FILL_HERE}}${ctx.textAfterCursor}\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n<COMPLETION>`;
+  }
+
+  prompt(params: AutoCompleteContext): PromptArgs {
+    return {
+      messages: [
+        { role: "system", content: this.systemPrompt() },
+        { role: "user", content: this.userPrompt(params) },
+      ],
+    };
+  }
+}
+

src/autocomplete/holeFiller.ts

@@ -1,110 +1,18 @@
-import { type AutoCompleteContext } from "./context";
+import { Prompt } from 'ai';
+import { ProviderOptions } from '@ai-sdk/provider-utils';
 
 export interface HoleFiller {
-  systemPrompt(): string
-  userPrompt(params: AutoCompleteContext): string
-}
-
-// Source: continue/core/autocomplete/templating/AutocompleteTemplate.ts (holeFillerTemplate)
-export class DefaultHoleFiller implements HoleFiller {
-	systemPrompt(): string  {
-		// From https://github.com/VictorTaelin/AI-scripts
-		return `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'.
-		Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.
-		All completions MUST be truthful, accurate, well-written and correct.
-## EXAMPLE QUERY:
-
-<QUERY>
-function sum_evens(lim) {
-  var sum = 0;
-  for (var i = 0; i < lim; ++i) {
-    {{FILL_HERE}}
-  }
-  return sum;
-}
-</QUERY>
-
-TASK: Fill the {{FILL_HERE}} hole.
-
-## CORRECT COMPLETION
-
-<COMPLETION>if (i % 2 === 0) {
-      sum += i;
-    }</COMPLETION>
-
-## EXAMPLE QUERY:
-
-<QUERY>
-def sum_list(lst):
-  total = 0
-  for x in lst:
-  {{FILL_HERE}}
-  return total
-
-print sum_list([1, 2, 3])
-</QUERY>
-
-## CORRECT COMPLETION:
-
-<COMPLETION>  total += x</COMPLETION>
-
-## EXAMPLE QUERY:
-
-<QUERY>
-// data Tree a = Node (Tree a) (Tree a) | Leaf a
-
-// sum :: Tree Int -> Int
-// sum (Node lft rgt) = sum lft + sum rgt
-// sum (Leaf val)     = val
-
-// convert to TypeScript:
-{{FILL_HERE}}
-</QUERY>
-
-## CORRECT COMPLETION:
-
-<COMPLETION>type Tree<T>
-  = {$:"Node", lft: Tree<T>, rgt: Tree<T>}
-  | {$:"Leaf", val: T};
-
-function sum(tree: Tree<number>): number {
-  switch (tree.$) {
-    case "Node":
-      return sum(tree.lft) + sum(tree.rgt);
-    case "Leaf":
-      return tree.val;
-  }
-}</COMPLETION>
-
-## EXAMPLE QUERY:
-
-The 5th {{FILL_HERE}} is Jupiter.
-
-## CORRECT COMPLETION:
-
-<COMPLETION>planet from the Sun</COMPLETION>
-
-## EXAMPLE QUERY:
-
-function hypothenuse(a, b) {
-  return Math.sqrt({{FILL_HERE}}b ** 2);
-}
-
-## CORRECT COMPLETION:
-
-<COMPLETION>a ** 2 + </COMPLETION>
-`;
-	}
-
-	userPrompt(ctx: AutoCompleteContext): string {
-    let context = '';
-    if (ctx.filename !== '') {
-      context += `// Filename: "${ctx.filename}" \n`;
-    }
-    if (ctx.language !== '') {
-      context += `// Programming language: "${ctx.language}" \n`;
-    }
-		return `${context}<QUERY>\n${ctx.textBeforeCursor}{{FILL_HERE}}${ctx.textAfterCursor}\n</QUERY>\nTASK: Fill the {{FILL_HERE}} hole. Answer only with the CORRECT completion, and NOTHING ELSE. Do it now.\n<COMPLETION>`;
-	}
+  prompt(params: AutoCompleteContext): PromptArgs
 }
 
+export type PromptArgs = Prompt & {
+    providerOptions?: ProviderOptions;
+};
+
+export type AutoCompleteContext = {
+  textBeforeCursor: string,
+  textAfterCursor: string,
+  currentLineText: string,
+  filename?: string,
+  language?: string,
+}

src/autocomplete/index.ts

@@ -0,0 +1,3 @@
+export * from './defaultHoleFiller';
+export * from './mistralfimHoleFiller';
+export * from './holeFiller';

src/autocomplete/mistralfimHoleFiller.ts

@@ -0,0 +1,14 @@
+import { HoleFiller, PromptArgs, AutoCompleteContext } from "./holeFiller";
+
+export class MistralFimHoleFiller implements HoleFiller {
+    prompt(params: AutoCompleteContext): PromptArgs {
+        return {
+            prompt: params.textBeforeCursor,
+            providerOptions: {
+                'mistral.fim': {
+                    suffix: params.textAfterCursor,
+                }
+            }
+        };
+    }
+}

src/providers/codestral.ts

@@ -0,0 +1,24 @@
+import { ProfileWithAPIKey, ProviderConnection, Model } from "../types";
+import { type LanguageModelV2 } from "@ai-sdk/provider";
+import { LanguageModelProvider } from "./providers";
+import { createMistralFim } from 'ai-sdk-mistral-fim';
+
+export class CodestralProvider implements LanguageModelProvider {
+    languageModel(profile: ProfileWithAPIKey): LanguageModelV2 {
+        return createMistralFim({
+            baseURL: profile.baseURL,
+            apiKey: profile.apiKey,
+        })(profile.modelId);
+    }
+
+    async listModels(_conn: ProviderConnection): Promise<Model[]> {
+        return new Promise((resolve) => {
+            resolve([
+                {
+                    id: 'codestral-latest',
+                    name: 'codestral-latest'
+                },
+            ]);
+        });
+    }
+}

src/providers/providers.ts

@@ -2,6 +2,7 @@ import { ProfileWithAPIKey, Provider, ProviderConnection, ProviderID, Model } fr
 import { OpenAICompatibleProvider } from "./openaiCompatible";
 import { OllamaProvider } from "./ollama";
 import { type LanguageModelV2 } from "@ai-sdk/provider";
+import { CodestralProvider } from "./codestral";
 
 export interface LanguageModelProvider {
     languageModel(profile: ProfileWithAPIKey): LanguageModelV2
@@ -17,10 +18,11 @@ function languageModelProvider(providerId: ProviderID): LanguageModelProvider {
         case 'groq':
         case 'openai-compatible':
         case 'mistral':
-        case 'mistral-codestral': // TODO: we should support FIM endpoint.
-            return new OpenAICompatibleProvider();
+          return new OpenAICompatibleProvider();
         case 'ollama':
-            return new OllamaProvider();
+          return new OllamaProvider();
+        case 'mistral-codestral':
+          return new CodestralProvider();
         default:
             throw new Error(`Unsupported provider: ${providerId}`);
     }
@@ -34,6 +36,9 @@ export function getLanguageModelFromProfile(profile: ProfileWithAPIKey): Languag
     return languageModelProvider(profile.provider).languageModel(profile);
 }
 
+export function isFimProvider(provider: ProviderID): boolean {
+  return provider === 'mistral-codestral';
+}
 
 export const providers: Provider[] = [
   {