choose provider and model

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<BedrockRegion, string> = {
   euc1: 'eu',
 }
 
-export type BedrockModel = 'anthropic.claude-3-7-sonnet-20250219-v1:0' | 'openai.gpt-oss-120b-1:0'
-
 export type RegionWeights = Record<BedrockRegion, number>
 
 /**
@@ -91,10 +90,8 @@ const modelRegionWeights: Record<BedrockModel, RegionWeights> = {
  * 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<LanguageModel> => {
     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
   }
 }

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')
+  })
 })

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<string, any>
 
 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<ModelResponse> {
-  // 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<ModelResponse> {
+  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<Model, ProviderModelConfig>
+
+  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<BedrockModel, ProviderModelConfig>
+    )[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}`) }
 }

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<string, any>
+  /** The default model for this provider (used when limited or no preferred specified) */
+  default: boolean
+}
+
+export type ProviderRegistry = {
+  bedrock: {
+    models: Record<BedrockModel, ProviderModelConfig>
+  }
+  openai: {
+    models: Record<OpenAIModel, ProviderModelConfig>
+  }
+  anthropic: {
+    models: Record<AnthropicModel, ProviderModelConfig>
+  }
+}
+
+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<Model, ProviderModelConfig>
+  if (!models) return undefined
+
+  return Object.keys(models).find((id) => models[id as Model]?.default) as Model | undefined
+}

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",

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',

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 })

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 })

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)

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',

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)

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