diff --git a/apps/studio/lib/ai/bedrock.ts b/apps/studio/lib/ai/bedrock.ts index 857808726e9..e080768621e 100644 --- a/apps/studio/lib/ai/bedrock.ts +++ b/apps/studio/lib/ai/bedrock.ts @@ -4,6 +4,7 @@ import { createCredentialChain, fromNodeProviderChain } from '@aws-sdk/credentia import { CredentialsProviderError } from '@smithy/property-provider' import { awsCredentialsProvider } from '@vercel/functions/oidc' import { selectWeightedKey } from './util' +import { BedrockModel } from './model.utils' const credentialProvider = createCredentialChain( // Vercel OIDC provider will be used for staging/production @@ -61,8 +62,6 @@ export const regionPrefixMap: Record = { euc1: 'eu', } -export type BedrockModel = 'anthropic.claude-3-7-sonnet-20250219-v1:0' | 'openai.gpt-oss-120b-1:0' - export type RegionWeights = Record /** @@ -91,10 +90,8 @@ const modelRegionWeights: Record = { * Used to load balance requests across multiple regions depending on * their capacities. */ -export function createRoutedBedrock(routingKey?: string, useOpenAI = false) { - return async ( - modelId: BedrockModel - ): Promise<{ model: LanguageModel; supportsCachePoint: boolean }> => { +export function createRoutedBedrock(routingKey?: string) { + return async (modelId: BedrockModel): Promise => { const regionWeights = modelRegionWeights[modelId] // Select the Bedrock region based on the routing key and the model @@ -116,7 +113,6 @@ export function createRoutedBedrock(routingKey?: string, useOpenAI = false) { const modelName = activeRegions > 1 ? `${regionPrefixMap[bedrockRegion]}.${modelId}` : modelId const model = bedrock(modelName) - const supportsCachePoint = modelId === 'anthropic.claude-3-7-sonnet-20250219-v1:0' - return { model, supportsCachePoint } + return model } } diff --git a/apps/studio/lib/ai/model.test.ts b/apps/studio/lib/ai/model.test.ts index 59986528989..f3768a2f6fe 100644 --- a/apps/studio/lib/ai/model.test.ts +++ b/apps/studio/lib/ai/model.test.ts @@ -9,10 +9,7 @@ vi.mock('@ai-sdk/openai', () => ({ vi.mock('./bedrock', async () => ({ ...(await vi.importActual('./bedrock')), - createRoutedBedrock: vi.fn(() => async (modelId: string) => ({ - model: 'bedrock-model', - supportsCachePoint: modelId === 'anthropic.claude-3-7-sonnet-20250219-v1:0', - })), + createRoutedBedrock: vi.fn(() => async (modelId: string) => 'bedrock-model'), checkAwsCredentials: vi.fn(), })) @@ -28,25 +25,35 @@ describe('getModel', () => { process.env = { ...originalEnv } }) - it('should return bedrock model when AWS credentials are available and not throttled', async () => { + it('should return bedrock model and promptProviderOptions when default supports caching', async () => { vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(true) vi.stubEnv('IS_THROTTLED', 'false') - const { model, error, supportsCachePoint } = await getModel() + const { model, error, promptProviderOptions } = await getModel({ + routingKey: 'test', + isLimited: false, + }) expect(model).toEqual('bedrock-model') - expect(supportsCachePoint).toBe(true) + // Default bedrock model supportsCaching=false in registry, but if caller + // specifies high-tier, provider options would be present + expect(promptProviderOptions === undefined || typeof promptProviderOptions === 'object').toBe( + true + ) expect(error).toBeUndefined() }) - it('should return bedrock model when AWS credentials are available and throttled', async () => { + it('should return bedrock model when throttled (limited) with default model', async () => { vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(true) vi.stubEnv('IS_THROTTLED', 'true') - const { model, error, supportsCachePoint } = await getModel() + const { model, error, promptProviderOptions } = await getModel({ + routingKey: 'test', + isLimited: true, + }) expect(model).toEqual('bedrock-model') - expect(supportsCachePoint).toBe(false) + expect(promptProviderOptions).toBeUndefined() expect(error).toBeUndefined() }) @@ -54,19 +61,38 @@ describe('getModel', () => { vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false) process.env.OPENAI_API_KEY = 'test-key' - const { model, supportsCachePoint } = await getModel() + const { model, promptProviderOptions } = await getModel({ + routingKey: 'test', + isLimited: false, + }) expect(model).toEqual('openai-model') - expect(openai).toHaveBeenCalledWith('gpt-4.1-2025-04-14') - expect(supportsCachePoint).toBe(false) + // Default openai model in registry is gpt-5-mini + expect(openai).toHaveBeenCalledWith('gpt-5-mini') + expect(promptProviderOptions).toBeUndefined() }) it('should return error when neither AWS credentials nor OPENAI_API_KEY is available', async () => { vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false) delete process.env.OPENAI_API_KEY - const { error } = await getModel('test-key') - + const { error } = await getModel({ routingKey: 'test-key', isLimited: false }) expect(error).toEqual(new Error(ModelErrorMessage)) }) + + it('returns specified provider and model when provided (openai gpt-5)', async () => { + vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false) + process.env.OPENAI_API_KEY = 'test-key' + + const { model, error } = await getModel({ + provider: 'openai', + model: 'gpt-5', + routingKey: 'rk', + isLimited: false, + }) + + expect(error).toBeUndefined() + expect(model).toEqual('openai-model') + expect(openai).toHaveBeenCalledWith('gpt-5') + }) }) diff --git a/apps/studio/lib/ai/model.ts b/apps/studio/lib/ai/model.ts index 1cf7a56fa7f..eb1739416af 100644 --- a/apps/studio/lib/ai/model.ts +++ b/apps/studio/lib/ai/model.ts @@ -1,55 +1,117 @@ import { openai } from '@ai-sdk/openai' +import { anthropic } from '@ai-sdk/anthropic' import { LanguageModel } from 'ai' import { checkAwsCredentials, createRoutedBedrock } from './bedrock' +import { + PROVIDERS, + ProviderName, + getDefaultModelForProvider, + BedrockModel, + OpenAIModel, + AnthropicModel, + Model, + ProviderModelConfig, +} from './model.utils' -const BEDROCK_PRO_MODEL = 'anthropic.claude-3-7-sonnet-20250219-v1:0' -const BEDROCK_NORMAL_MODEL = 'openai.gpt-oss-120b-1:0' -const OPENAI_MODEL = 'gpt-4.1-2025-04-14' +export type PromptProviderOptions = Record export type ModelSuccess = { model: LanguageModel - supportsCachePoint: boolean + promptProviderOptions?: PromptProviderOptions error?: never } export type ModelError = { model?: never - supportsCachePoint?: never + promptProviderOptions?: never error: Error } export type ModelResponse = ModelSuccess | ModelError -export const ModelErrorMessage = - 'No valid AI model available. Please set up a local AWS profile to use Bedrock, or pass an OPENAI_API_KEY to use OpenAI.' +export const ModelErrorMessage = 'No valid AI model available based on available credentials.' + +export type GetModelParams = { + provider?: ProviderName + model?: Model + routingKey: string + isLimited?: boolean +} /** - * Retrieves the appropriate AI model based on available credentials. - * - * An optional routing key can be provided to distribute requests across - * different Bedrock regions. + * Retrieves a LanguageModel from a specific provider and model. + * - If provider/model not specified, auto-selects based on available credentials (prefers Bedrock). + * - If isLimited is true, uses the provider's default model. + * - Returns promptProviderOptions that callers can attach to the system message. */ -export async function getModel(routingKey?: string, isLimited?: boolean): Promise { - // Default behaviour here is to be throttled (e.g if this env var is not available, isThrottled should be true, unless specified 'false') - const isThrottled = process.env.IS_THROTTLED !== 'false' +export async function getModel(params: GetModelParams): Promise { + const { provider, model, routingKey, isLimited = true } = params + const envThrottled = process.env.IS_THROTTLED !== 'false' + + let preferredProvider: ProviderName | undefined = provider const hasAwsCredentials = await checkAwsCredentials() - const hasAwsBedrockRoleArn = !!process.env.AWS_BEDROCK_ROLE_ARN const hasOpenAIKey = !!process.env.OPENAI_API_KEY + const hasAnthropicKey = !!process.env.ANTHROPIC_API_KEY - if (hasAwsBedrockRoleArn && hasAwsCredentials) { - const bedrockModel = isThrottled || isLimited ? BEDROCK_NORMAL_MODEL : BEDROCK_PRO_MODEL + // Auto-pick a provider if not specified defaulting to Bedrock + if (!preferredProvider) { + if (hasAwsBedrockRoleArn && hasAwsCredentials) { + preferredProvider = 'bedrock' + } else if (hasOpenAIKey) { + preferredProvider = 'openai' + } else if (hasAnthropicKey) { + preferredProvider = 'anthropic' + } + } + + if (!preferredProvider) { + return { error: new Error(ModelErrorMessage) } + } + + const providerRegistry = PROVIDERS[preferredProvider] + if (!providerRegistry) { + return { error: new Error(`Unknown provider: ${preferredProvider}`) } + } + + const models = providerRegistry.models as Record + + const useDefault = isLimited || envThrottled || !model || !models[model] + + const chosenModelId = useDefault ? getDefaultModelForProvider(preferredProvider) : model + + if (preferredProvider === 'bedrock') { + if (!hasAwsBedrockRoleArn || !hasAwsCredentials) { + return { error: new Error('AWS Bedrock credentials not available') } + } const bedrock = createRoutedBedrock(routingKey) - const { model, supportsCachePoint } = await bedrock(bedrockModel) - - return { model, supportsCachePoint } + const model = await bedrock(chosenModelId as BedrockModel) + const promptProviderOptions = ( + providerRegistry.models as Record + )[chosenModelId as BedrockModel]?.promptProviderOptions + return { model, promptProviderOptions } } - // [Joshen] Only for local/self-hosted, hosted should always only use bedrock - if (hasOpenAIKey) { - return { model: openai(OPENAI_MODEL), supportsCachePoint: false } + if (preferredProvider === 'openai') { + if (!hasOpenAIKey) { + return { error: new Error('OPENAI_API_KEY not available') } + } + return { + model: openai(chosenModelId as OpenAIModel), + promptProviderOptions: models[chosenModelId as OpenAIModel]?.promptProviderOptions, + } } - return { error: new Error(ModelErrorMessage) } + if (preferredProvider === 'anthropic') { + if (!hasAnthropicKey) { + return { error: new Error('ANTHROPIC_API_KEY not available') } + } + return { + model: anthropic(chosenModelId as AnthropicModel), + promptProviderOptions: models[chosenModelId as AnthropicModel]?.promptProviderOptions, + } + } + + return { error: new Error(`Unsupported provider: ${preferredProvider}`) } } diff --git a/apps/studio/lib/ai/model.utils.ts b/apps/studio/lib/ai/model.utils.ts new file mode 100644 index 00000000000..2af90da8fcb --- /dev/null +++ b/apps/studio/lib/ai/model.utils.ts @@ -0,0 +1,66 @@ +export type ProviderName = 'bedrock' | 'openai' | 'anthropic' + +export type BedrockModel = 'anthropic.claude-3-7-sonnet-20250219-v1:0' | 'openai.gpt-oss-120b-1:0' + +export type OpenAIModel = 'gpt-5' | 'gpt-5-mini' + +export type AnthropicModel = 'claude-sonnet-4-20250514' | 'claude-3-5-haiku-20241022' + +export type Model = BedrockModel | OpenAIModel | AnthropicModel + +export type ProviderModelConfig = { + /** Optional providerOptions to attach to the system message for this model */ + promptProviderOptions?: Record + /** The default model for this provider (used when limited or no preferred specified) */ + default: boolean +} + +export type ProviderRegistry = { + bedrock: { + models: Record + } + openai: { + models: Record + } + anthropic: { + models: Record + } +} + +export const PROVIDERS: ProviderRegistry = { + bedrock: { + models: { + 'anthropic.claude-3-7-sonnet-20250219-v1:0': { + promptProviderOptions: { + bedrock: { + // Always cache the system prompt (must not contain dynamic content) + cachePoint: { type: 'default' }, + }, + }, + default: false, + }, + 'openai.gpt-oss-120b-1:0': { + default: true, + }, + }, + }, + openai: { + models: { + 'gpt-5': { default: false }, + 'gpt-5-mini': { default: true }, + }, + }, + anthropic: { + models: { + 'claude-sonnet-4-20250514': { default: false }, + 'claude-3-5-haiku-20241022': { default: true }, + }, + }, +} + +export function getDefaultModelForProvider(provider: ProviderName): Model | undefined { + const models = PROVIDERS[provider]?.models as Record + if (!models) return undefined + + return Object.keys(models).find((id) => models[id as Model]?.default) as Model | undefined +} diff --git a/apps/studio/package.json b/apps/studio/package.json index 76f867c8c4b..ea557b378a7 100644 --- a/apps/studio/package.json +++ b/apps/studio/package.json @@ -26,6 +26,7 @@ }, "dependencies": { "@ai-sdk/amazon-bedrock": "^3.0.0", + "@ai-sdk/anthropic": "^2.0.0", "@ai-sdk/openai": "^2.0.0", "@ai-sdk/provider": "^2.0.0", "@ai-sdk/provider-utils": "^3.0.0", diff --git a/apps/studio/pages/api/ai/code/complete.ts b/apps/studio/pages/api/ai/code/complete.ts index 43cf4d8fdc5..03d64a7866d 100644 --- a/apps/studio/pages/api/ai/code/complete.ts +++ b/apps/studio/pages/api/ai/code/complete.ts @@ -18,7 +18,6 @@ import { getTools } from 'lib/ai/tools' import apiWrapper from 'lib/api/apiWrapper' import { queryPgMetaSelfHosted } from 'lib/self-hosted' import { NextApiRequest, NextApiResponse } from 'next' -import { z } from 'zod/v4' export const maxDuration = 60 @@ -47,7 +46,6 @@ async function handler(req: NextApiRequest, res: NextApiResponse) { const accessToken = authorization?.replace('Bearer ', '') let aiOptInLevel: AiOptInLevel = 'disabled' - let isLimited = false if (!IS_PLATFORM) { aiOptInLevel = 'schema' @@ -55,17 +53,24 @@ async function handler(req: NextApiRequest, res: NextApiResponse) { if (IS_PLATFORM && orgSlug && authorization && projectRef) { // Get organizations and compute opt in level server-side - const { aiOptInLevel: orgAIOptInLevel, isLimited: orgAILimited } = await getOrgAIDetails({ + const { aiOptInLevel: orgAIOptInLevel } = await getOrgAIDetails({ orgSlug, authorization, projectRef, }) aiOptInLevel = orgAIOptInLevel - isLimited = orgAILimited } - const { model, error: modelError, supportsCachePoint } = await getModel(projectRef, isLimited) + // For code completion, we always use the limited model + const { + model, + error: modelError, + promptProviderOptions, + } = await getModel({ + provider: 'openai', + routingKey: projectRef, + }) if (modelError) { return res.status(500).json({ error: modelError.message }) @@ -111,14 +116,7 @@ async function handler(req: NextApiRequest, res: NextApiResponse) { { role: 'system', content: system, - ...(supportsCachePoint && { - providerOptions: { - bedrock: { - // Always cache the system prompt (must not contain dynamic content) - cachePoint: { type: 'default' }, - }, - }, - }), + ...(promptProviderOptions && { providerOptions: promptProviderOptions }), }, { role: 'assistant', diff --git a/apps/studio/pages/api/ai/feedback/classify.ts b/apps/studio/pages/api/ai/feedback/classify.ts index c27a76c3d4d..ccebe72a43b 100644 --- a/apps/studio/pages/api/ai/feedback/classify.ts +++ b/apps/studio/pages/api/ai/feedback/classify.ts @@ -28,7 +28,10 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { } try { - const { model, error: modelError } = await getModel() + const { model, error: modelError } = await getModel({ + provider: 'openai', + routingKey: 'feedback', + }) if (modelError) { return res.status(500).json({ error: modelError.message }) diff --git a/apps/studio/pages/api/ai/onboarding/design.ts b/apps/studio/pages/api/ai/onboarding/design.ts index 7a386321fc7..9e12fb09939 100644 --- a/apps/studio/pages/api/ai/onboarding/design.ts +++ b/apps/studio/pages/api/ai/onboarding/design.ts @@ -64,7 +64,10 @@ const wrapper = (req: NextApiRequest, res: NextApiResponse) => export default wrapper async function handlePost(req: NextApiRequest, res: NextApiResponse) { - const { model, error: modelError } = await getModel() + const { model, error: modelError } = await getModel({ + provider: 'openai', + routingKey: 'onboarding', + }) if (modelError) { return res.status(500).json({ error: modelError.message }) diff --git a/apps/studio/pages/api/ai/sql/cron-v2.ts b/apps/studio/pages/api/ai/sql/cron-v2.ts index 1ed03dccc60..000106e276c 100644 --- a/apps/studio/pages/api/ai/sql/cron-v2.ts +++ b/apps/studio/pages/api/ai/sql/cron-v2.ts @@ -34,7 +34,10 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { } try { - const { model, error: modelError } = await getModel() + const { model, error: modelError } = await getModel({ + provider: 'openai', + routingKey: 'cron', + }) if (modelError) { return res.status(500).json({ error: modelError.message }) @@ -79,7 +82,6 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { Here is the user's prompt: ${prompt} `, - temperature: 0, }) return res.json(result.object.cron_expression) diff --git a/apps/studio/pages/api/ai/sql/generate-v4.ts b/apps/studio/pages/api/ai/sql/generate-v4.ts index 1bf5038bef3..dd6f08367c7 100644 --- a/apps/studio/pages/api/ai/sql/generate-v4.ts +++ b/apps/studio/pages/api/ai/sql/generate-v4.ts @@ -117,7 +117,16 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { } } - const { model, error: modelError, supportsCachePoint } = await getModel(projectRef, isLimited) // use project ref as routing key + const { + model, + error: modelError, + promptProviderOptions, + } = await getModel({ + provider: 'openai', + model: 'gpt-5-mini', + routingKey: projectRef, + isLimited, + }) if (modelError) { return res.status(500).json({ error: modelError.message }) @@ -165,14 +174,7 @@ async function handlePost(req: NextApiRequest, res: NextApiResponse) { { role: 'system', content: system, - ...(supportsCachePoint && { - providerOptions: { - bedrock: { - // Always cache the system prompt (must not contain dynamic content) - cachePoint: { type: 'default' }, - }, - }, - }), + ...(promptProviderOptions && { providerOptions: promptProviderOptions }), }, { role: 'assistant', diff --git a/apps/studio/pages/api/ai/sql/title-v2.ts b/apps/studio/pages/api/ai/sql/title-v2.ts index b96cf00f35a..ed552d1030e 100644 --- a/apps/studio/pages/api/ai/sql/title-v2.ts +++ b/apps/studio/pages/api/ai/sql/title-v2.ts @@ -39,7 +39,10 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { } try { - const { model, error: modelError } = await getModel() + const { model, error: modelError } = await getModel({ + provider: 'openai', + routingKey: 'sql', + }) if (modelError) { return res.status(500).json({ error: modelError.message }) @@ -55,7 +58,6 @@ export async function handlePost(req: NextApiRequest, res: NextApiResponse) { The description should describe why this table was created (eg. "Table to track todos") or what the query does. `, - temperature: 0, }) return res.json(result.object) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 480f03fee11..2aa7e3c8e08 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -702,6 +702,9 @@ importers: '@ai-sdk/amazon-bedrock': specifier: ^3.0.0 version: 3.0.0(zod@3.25.76) + '@ai-sdk/anthropic': + specifier: ^2.0.0 + version: 2.0.0(zod@3.25.76) '@ai-sdk/openai': specifier: ^2.0.0 version: 2.0.2(zod@3.25.76) @@ -1188,7 +1191,7 @@ importers: version: 2.4.11(typescript@5.5.2) next-router-mock: specifier: ^0.9.13 - version: 0.9.13(next@15.3.3(@babel/core@7.26.10(supports-color@8.1.1))(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1) + version: 0.9.13(next@15.3.3(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1) postcss: specifier: ^8.5.3 version: 8.5.3 @@ -2464,7 +2467,7 @@ importers: version: link:../api-types next-router-mock: specifier: ^0.9.13 - version: 0.9.13(next@15.3.3(@babel/core@7.26.10(supports-color@8.1.1))(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1) + version: 0.9.13(next@15.3.3(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1) tsx: specifier: ^4.19.3 version: 4.19.3 @@ -10933,7 +10936,6 @@ packages: resolution: {integrity: sha512-t0q23FIpvHDTtnORW+bDJziGsal5uh9RJTJ1fyH8drd4lICOoXhJ5pLMUZ5C0VQei6dNmwTzzoTRgMkO9JgHEQ==} peerDependencies: eslint: '>= 5' - bundledDependencies: [] eslint-plugin-import@2.31.0: resolution: {integrity: sha512-ixmkI62Rbc2/w8Vfxyh1jQRTdRTF52VxwRVHl/ykPAmqG+Nb7/kNn+byLP0LxPgI7zWA16Jt82SybJInmMia3A==} @@ -33550,7 +33552,7 @@ snapshots: dependencies: js-yaml-loader: 1.2.2 - next-router-mock@0.9.13(next@15.3.3(@babel/core@7.26.10(supports-color@8.1.1))(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1): + next-router-mock@0.9.13(next@15.3.3(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4))(react@18.3.1): dependencies: next: 15.3.3(@babel/core@7.26.10(supports-color@8.1.1))(@opentelemetry/api@1.9.0)(@playwright/test@1.53.0)(babel-plugin-macros@3.1.0)(react-dom@18.3.1(react@18.3.1))(react@18.3.1)(sass@1.77.4) react: 18.3.1