mirror of
https://github.com/supabase/supabase.git
synced 2026-10-05 09:25:06 +03:00
Saxon FletcherandJoshen Lim
bf004b8a1d
* choose provider and model * fix test * Fix lint * Nit clean up * add provider options and remove anthropic * remove package * use chat model and update prompt * prompt rendering * default to pro * fix test * low * remove log * remove redundant part.text --------- Co-authored-by: Joshen Lim <joshenlimek@gmail.com>
100 lines
3.3 KiB
TypeScript
100 lines
3.3 KiB
TypeScript
import { openai } from '@ai-sdk/openai'
|
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
|
|
import * as bedrockModule from './bedrock'
|
|
import { getModel, ModelErrorMessage } from './model'
|
|
|
|
vi.mock('@ai-sdk/openai', () => ({
|
|
openai: { chat: vi.fn(() => 'openai-model') },
|
|
}))
|
|
|
|
vi.mock('./bedrock', async () => ({
|
|
...(await vi.importActual('./bedrock')),
|
|
createRoutedBedrock: vi.fn(() => async (modelId: string) => 'bedrock-model'),
|
|
checkAwsCredentials: vi.fn(),
|
|
}))
|
|
|
|
describe('getModel', () => {
|
|
const originalEnv = { ...process.env }
|
|
|
|
beforeEach(() => {
|
|
vi.resetAllMocks()
|
|
vi.stubEnv('AWS_BEDROCK_ROLE_ARN', 'test')
|
|
})
|
|
|
|
afterEach(() => {
|
|
process.env = { ...originalEnv }
|
|
})
|
|
|
|
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, promptProviderOptions } = await getModel({
|
|
routingKey: 'test',
|
|
isLimited: false,
|
|
})
|
|
|
|
expect(model).toEqual('bedrock-model')
|
|
// 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 throttled (limited) with default model', async () => {
|
|
vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(true)
|
|
vi.stubEnv('IS_THROTTLED', 'true')
|
|
|
|
const { model, error, promptProviderOptions } = await getModel({
|
|
routingKey: 'test',
|
|
isLimited: true,
|
|
})
|
|
|
|
expect(model).toEqual('bedrock-model')
|
|
expect(promptProviderOptions).toBeUndefined()
|
|
expect(error).toBeUndefined()
|
|
})
|
|
|
|
it('should return OpenAI model when AWS credentials are not available but OPENAI_API_KEY is set', async () => {
|
|
vi.mocked(bedrockModule.checkAwsCredentials).mockResolvedValue(false)
|
|
process.env.OPENAI_API_KEY = 'test-key'
|
|
|
|
const { model, promptProviderOptions } = await getModel({
|
|
routingKey: 'test',
|
|
isLimited: false,
|
|
})
|
|
|
|
expect(model).toEqual('openai-model')
|
|
// Default openai model in registry is gpt-5-mini
|
|
expect(openai.chat).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({ 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'
|
|
process.env.IS_THROTTLED = 'false'
|
|
|
|
const { model, error } = await getModel({
|
|
provider: 'openai',
|
|
model: 'gpt-5',
|
|
routingKey: 'rk',
|
|
isLimited: false,
|
|
})
|
|
|
|
expect(error).toBeUndefined()
|
|
expect(model).toEqual('openai-model')
|
|
expect(openai.chat).toHaveBeenCalledWith('gpt-5')
|
|
})
|
|
})
|