Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 78 additions & 2 deletions apps/sim/providers/utils.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ import {
MODELS_WITH_VERBOSITY,
PROVIDERS_WITH_TOOL_USAGE_CONTROL,
prepareToolsWithUsageControl,
shouldBillModelUsage,
supportsTemperature,
supportsToolUsageControl,
transformCustomTool,
Expand All @@ -40,6 +41,7 @@ describe('getApiKey', () => {
beforeEach(() => {
vi.clearAllMocks()

// @ts-expect-error - mocking boolean with different value
isHostedSpy.mockReturnValue(false)

module.require = vi.fn(() => ({
Expand All @@ -53,6 +55,7 @@ describe('getApiKey', () => {
})

it('should return user-provided key when not in hosted environment', () => {
// @ts-expect-error - mocking boolean with different value
isHostedSpy.mockReturnValue(false)

// For OpenAI
Expand All @@ -65,6 +68,7 @@ describe('getApiKey', () => {
})

it('should throw error if no key provided in non-hosted environment', () => {
// @ts-expect-error - mocking boolean with different value
isHostedSpy.mockReturnValue(false)

expect(() => getApiKey('openai', 'gpt-4')).toThrow('API key is required for openai gpt-4')
Expand All @@ -80,7 +84,8 @@ describe('getApiKey', () => {
throw new Error('Rotation failed')
})

const key = getApiKey('openai', 'gpt-4', 'user-fallback-key')
// Use gpt-4o which IS in the hosted models list
const key = getApiKey('openai', 'gpt-4o', 'user-fallback-key')
expect(key).toBe('user-fallback-key')
})

Expand All @@ -91,7 +96,8 @@ describe('getApiKey', () => {
throw new Error('Rotation failed')
})

expect(() => getApiKey('openai', 'gpt-4')).toThrow('No API key available for openai gpt-4')
// Use gpt-4o which IS in the hosted models list
expect(() => getApiKey('openai', 'gpt-4o')).toThrow('No API key available for openai gpt-4o')
})

it('should require user key for non-OpenAI/Anthropic providers even in hosted environment', () => {
Expand All @@ -104,6 +110,30 @@ describe('getApiKey', () => {
'API key is required for other-provider some-model'
)
})

it('should require user key for models NOT in hosted list even if provider matches', () => {
isHostedSpy.mockReturnValue(true)

// Models with version suffixes that are NOT in the hosted list should require user API key
// even though they're from anthropic/openai providers

// User provides their own key - should work
const key1 = getApiKey('anthropic', 'claude-sonnet-4-20250514', 'user-key-anthropic')
expect(key1).toBe('user-key-anthropic')

// No user key - should throw, NOT use server key
expect(() => getApiKey('anthropic', 'claude-sonnet-4-20250514')).toThrow(
'API key is required for anthropic claude-sonnet-4-20250514'
)

// Same for OpenAI versioned models not in list
const key2 = getApiKey('openai', 'gpt-4o-2024-08-06', 'user-key-openai')
expect(key2).toBe('user-key-openai')

expect(() => getApiKey('openai', 'gpt-4o-2024-08-06')).toThrow(
'API key is required for openai gpt-4o-2024-08-06'
)
})
})

describe('Model Capabilities', () => {
Expand Down Expand Up @@ -476,6 +506,52 @@ describe('getHostedModels', () => {
})
})

describe('shouldBillModelUsage', () => {
it.concurrent('should return true for exact matches of hosted models', () => {
// OpenAI models
expect(shouldBillModelUsage('gpt-4o')).toBe(true)
expect(shouldBillModelUsage('o1')).toBe(true)

// Anthropic models
expect(shouldBillModelUsage('claude-sonnet-4-0')).toBe(true)
expect(shouldBillModelUsage('claude-opus-4-0')).toBe(true)

// Google models
expect(shouldBillModelUsage('gemini-2.5-pro')).toBe(true)
expect(shouldBillModelUsage('gemini-2.5-flash')).toBe(true)
})

it.concurrent('should return false for non-hosted models', () => {
// Other providers
expect(shouldBillModelUsage('deepseek-v3')).toBe(false)
expect(shouldBillModelUsage('grok-4-latest')).toBe(false)

// Unknown models
expect(shouldBillModelUsage('unknown-model')).toBe(false)
})

it.concurrent('should return false for versioned model names not in hosted list', () => {
// Versioned model names that are NOT in the hosted list
// These should NOT be billed (user provides own API key)
expect(shouldBillModelUsage('claude-sonnet-4-20250514')).toBe(false)
expect(shouldBillModelUsage('gpt-4o-2024-08-06')).toBe(false)
expect(shouldBillModelUsage('claude-3-5-sonnet-20241022')).toBe(false)
})

it.concurrent('should be case insensitive', () => {
expect(shouldBillModelUsage('GPT-4O')).toBe(true)
expect(shouldBillModelUsage('Claude-Sonnet-4-0')).toBe(true)
expect(shouldBillModelUsage('GEMINI-2.5-PRO')).toBe(true)
})

it.concurrent('should not match partial model names', () => {
// Should not match partial/prefix models
expect(shouldBillModelUsage('gpt-4')).toBe(false) // gpt-4o is hosted, not gpt-4
expect(shouldBillModelUsage('claude-sonnet')).toBe(false)
expect(shouldBillModelUsage('gemini')).toBe(false)
})
})

describe('Provider Management', () => {
describe('getProviderFromModel', () => {
it.concurrent('should return correct provider for known models', () => {
Expand Down
29 changes: 16 additions & 13 deletions apps/sim/providers/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -619,7 +619,7 @@ export function getHostedModels(): string[] {
*/
export function shouldBillModelUsage(model: string): boolean {
const hostedModels = getHostedModels()
return hostedModels.includes(model)
return hostedModels.some((hostedModel) => model.toLowerCase() === hostedModel.toLowerCase())
}

/**
Expand All @@ -643,19 +643,22 @@ export function getApiKey(provider: string, model: string, userProvidedKey?: str
const isGeminiModel = provider === 'google'

if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel)) {
try {
// Import the key rotation function
const { getRotatingApiKey } = require('@/lib/core/config/api-keys')
const serverKey = getRotatingApiKey(isGeminiModel ? 'gemini' : provider)
return serverKey
} catch (_error) {
// If server key fails and we have a user key, fallback to that
if (hasUserKey) {
return userProvidedKey!
}
// Only use server key if model is explicitly in our hosted list
const hostedModels = getHostedModels()
const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())

if (isModelHosted) {
try {
const { getRotatingApiKey } = require('@/lib/core/config/api-keys')
const serverKey = getRotatingApiKey(isGeminiModel ? 'gemini' : provider)
return serverKey
} catch (_error) {
if (hasUserKey) {
return userProvidedKey!
}

// Otherwise, throw an error
throw new Error(`No API key available for ${provider} ${model}`)
throw new Error(`No API key available for ${provider} ${model}`)
}
}
}

Expand Down