Skip to content

Commit e27be8f

Browse files
alexandrevilainalexandrevilain
authored
Merge pull request #19 from alexandrevilain/feat/fim-completion
feat(autocomplete): add FIM support with Mistral Codestral provider
2 parents 49e9dea + 1484ffb commit e27be8f

File tree

11 files changed

+225
-135
lines changed

11 files changed

+225
-135
lines changed

‎package-lock.json‎

Lines changed: 35 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,7 @@
150150
"@ai-sdk/openai-compatible": "^1.0.7",
151151
"@ai-sdk/provider": "^2.0.0",
152152
"ai": "^5.0.14",
153+
"ai-sdk-mistral-fim": "^0.0.1",
153154
"ai-sdk-ollama": "^0.5.0",
154155
"uuid": "^11.1.0"
155156
}

‎src/autocomplete/context.ts‎

Lines changed: 0 additions & 7 deletions
This file was deleted.

‎src/autocomplete/defaultHoleFiller.ts‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import { HoleFiller, PromptArgs, AutoCompleteContext } from "./holeFiller";
2+
3+
// Source: continue/core/autocomplete/templating/AutocompleteTemplate.ts (holeFillerTemplate)
4+
export class DefaultHoleFiller implements HoleFiller {
5+
systemPrompt(): string {
6+
// From https://github.com/VictorTaelin/AI-scripts
7+
return `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'.
8+
Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.
9+
All completions MUST be truthful, accurate, well-written and correct.
10+
## EXAMPLE QUERY:
11+
12+
<QUERY>
13+
function sum_evens(lim) {
14+
var sum = 0;
15+
for (var i = 0; i < lim; ++i) {
16+
{{FILL_HERE}}
17+
}
18+
return sum;
19+
}
20+
</QUERY>
21+
22+
TASK: Fill the {{FILL_HERE}} hole.
23+
24+
## CORRECT COMPLETION
25+
26+
<COMPLETION>if (i % 2 === 0) {
27+
sum += i;
28+
}</COMPLETION>
29+
30+
## EXAMPLE QUERY:
31+
32+
<QUERY>
33+
def sum_list(lst):
34+
total = 0
35+
for x in lst:
36+
{{FILL_HERE}}
37+
return total
38+
39+
print sum_list([1, 2, 3])
40+
</QUERY>
41+
42+
## CORRECT COMPLETION:
43+
44+
<COMPLETION> total += x</COMPLETION>
45+
46+
## EXAMPLE QUERY:
47+
48+
<QUERY>
49+
// data Tree a = Node (Tree a) (Tree a) | Leaf a
50+
51+
// sum :: Tree Int -> Int
52+
// sum (Node lft rgt) = sum lft + sum rgt
53+
// sum (Leaf val) = val
54+
55+
// convert to TypeScript:
56+
{{FILL_HERE}}
57+
</QUERY>
58+
59+
## CORRECT COMPLETION:
60+
61+
<COMPLETION>type Tree<T>
62+
= {$:"Node", lft: Tree<T>, rgt: Tree<T>}
63+
| {$:"Leaf", val: T};
64+
65+
function sum(tree: Tree<number>): number {
66+
switch (tree.$) {
67+
case "Node":
68+
return sum(tree.lft) + sum(tree.rgt);
69+
case "Leaf":
70+
return tree.val;
71+
}
72+
}</COMPLETION>
73+
74+
## EXAMPLE QUERY:
75+
76+
The 5th {{FILL_HERE}} is Jupiter.
77+
78+
## CORRECT COMPLETION:
79+
80+
<COMPLETION>planet from the Sun</COMPLETION>
81+
82+
## EXAMPLE QUERY:
83+
84+
function hypothenuse(a, b) {
85+
return Math.sqrt({{FILL_HERE}}b ** 2);
86+
}
87+
88+
## CORRECT COMPLETION:
89+
90+
<COMPLETION>a ** 2 + </COMPLETION>
91+
`;
92+
}
93+
94+
userPrompt(ctx: AutoCompleteContext): string {
95+
let context = '';
96+
if (ctx.filename !== '') {
97+
context += `// Filename: "${ctx.filename}" \n`;
98+
}
99+
if (ctx.language !== '') {
100+
context += `// Programming language: "${ctx.language}" \n`;
101+
}
102+
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>`;
103+
}
104+
105+
prompt(params: AutoCompleteContext): PromptArgs {
106+
return {
107+
messages: [
108+
{ role: "system", content: this.systemPrompt() },
109+
{ role: "user", content: this.userPrompt(params) },
110+
],
111+
};
112+
}
113+
}
114+

‎src/autocomplete/holeFiller.ts‎

Lines changed: 14 additions & 106 deletions
Original file line numberDiff line numberDiff line change
@@ -1,110 +1,18 @@
1-
import { type AutoCompleteContext } from "./context";
1+
import { Prompt } from 'ai';
2+
import { ProviderOptions } from '@ai-sdk/provider-utils';
23

34
export interface HoleFiller {
4-
systemPrompt(): string
5-
userPrompt(params: AutoCompleteContext): string
6-
}
7-
8-
// Source: continue/core/autocomplete/templating/AutocompleteTemplate.ts (holeFillerTemplate)
9-
export class DefaultHoleFiller implements HoleFiller {
10-
systemPrompt(): string {
11-
// From https://github.com/VictorTaelin/AI-scripts
12-
return `You are a HOLE FILLER. You are provided with a file containing holes, formatted as '{{HOLE_NAME}}'.
13-
Your TASK is to complete with a string to replace this hole with, inside a <COMPLETION/> XML tag, including context-aware indentation, if needed.
14-
All completions MUST be truthful, accurate, well-written and correct.
15-
## EXAMPLE QUERY:
16-
17-
<QUERY>
18-
function sum_evens(lim) {
19-
var sum = 0;
20-
for (var i = 0; i < lim; ++i) {
21-
{{FILL_HERE}}
22-
}
23-
return sum;
24-
}
25-
</QUERY>
26-
27-
TASK: Fill the {{FILL_HERE}} hole.
28-
29-
## CORRECT COMPLETION
30-
31-
<COMPLETION>if (i % 2 === 0) {
32-
sum += i;
33-
}</COMPLETION>
34-
35-
## EXAMPLE QUERY:
36-
37-
<QUERY>
38-
def sum_list(lst):
39-
total = 0
40-
for x in lst:
41-
{{FILL_HERE}}
42-
return total
43-
44-
print sum_list([1, 2, 3])
45-
</QUERY>
46-
47-
## CORRECT COMPLETION:
48-
49-
<COMPLETION> total += x</COMPLETION>
50-
51-
## EXAMPLE QUERY:
52-
53-
<QUERY>
54-
// data Tree a = Node (Tree a) (Tree a) | Leaf a
55-
56-
// sum :: Tree Int -> Int
57-
// sum (Node lft rgt) = sum lft + sum rgt
58-
// sum (Leaf val) = val
59-
60-
// convert to TypeScript:
61-
{{FILL_HERE}}
62-
</QUERY>
63-
64-
## CORRECT COMPLETION:
65-
66-
<COMPLETION>type Tree<T>
67-
= {$:"Node", lft: Tree<T>, rgt: Tree<T>}
68-
| {$:"Leaf", val: T};
69-
70-
function sum(tree: Tree<number>): number {
71-
switch (tree.$) {
72-
case "Node":
73-
return sum(tree.lft) + sum(tree.rgt);
74-
case "Leaf":
75-
return tree.val;
76-
}
77-
}</COMPLETION>
78-
79-
## EXAMPLE QUERY:
80-
81-
The 5th {{FILL_HERE}} is Jupiter.
82-
83-
## CORRECT COMPLETION:
84-
85-
<COMPLETION>planet from the Sun</COMPLETION>
86-
87-
## EXAMPLE QUERY:
88-
89-
function hypothenuse(a, b) {
90-
return Math.sqrt({{FILL_HERE}}b ** 2);
91-
}
92-
93-
## CORRECT COMPLETION:
94-
95-
<COMPLETION>a ** 2 + </COMPLETION>
96-
`;
97-
}
98-
99-
userPrompt(ctx: AutoCompleteContext): string {
100-
let context = '';
101-
if (ctx.filename !== '') {
102-
context += `// Filename: "${ctx.filename}" \n`;
103-
}
104-
if (ctx.language !== '') {
105-
context += `// Programming language: "${ctx.language}" \n`;
106-
}
107-
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>`;
108-
}
5+
prompt(params: AutoCompleteContext): PromptArgs
1096
}
1107

8+
export type PromptArgs = Prompt & {
9+
providerOptions?: ProviderOptions;
10+
};
11+
12+
export type AutoCompleteContext = {
13+
textBeforeCursor: string,
14+
textAfterCursor: string,
15+
currentLineText: string,
16+
filename?: string,
17+
language?: string,
18+
}

‎src/autocomplete/index.ts‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
export * from './defaultHoleFiller';
2+
export * from './mistralfimHoleFiller';
3+
export * from './holeFiller';

‎src/autocomplete/mistralfimHoleFiller.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import { HoleFiller, PromptArgs, AutoCompleteContext } from "./holeFiller";
2+
3+
export class MistralFimHoleFiller implements HoleFiller {
4+
prompt(params: AutoCompleteContext): PromptArgs {
5+
return {
6+
prompt: params.textBeforeCursor,
7+
providerOptions: {
8+
'mistral.fim': {
9+
suffix: params.textAfterCursor,
10+
}
11+
}
12+
};
13+
}
14+
}

‎src/providers/codestral.ts‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import { ProfileWithAPIKey, ProviderConnection, Model } from "../types";
2+
import { type LanguageModelV2 } from "@ai-sdk/provider";
3+
import { LanguageModelProvider } from "./providers";
4+
import { createMistralFim } from 'ai-sdk-mistral-fim';
5+
6+
export class CodestralProvider implements LanguageModelProvider {
7+
languageModel(profile: ProfileWithAPIKey): LanguageModelV2 {
8+
return createMistralFim({
9+
baseURL: profile.baseURL,
10+
apiKey: profile.apiKey,
11+
})(profile.modelId);
12+
}
13+
14+
async listModels(_conn: ProviderConnection): Promise<Model[]> {
15+
return new Promise((resolve) => {
16+
resolve([
17+
{
18+
id: 'codestral-latest',
19+
name: 'codestral-latest'
20+
},
21+
]);
22+
});
23+
}
24+
}

‎src/providers/providers.ts‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ import { ProfileWithAPIKey, Provider, ProviderConnection, ProviderID, Model } fr
22
import { OpenAICompatibleProvider } from "./openaiCompatible";
33
import { OllamaProvider } from "./ollama";
44
import { type LanguageModelV2 } from "@ai-sdk/provider";
5+
import { CodestralProvider } from "./codestral";
56

67
export interface LanguageModelProvider {
78
languageModel(profile: ProfileWithAPIKey): LanguageModelV2
@@ -17,10 +18,11 @@ function languageModelProvider(providerId: ProviderID): LanguageModelProvider {
1718
case 'groq':
1819
case 'openai-compatible':
1920
case 'mistral':
20-
case 'mistral-codestral': // TODO: we should support FIM endpoint.
21-
return new OpenAICompatibleProvider();
21+
return new OpenAICompatibleProvider();
2222
case 'ollama':
23-
return new OllamaProvider();
23+
return new OllamaProvider();
24+
case 'mistral-codestral':
25+
return new CodestralProvider();
2426
default:
2527
throw new Error(`Unsupported provider: ${providerId}`);
2628
}
@@ -34,6 +36,9 @@ export function getLanguageModelFromProfile(profile: ProfileWithAPIKey): Languag
3436
return languageModelProvider(profile.provider).languageModel(profile);
3537
}
3638

39+
export function isFimProvider(provider: ProviderID): boolean {
40+
return provider === 'mistral-codestral';
41+
}
3742

3843
export const providers: Provider[] = [
3944
{

0 commit comments

Comments
 (0)