mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-24 15:45:35 +08:00
feat(providers): add xAI to hosted key rotation pool (#5574)
* feat(providers): add xAI to hosted key rotation pool Wires xai into the same hosted-key mechanism as OpenAI, Anthropic, and Z.ai so Sim can serve Grok models without users bringing their own key. * fix(pi): include xai in Pi cloud-mode workspace BYOK read-back xai was fully wired as a Pi-supported provider but missing from WORKSPACE_BYOK_PROVIDERS, so a stored workspace xAI key was never read back for cloud-mode Pi runs. * fix(byok): add xai settings UI row xai is both hosted (Pi block hides its inline API key field for hosted models) and Pi-supported (cloud mode requires a user key), so without a Settings > BYOK row users had no way to supply an xai key for Pi cloud runs.
This commit is contained in:
@@ -34,6 +34,7 @@ import {
|
||||
SerperIcon,
|
||||
TogetherIcon,
|
||||
WizaIcon,
|
||||
xAIIcon,
|
||||
ZeroBounceIcon,
|
||||
} from '@/components/icons'
|
||||
import { MAX_BYOK_KEYS_PER_PROVIDER } from '@/lib/api/contracts/byok-keys'
|
||||
@@ -76,6 +77,13 @@ const PROVIDERS: (BYOKManagerProvider & { id: BYOKProviderId })[] = [
|
||||
description: 'LLM calls and Knowledge Base OCR',
|
||||
placeholder: 'Enter your API key',
|
||||
},
|
||||
{
|
||||
id: 'xai',
|
||||
name: 'xAI',
|
||||
icon: xAIIcon,
|
||||
description: 'LLM calls',
|
||||
placeholder: 'xai-...',
|
||||
},
|
||||
{
|
||||
id: 'fireworks',
|
||||
name: 'Fireworks',
|
||||
@@ -287,6 +295,7 @@ const PROVIDER_SECTIONS: BYOKProviderSection[] = [
|
||||
'anthropic',
|
||||
'google',
|
||||
'mistral',
|
||||
'xai',
|
||||
'fireworks',
|
||||
'together',
|
||||
'baseten',
|
||||
|
||||
@@ -134,6 +134,21 @@ describe('resolvePiModelKey', () => {
|
||||
expect(mockGetApiKeyWithBYOK).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('cloud mode falls back to a stored workspace key for xAI', async () => {
|
||||
mockGetProviderFromModel.mockReturnValue('xai')
|
||||
mockGetBYOKKey.mockResolvedValue({ apiKey: 'xai-workspace-key', isBYOK: true })
|
||||
|
||||
const result = await resolvePiModelKey({
|
||||
model: 'grok-4.5',
|
||||
mode: 'cloud',
|
||||
workspaceId: 'ws-1',
|
||||
})
|
||||
|
||||
expect(result).toEqual({ providerId: 'xai', apiKey: 'xai-workspace-key', isBYOK: true })
|
||||
expect(mockGetBYOKKey).toHaveBeenCalledWith('ws-1', 'xai')
|
||||
expect(mockGetApiKeyWithBYOK).not.toHaveBeenCalled()
|
||||
})
|
||||
|
||||
it('cloud mode rejects when no user key is available (never a hosted key)', async () => {
|
||||
mockGetProviderFromModel.mockReturnValue('anthropic')
|
||||
mockGetBYOKKey.mockResolvedValue(null)
|
||||
|
||||
@@ -33,7 +33,13 @@ interface ResolvePiModelKeyParams {
|
||||
}
|
||||
|
||||
/** Providers whose key Sim can store as a workspace BYOK key (read back for cloud). */
|
||||
const WORKSPACE_BYOK_PROVIDERS = new Set<string>(['anthropic', 'openai', 'google', 'mistral'])
|
||||
const WORKSPACE_BYOK_PROVIDERS = new Set<string>([
|
||||
'anthropic',
|
||||
'openai',
|
||||
'google',
|
||||
'mistral',
|
||||
'xai',
|
||||
])
|
||||
|
||||
/** Resolves the provider and a usable API key for the selected model. */
|
||||
export async function resolvePiModelKey(params: ResolvePiModelKeyParams): Promise<PiKeyResolution> {
|
||||
|
||||
@@ -205,13 +205,14 @@ export async function getApiKeyWithBYOK(
|
||||
const isGeminiModel = provider === 'google'
|
||||
const isMistralModel = provider === 'mistral'
|
||||
const isZaiModel = provider === 'zai'
|
||||
const isXaiModel = provider === 'xai'
|
||||
|
||||
const byokProviderId = isGeminiModel ? 'google' : (provider as BYOKProviderId)
|
||||
|
||||
if (
|
||||
isHosted &&
|
||||
workspaceId &&
|
||||
(isOpenAIModel || isClaudeModel || isGeminiModel || isMistralModel || isZaiModel)
|
||||
(isOpenAIModel || isClaudeModel || isGeminiModel || isMistralModel || isZaiModel || isXaiModel)
|
||||
) {
|
||||
const hostedModels = getHostedModels()
|
||||
const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())
|
||||
|
||||
@@ -7,6 +7,7 @@ export const byokProviderIdSchema = z.enum([
|
||||
'google',
|
||||
'mistral',
|
||||
'zai',
|
||||
'xai',
|
||||
'fireworks',
|
||||
'together',
|
||||
'baseten',
|
||||
|
||||
@@ -12,7 +12,8 @@ export function getRotatingApiKey(provider: string): string {
|
||||
provider !== 'anthropic' &&
|
||||
provider !== 'gemini' &&
|
||||
provider !== 'cohere' &&
|
||||
provider !== 'zai'
|
||||
provider !== 'zai' &&
|
||||
provider !== 'xai'
|
||||
) {
|
||||
throw new Error(`No rotation implemented for provider: ${provider}`)
|
||||
}
|
||||
@@ -39,6 +40,10 @@ export function getRotatingApiKey(provider: string): string {
|
||||
if (env.ZAI_API_KEY_1) keys.push(env.ZAI_API_KEY_1)
|
||||
if (env.ZAI_API_KEY_2) keys.push(env.ZAI_API_KEY_2)
|
||||
if (env.ZAI_API_KEY_3) keys.push(env.ZAI_API_KEY_3)
|
||||
} else if (provider === 'xai') {
|
||||
if (env.XAI_API_KEY_1) keys.push(env.XAI_API_KEY_1)
|
||||
if (env.XAI_API_KEY_2) keys.push(env.XAI_API_KEY_2)
|
||||
if (env.XAI_API_KEY_3) keys.push(env.XAI_API_KEY_3)
|
||||
}
|
||||
|
||||
if (keys.length === 0) {
|
||||
|
||||
@@ -147,6 +147,9 @@ export const env = createEnv({
|
||||
ZAI_API_KEY_1: z.string().min(1).optional(), // Primary Z.ai API key for load balancing
|
||||
ZAI_API_KEY_2: z.string().min(1).optional(), // Additional Z.ai API key for load balancing
|
||||
ZAI_API_KEY_3: z.string().min(1).optional(), // Additional Z.ai API key for load balancing
|
||||
XAI_API_KEY_1: z.string().min(1).optional(), // Primary xAI API key for load balancing
|
||||
XAI_API_KEY_2: z.string().min(1).optional(), // Additional xAI API key for load balancing
|
||||
XAI_API_KEY_3: z.string().min(1).optional(), // Additional xAI API key for load balancing
|
||||
OLLAMA_URL: z.string().url().optional(), // Ollama local LLM server URL
|
||||
VLLM_BASE_URL: z.string().url().optional(), // vLLM self-hosted base URL (OpenAI-compatible)
|
||||
VLLM_API_KEY: z.string().optional(), // Optional bearer token for vLLM
|
||||
|
||||
@@ -43,6 +43,9 @@ vi.mock('@/lib/core/config/env', () =>
|
||||
GEMINI_API_KEY_1: 'test-gemini-key-1',
|
||||
GEMINI_API_KEY_2: 'test-gemini-key-2',
|
||||
GEMINI_API_KEY_3: 'test-gemini-key-3',
|
||||
XAI_API_KEY_1: 'test-xai-key-1',
|
||||
XAI_API_KEY_2: 'test-xai-key-2',
|
||||
XAI_API_KEY_3: 'test-xai-key-3',
|
||||
})
|
||||
)
|
||||
|
||||
@@ -327,6 +330,11 @@ describe('getRotatingApiKey', () => {
|
||||
expect(result).toMatch(/^test-gemini-key-[1-3]$/)
|
||||
})
|
||||
|
||||
it.concurrent('should return xAI API key based on current minute', () => {
|
||||
const result = getRotatingApiKey('xai')
|
||||
expect(result).toMatch(/^test-xai-key-[1-3]$/)
|
||||
})
|
||||
|
||||
it.concurrent('should throw error for unsupported provider', () => {
|
||||
expect(() => getRotatingApiKey('unsupported')).toThrow('No rotation implemented for provider')
|
||||
})
|
||||
|
||||
@@ -218,3 +218,17 @@ describe('zai provider definition', () => {
|
||||
expect(getHostedModels()).toContain('glm-4.6')
|
||||
})
|
||||
})
|
||||
|
||||
describe('xai provider definition', () => {
|
||||
const xai = PROVIDER_DEFINITIONS.xai
|
||||
|
||||
it('is registered with grok-4.5 as the default model', () => {
|
||||
expect(xai).toBeDefined()
|
||||
expect(xai.id).toBe('xai')
|
||||
expect(xai.defaultModel).toBe('grok-4.5')
|
||||
})
|
||||
|
||||
it('is included in getHostedModels since Sim provides the xAI key server-side', () => {
|
||||
expect(getHostedModels()).toContain('grok-4.5')
|
||||
})
|
||||
})
|
||||
|
||||
@@ -3921,6 +3921,7 @@ export function getHostedModels(): string[] {
|
||||
...getProviderModels('anthropic'),
|
||||
...getProviderModels('google'),
|
||||
...getProviderModels('zai'),
|
||||
...getProviderModels('xai'),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@@ -865,7 +865,7 @@ describe('Cost Calculation', () => {
|
||||
})
|
||||
|
||||
describe('getHostedModels', () => {
|
||||
it.concurrent('should return OpenAI, Anthropic, and Google models as hosted', () => {
|
||||
it.concurrent('should return OpenAI, Anthropic, Google, and xAI models as hosted', () => {
|
||||
const hostedModels = getHostedModels()
|
||||
|
||||
expect(hostedModels).toContain('gpt-4o')
|
||||
@@ -877,8 +877,9 @@ describe('getHostedModels', () => {
|
||||
expect(hostedModels).toContain('gemini-2.5-pro')
|
||||
expect(hostedModels).toContain('gemini-2.5-flash')
|
||||
|
||||
expect(hostedModels).toContain('grok-4.5')
|
||||
|
||||
expect(hostedModels).not.toContain('deepseek-v3')
|
||||
expect(hostedModels).not.toContain('grok-4-latest')
|
||||
})
|
||||
|
||||
it.concurrent('should return an array of strings', () => {
|
||||
@@ -902,11 +903,12 @@ describe('shouldBillModelUsage', () => {
|
||||
|
||||
expect(shouldBillModelUsage('gemini-2.5-pro')).toBe(true)
|
||||
expect(shouldBillModelUsage('gemini-2.5-flash')).toBe(true)
|
||||
|
||||
expect(shouldBillModelUsage('grok-4.5')).toBe(true)
|
||||
})
|
||||
|
||||
it.concurrent('should return false for non-hosted models', () => {
|
||||
expect(shouldBillModelUsage('deepseek-v3')).toBe(false)
|
||||
expect(shouldBillModelUsage('grok-4-latest')).toBe(false)
|
||||
|
||||
expect(shouldBillModelUsage('unknown-model')).toBe(false)
|
||||
})
|
||||
|
||||
@@ -903,8 +903,9 @@ export function getApiKey(provider: string, model: string, userProvidedKey?: str
|
||||
const isClaudeModel = provider === 'anthropic'
|
||||
const isGeminiModel = provider === 'google'
|
||||
const isZaiModel = provider === 'zai'
|
||||
const isXaiModel = provider === 'xai'
|
||||
|
||||
if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel || isZaiModel)) {
|
||||
if (isHosted && (isOpenAIModel || isClaudeModel || isGeminiModel || isZaiModel || isXaiModel)) {
|
||||
const hostedModels = getHostedModels()
|
||||
const isModelHosted = hostedModels.some((m) => m.toLowerCase() === model.toLowerCase())
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ export type BYOKProviderId =
|
||||
| 'google'
|
||||
| 'mistral'
|
||||
| 'zai'
|
||||
| 'xai'
|
||||
| 'fireworks'
|
||||
| 'together'
|
||||
| 'baseten'
|
||||
|
||||
Reference in New Issue
Block a user