Skip to content

Commit 8d69eb9

Browse files
author
Liang Chen
authored
feat: integrate Mistral AI adapter (#69)
1 parent 9553e10 commit 8d69eb9

2 files changed

Lines changed: 115 additions & 5 deletions

File tree

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,80 @@
11
import { smokeTest } from '@profullstack/sh1pt-core/testing';
2+
import { afterEach, describe, expect, it, vi } from 'vitest';
23
import adapter from './index.js';
34

45
smokeTest(adapter, { idPrefix: 'ai' });
6+
7+
const ctx = (secrets: Record<string, string> = { MISTRAL_API_KEY: 'test-key' }, dryRun = false) => ({
8+
secret: (key: string) => secrets[key],
9+
log: () => {},
10+
dryRun,
11+
});
12+
13+
describe('Mistral OpenAI-compatible generation', () => {
14+
afterEach(() => {
15+
vi.unstubAllGlobals();
16+
});
17+
18+
it('short-circuits dry-run before network calls', async () => {
19+
const fetchMock = vi.fn();
20+
vi.stubGlobal('fetch', fetchMock);
21+
22+
const result = await adapter.generate(ctx({ MISTRAL_API_KEY: 'test-key' }, true), 'hello', {}, {});
23+
24+
expect(result).toEqual({ text: '[dry-run]', model: 'mistral-large-latest' });
25+
expect(fetchMock).not.toHaveBeenCalled();
26+
});
27+
28+
it('posts chat completions requests and maps usage tokens', async () => {
29+
const fetchMock = vi.fn().mockResolvedValue({
30+
ok: true,
31+
json: async () => ({
32+
choices: [{ message: { content: 'hi from mistral' } }],
33+
model: 'mistral-small-latest',
34+
usage: { prompt_tokens: 8, completion_tokens: 4 },
35+
}),
36+
});
37+
vi.stubGlobal('fetch', fetchMock);
38+
39+
const result = await adapter.generate(ctx(), 'hello', {
40+
model: 'mistral-small-latest',
41+
system: 'be direct',
42+
maxTokens: 24,
43+
temperature: 0.3,
44+
extra: { top_p: 0.7 },
45+
}, {});
46+
47+
expect(fetchMock).toHaveBeenCalledOnce();
48+
const call = fetchMock.mock.calls[0];
49+
expect(call).toBeDefined();
50+
const [url, request] = call!;
51+
expect(url).toBe('https://api.mistral.ai/v1/chat/completions');
52+
expect(request.headers.authorization).toBe('Bearer test-key');
53+
expect(JSON.parse(request.body)).toEqual({
54+
model: 'mistral-small-latest',
55+
messages: [
56+
{ role: 'system', content: 'be direct' },
57+
{ role: 'user', content: 'hello' },
58+
],
59+
max_tokens: 24,
60+
temperature: 0.3,
61+
top_p: 0.7,
62+
});
63+
expect(result).toEqual({
64+
text: 'hi from mistral',
65+
model: 'mistral-small-latest',
66+
inputTokens: 8,
67+
outputTokens: 4,
68+
});
69+
});
70+
71+
it('includes status and response body excerpt on errors', async () => {
72+
vi.stubGlobal('fetch', vi.fn().mockResolvedValue({
73+
ok: false,
74+
status: 400,
75+
text: async () => 'bad request'.repeat(30),
76+
}));
77+
78+
await expect(adapter.generate(ctx(), 'hello', {}, {})).rejects.toThrow(/Mistral 400: bad request/);
79+
});
80+
});

packages/ai/mistral/src/index.ts

Lines changed: 39 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,17 +4,51 @@ interface Config {
44
baseUrl?: string;
55
}
66

7+
const DEFAULT_BASE = 'https://api.mistral.ai';
8+
79
export default defineAi<Config>({
810
id: 'ai-mistral',
911
label: 'Mistral',
1012
defaultModel: 'mistral-large-latest',
11-
models: ['mistral-large-latest'],
13+
models: ['mistral-large-latest', 'mistral-small-latest', 'codestral-latest', 'ministral-8b-latest'],
1214

13-
async generate(ctx, prompt, _opts, _config) {
15+
async generate(ctx, prompt, opts, config) {
1416
const apiKey = ctx.secret('MISTRAL_API_KEY');
15-
if (!apiKey) throw new Error('MISTRAL_API_KEY not in vault — run `sh1pt promote ai setup`');
16-
ctx.log(`[stub] ai-mistral · ${prompt.length} chars in — integration pending`);
17-
return { text: '[stub — ai-mistral integration not yet implemented]', model: 'mistral-large-latest' };
17+
if (!apiKey) throw new Error('MISTRAL_API_KEY not in vault');
18+
const model = opts.model ?? 'mistral-large-latest';
19+
ctx.log(`mistral · model=${model} · ${prompt.length} chars in`);
20+
if (ctx.dryRun) return { text: '[dry-run]', model };
21+
22+
const messages: Array<{ role: string; content: string }> = [];
23+
if (opts.system) messages.push({ role: 'system', content: opts.system });
24+
messages.push({ role: 'user', content: prompt });
25+
26+
const res = await fetch(`${config.baseUrl ?? DEFAULT_BASE}/v1/chat/completions`, {
27+
method: 'POST',
28+
headers: {
29+
authorization: `Bearer ${apiKey}`,
30+
'content-type': 'application/json',
31+
},
32+
body: JSON.stringify({
33+
model,
34+
messages,
35+
...(opts.maxTokens !== undefined ? { max_tokens: opts.maxTokens } : {}),
36+
...(opts.temperature !== undefined ? { temperature: opts.temperature } : {}),
37+
...opts.extra,
38+
}),
39+
});
40+
if (!res.ok) throw new Error(`Mistral ${res.status}: ${(await res.text()).slice(0, 200)}`);
41+
const data = (await res.json()) as {
42+
choices: Array<{ message?: { content?: string } }>;
43+
model: string;
44+
usage?: { prompt_tokens?: number; completion_tokens?: number };
45+
};
46+
return {
47+
text: data.choices[0]?.message?.content ?? '',
48+
model: data.model,
49+
inputTokens: data.usage?.prompt_tokens,
50+
outputTokens: data.usage?.completion_tokens,
51+
};
1852
},
1953

2054
setup: tokenSetup<Config>({

0 commit comments

Comments
 (0)