Skip to content

Commit 3c595ac

Browse files
committed
feat: support Gemini in AI-SDK runner
1 parent 5d1ab29 commit 3c595ac

File tree

3 files changed

+37
-3
lines changed

3 files changed

+37
-3
lines changed

package.json

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
},
5353
"dependencies": {
5454
"@ai-sdk/anthropic": "^2.0.45",
55+
"@ai-sdk/google": "^2.0.39",
5556
"@anthropic-ai/sdk": "^0.68.0",
5657
"@axe-core/puppeteer": "^4.10.2",
5758
"@genkit-ai/compat-oai": "1.23.0",

pnpm-lock.yaml

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

runner/codegen/ai-sdk-runner.ts

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
SystemModelMessage,
1818
TextPart,
1919
} from 'ai';
20+
import {google, GoogleGenerativeAIProviderOptions} from '@ai-sdk/google';
2021
import {anthropic, AnthropicProviderOptions} from '@ai-sdk/anthropic';
2122
import z from 'zod';
2223
import {callWithTimeout} from '../utils/timeout.js';
@@ -27,6 +28,10 @@ const SUPPORTED_MODELS = [
2728
'claude-opus-4.1-with-thinking',
2829
'claude-sonnet-4.5-no-thinking',
2930
'claude-sonnet-4.5-with-thinking',
31+
'gemini-2.5-flash-lite',
32+
'gemini-2.5-flash',
33+
'gemini-2.5-pro',
34+
'gemini-3-pro-preview',
3035
] as const;
3136

3237
// Increased to a very high value as we rely on an actual timeout
@@ -131,7 +136,8 @@ export class AiSDKRunner implements LlmRunner {
131136
private async _getAiSdkModelOptions(
132137
request: LocalLlmGenerateTextRequestOptions,
133138
): Promise<{model: LanguageModel; providerOptions: {}}> {
134-
switch (request.model) {
139+
const modelName = request.model as (typeof SUPPORTED_MODELS)[number];
140+
switch (modelName) {
135141
case 'claude-opus-4.1-no-thinking':
136142
case 'claude-opus-4.1-with-thinking': {
137143
const thinkingEnabled = request.model.endsWith('with-thinking');
@@ -149,11 +155,23 @@ export class AiSDKRunner implements LlmRunner {
149155
return {
150156
model: anthropic('claude-sonnet-4-5'),
151157
providerOptions: {
152-
sendReasoning: true,
153-
thinking: {type: 'enabled'},
158+
sendReasoning: thinkingEnabled,
159+
thinking: {type: thinkingEnabled ? 'enabled' : 'disabled'},
154160
} satisfies AnthropicProviderOptions,
155161
};
156162
}
163+
case 'gemini-2.5-flash-lite':
164+
case 'gemini-2.5-flash':
165+
case 'gemini-2.5-pro':
166+
case 'gemini-3-pro-preview':
167+
return {
168+
model: google(modelName),
169+
providerOptions: {
170+
thinkingConfig: {
171+
includeThoughts: request.thinkingConfig?.includeThoughts,
172+
},
173+
} satisfies GoogleGenerativeAIProviderOptions,
174+
};
157175
default:
158176
throw new Error(`Unexpected model in AI SDK runner: ${request.model}.`);
159177
}

0 commit comments

Comments
 (0)