Skip to content

Commit 7b02377

Browse files
committed
feat(providers): add Together AI, Baseten, and Ollama Cloud model providers
1 parent 919fa52 commit 7b02377

39 files changed

Lines changed: 4411 additions & 589 deletions

File tree

apps/sim/app/(landing)/models/utils.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ const PROVIDER_PREFIXES: Record<string, string[]> = {
88
bedrock: ['bedrock/'],
99
cerebras: ['cerebras/'],
1010
fireworks: ['fireworks/'],
11+
together: ['together/'],
12+
baseten: ['baseten/'],
13+
'ollama-cloud': ['ollama-cloud/'],
1114
groq: ['groq/'],
1215
openrouter: ['openrouter/'],
1316
vllm: ['vllm/'],
Lines changed: 252 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
/**
2+
* @vitest-environment node
3+
*/
4+
import { createMockRequest } from '@sim/testing'
5+
import { beforeEach, describe, expect, it, vi } from 'vitest'
6+
7+
const {
8+
mockFilterBlacklistedModels,
9+
mockIsProviderBlacklisted,
10+
mockGetBYOKKey,
11+
mockGetSession,
12+
mockGetUserEntityPermissions,
13+
mutableEnv,
14+
} = vi.hoisted(() => ({
15+
mockFilterBlacklistedModels: vi.fn(),
16+
mockIsProviderBlacklisted: vi.fn(),
17+
mockGetBYOKKey: vi.fn(),
18+
mockGetSession: vi.fn(),
19+
mockGetUserEntityPermissions: vi.fn(),
20+
mutableEnv: { BASETEN_API_KEY: undefined as string | undefined },
21+
}))
22+
23+
vi.mock('@/lib/core/config/env', () => ({ env: mutableEnv }))
24+
25+
vi.mock('@/providers/utils', () => ({
26+
filterBlacklistedModels: mockFilterBlacklistedModels,
27+
isProviderBlacklisted: mockIsProviderBlacklisted,
28+
}))
29+
30+
vi.mock('@/lib/api-key/byok', () => ({
31+
getBYOKKey: mockGetBYOKKey,
32+
}))
33+
34+
vi.mock('@/lib/auth', () => ({
35+
getSession: mockGetSession,
36+
}))
37+
38+
vi.mock('@/lib/workspaces/permissions/utils', () => ({
39+
getUserEntityPermissions: mockGetUserEntityPermissions,
40+
}))
41+
42+
import { GET } from '@/app/api/providers/baseten/models/route'
43+
44+
const BASETEN_MODELS_URL = 'https://inference.baseten.co/v1/models'
45+
46+
function jsonResponse(body: unknown, init: { ok?: boolean; status?: number } = {}): Response {
47+
const status = init.status ?? 200
48+
const ok = init.ok ?? (status >= 200 && status < 300)
49+
return {
50+
ok,
51+
status,
52+
statusText: ok ? 'OK' : 'Error',
53+
json: vi.fn(async () => body),
54+
} as unknown as Response
55+
}
56+
57+
function setEnvKey(value: string | undefined): void {
58+
mutableEnv.BASETEN_API_KEY = value
59+
}
60+
61+
function authHeaderFromLastFetch(mockFetch: ReturnType<typeof vi.fn>): unknown {
62+
const init = mockFetch.mock.calls.at(-1)?.[1] as RequestInit | undefined
63+
return (init?.headers as Record<string, string> | undefined)?.Authorization
64+
}
65+
66+
describe('GET /api/providers/baseten/models', () => {
67+
let mockFetch: ReturnType<typeof vi.fn>
68+
69+
beforeEach(() => {
70+
vi.clearAllMocks()
71+
72+
mockFetch = vi.fn()
73+
vi.stubGlobal('fetch', mockFetch)
74+
75+
mockIsProviderBlacklisted.mockReturnValue(false)
76+
mockFilterBlacklistedModels.mockImplementation((models: string[]) => models)
77+
mockGetBYOKKey.mockResolvedValue(null)
78+
mockGetSession.mockResolvedValue(null)
79+
mockGetUserEntityPermissions.mockResolvedValue(null)
80+
setEnvKey(undefined)
81+
})
82+
83+
it('returns empty models without fetching when the provider is blacklisted', async () => {
84+
mockIsProviderBlacklisted.mockReturnValue(true)
85+
setEnvKey('env-key')
86+
87+
const res = await GET(createMockRequest('GET'))
88+
89+
expect(res.status).toBe(200)
90+
expect(await res.json()).toEqual({ models: [] })
91+
expect(mockFetch).not.toHaveBeenCalled()
92+
})
93+
94+
it('returns empty models when no workspaceId and no env key are available', async () => {
95+
const res = await GET(createMockRequest('GET'))
96+
97+
expect(res.status).toBe(200)
98+
expect(await res.json()).toEqual({ models: [] })
99+
expect(mockFetch).not.toHaveBeenCalled()
100+
})
101+
102+
it('fetches models with the env key and prefixes each id with baseten/', async () => {
103+
setEnvKey('env-key')
104+
mockFetch.mockResolvedValueOnce(
105+
jsonResponse({
106+
data: [{ id: 'openai/gpt-oss-120b' }, { id: 'deepseek-ai/DeepSeek-V3' }],
107+
})
108+
)
109+
110+
const res = await GET(createMockRequest('GET'))
111+
112+
expect(res.status).toBe(200)
113+
expect(await res.json()).toEqual({
114+
models: ['baseten/openai/gpt-oss-120b', 'baseten/deepseek-ai/DeepSeek-V3'],
115+
})
116+
117+
expect(mockFetch).toHaveBeenCalledTimes(1)
118+
const [url, init] = mockFetch.mock.calls[0]
119+
expect(url).toBe(BASETEN_MODELS_URL)
120+
expect((init.headers as Record<string, string>).Authorization).toBe('Bearer env-key')
121+
})
122+
123+
it('uses the BYOK key when workspaceId, session, and permission are present', async () => {
124+
setEnvKey('env-key')
125+
mockGetSession.mockResolvedValue({ user: { id: 'user-1' } })
126+
mockGetUserEntityPermissions.mockResolvedValue('admin')
127+
mockGetBYOKKey.mockResolvedValue({ apiKey: 'byok-key', isBYOK: true })
128+
mockFetch.mockResolvedValueOnce(jsonResponse({ data: [{ id: 'model-a' }] }))
129+
130+
const res = await GET(
131+
createMockRequest('GET', undefined, {}, 'http://localhost:3000/api/test?workspaceId=ws-1')
132+
)
133+
134+
expect(res.status).toBe(200)
135+
expect(await res.json()).toEqual({ models: ['baseten/model-a'] })
136+
137+
expect(mockGetBYOKKey).toHaveBeenCalledWith('ws-1', 'baseten')
138+
expect(authHeaderFromLastFetch(mockFetch)).toBe('Bearer byok-key')
139+
})
140+
141+
it('falls back to the env key when there is a workspaceId but no session', async () => {
142+
setEnvKey('env-key')
143+
mockGetSession.mockResolvedValue(null)
144+
mockFetch.mockResolvedValueOnce(jsonResponse({ data: [{ id: 'model-a' }] }))
145+
146+
const res = await GET(
147+
createMockRequest('GET', undefined, {}, 'http://localhost:3000/api/test?workspaceId=ws-1')
148+
)
149+
150+
expect(res.status).toBe(200)
151+
expect(await res.json()).toEqual({ models: ['baseten/model-a'] })
152+
expect(mockGetBYOKKey).not.toHaveBeenCalled()
153+
expect(authHeaderFromLastFetch(mockFetch)).toBe('Bearer env-key')
154+
})
155+
156+
it('falls back to the env key when the user lacks workspace permission', async () => {
157+
setEnvKey('env-key')
158+
mockGetSession.mockResolvedValue({ user: { id: 'user-1' } })
159+
mockGetUserEntityPermissions.mockResolvedValue(null)
160+
mockFetch.mockResolvedValueOnce(jsonResponse({ data: [{ id: 'model-a' }] }))
161+
162+
const res = await GET(
163+
createMockRequest('GET', undefined, {}, 'http://localhost:3000/api/test?workspaceId=ws-1')
164+
)
165+
166+
expect(res.status).toBe(200)
167+
expect(await res.json()).toEqual({ models: ['baseten/model-a'] })
168+
expect(mockGetBYOKKey).not.toHaveBeenCalled()
169+
expect(authHeaderFromLastFetch(mockFetch)).toBe('Bearer env-key')
170+
})
171+
172+
it('returns empty models when the upstream responds 401', async () => {
173+
setEnvKey('env-key')
174+
mockFetch.mockResolvedValueOnce(jsonResponse({}, { ok: false, status: 401 }))
175+
176+
const res = await GET(createMockRequest('GET'))
177+
178+
expect(res.status).toBe(200)
179+
expect(await res.json()).toEqual({ models: [] })
180+
})
181+
182+
it('returns empty models when the upstream responds 500', async () => {
183+
setEnvKey('env-key')
184+
mockFetch.mockResolvedValueOnce(jsonResponse({}, { ok: false, status: 500 }))
185+
186+
const res = await GET(createMockRequest('GET'))
187+
188+
expect(res.status).toBe(200)
189+
expect(await res.json()).toEqual({ models: [] })
190+
})
191+
192+
it('returns empty models when fetch throws', async () => {
193+
setEnvKey('env-key')
194+
mockFetch.mockRejectedValueOnce(new Error('network down'))
195+
196+
const res = await GET(createMockRequest('GET'))
197+
198+
expect(res.status).toBe(200)
199+
expect(await res.json()).toEqual({ models: [] })
200+
})
201+
202+
it('returns empty models when the upstream data array is empty', async () => {
203+
setEnvKey('env-key')
204+
mockFetch.mockResolvedValueOnce(jsonResponse({ data: [] }))
205+
206+
const res = await GET(createMockRequest('GET'))
207+
208+
expect(res.status).toBe(200)
209+
expect(await res.json()).toEqual({ models: [] })
210+
})
211+
212+
it('returns empty models when the upstream omits the data field', async () => {
213+
setEnvKey('env-key')
214+
mockFetch.mockResolvedValueOnce(jsonResponse({ object: 'list' }))
215+
216+
const res = await GET(createMockRequest('GET'))
217+
218+
expect(res.status).toBe(200)
219+
expect(await res.json()).toEqual({ models: [] })
220+
})
221+
222+
it('dedupes repeated model ids', async () => {
223+
setEnvKey('env-key')
224+
mockFetch.mockResolvedValueOnce(
225+
jsonResponse({
226+
data: [{ id: 'model-a' }, { id: 'model-a' }, { id: 'model-b' }],
227+
})
228+
)
229+
230+
const res = await GET(createMockRequest('GET'))
231+
232+
expect(res.status).toBe(200)
233+
expect(await res.json()).toEqual({ models: ['baseten/model-a', 'baseten/model-b'] })
234+
})
235+
236+
it('drops models removed by the blacklist filter', async () => {
237+
setEnvKey('env-key')
238+
mockFilterBlacklistedModels.mockImplementation((models: string[]) =>
239+
models.filter((m) => m !== 'baseten/blocked-model')
240+
)
241+
mockFetch.mockResolvedValueOnce(
242+
jsonResponse({
243+
data: [{ id: 'allowed-model' }, { id: 'blocked-model' }],
244+
})
245+
)
246+
247+
const res = await GET(createMockRequest('GET'))
248+
249+
expect(res.status).toBe(200)
250+
expect(await res.json()).toEqual({ models: ['baseten/allowed-model'] })
251+
})
252+
})
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
import { createLogger } from '@sim/logger'
2+
import { getErrorMessage } from '@sim/utils/errors'
3+
import { type NextRequest, NextResponse } from 'next/server'
4+
import {
5+
basetenProviderModelsQuerySchema,
6+
basetenUpstreamResponseSchema,
7+
providerModelsResponseSchema,
8+
} from '@/lib/api/contracts/providers'
9+
import { validationErrorResponse } from '@/lib/api/server'
10+
import { getBYOKKey } from '@/lib/api-key/byok'
11+
import { getSession } from '@/lib/auth'
12+
import { env } from '@/lib/core/config/env'
13+
import { withRouteHandler } from '@/lib/core/utils/with-route-handler'
14+
import { getUserEntityPermissions } from '@/lib/workspaces/permissions/utils'
15+
import { filterBlacklistedModels, isProviderBlacklisted } from '@/providers/utils'
16+
17+
const logger = createLogger('BasetenModelsAPI')
18+
19+
export const GET = withRouteHandler(async (request: NextRequest) => {
20+
if (isProviderBlacklisted('baseten')) {
21+
logger.info('Baseten provider is blacklisted, returning empty models')
22+
return NextResponse.json({ models: [] })
23+
}
24+
25+
let apiKey: string | undefined
26+
27+
const queryValidation = basetenProviderModelsQuerySchema.safeParse({
28+
workspaceId: request.nextUrl.searchParams.get('workspaceId') ?? undefined,
29+
})
30+
if (!queryValidation.success) return validationErrorResponse(queryValidation.error)
31+
const { workspaceId } = queryValidation.data
32+
if (workspaceId) {
33+
const session = await getSession()
34+
if (session?.user?.id) {
35+
const permission = await getUserEntityPermissions(session.user.id, 'workspace', workspaceId)
36+
if (permission) {
37+
const byokResult = await getBYOKKey(workspaceId, 'baseten')
38+
if (byokResult) {
39+
apiKey = byokResult.apiKey
40+
}
41+
}
42+
}
43+
}
44+
45+
if (!apiKey) {
46+
apiKey = env.BASETEN_API_KEY
47+
}
48+
49+
if (!apiKey) {
50+
logger.info('No Baseten API key available, returning empty models')
51+
return NextResponse.json({ models: [] })
52+
}
53+
54+
try {
55+
const response = await fetch('https://inference.baseten.co/v1/models', {
56+
headers: {
57+
Authorization: `Bearer ${apiKey}`,
58+
'Content-Type': 'application/json',
59+
},
60+
cache: 'no-store',
61+
})
62+
63+
if (!response.ok) {
64+
logger.warn('Failed to fetch Baseten models', {
65+
status: response.status,
66+
statusText: response.statusText,
67+
})
68+
return NextResponse.json({ models: [] })
69+
}
70+
71+
const data = basetenUpstreamResponseSchema.parse(await response.json())
72+
73+
const allModels: string[] = []
74+
for (const model of data.data ?? []) {
75+
allModels.push(`baseten/${model.id}`)
76+
}
77+
78+
const uniqueModels = Array.from(new Set(allModels))
79+
const models = filterBlacklistedModels(uniqueModels)
80+
81+
logger.info('Successfully fetched Baseten models', {
82+
count: models.length,
83+
filtered: uniqueModels.length - models.length,
84+
})
85+
86+
return NextResponse.json(providerModelsResponseSchema.parse({ models }))
87+
} catch (error) {
88+
logger.error('Error fetching Baseten models', {
89+
error: getErrorMessage(error, 'Unknown error'),
90+
})
91+
return NextResponse.json({ models: [] })
92+
}
93+
})

0 commit comments

Comments
 (0)