Skip to content

Commit a981925

Browse files
committed
feat(xai): add dynamic model discovery with correct pricing
- Implemented dynamic model discovery for xAI provider using /v1/language-models endpoint - Models are now fetched at runtime with real-time pricing and capabilities - Fixed pricing conversion: XAI API returns fractional cents (basis points), divide by 10,000 not 100 - This fixes pricing display showing $20.00 instead of $0.20 per 1M tokens - Removed static xAI models from MODELS_BY_PROVIDER to rely on dynamic discovery - Enhanced error logging with detailed status and URL information - Support dynamic model context window overrides from API - Fixed parseApiPrice to handle zero values correctly (for free models) - Provided complete ModelInfo fallback in useSelectedModel for UI type safety - Added comprehensive test coverage for cost utilities and XAI fetcher - Updated all tests to reflect correct pricing scale and dynamic model behavior
1 parent 8e4b145 commit a981925

File tree

19 files changed

+619
-91
lines changed

19 files changed

+619
-91
lines changed

packages/types/src/provider-settings.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ import {
2121
sambaNovaModels,
2222
vertexModels,
2323
vscodeLlmModels,
24-
xaiModels,
2524
internationalZAiModels,
2625
minimaxModels,
2726
} from "./providers/index.js"
@@ -50,6 +49,7 @@ export const dynamicProviders = [
5049
"glama",
5150
"roo",
5251
"chutes",
52+
"xai",
5353
] as const
5454

5555
export type DynamicProvider = (typeof dynamicProviders)[number]
@@ -137,7 +137,6 @@ export const providerNames = [
137137
"roo",
138138
"sambanova",
139139
"vertex",
140-
"xai",
141140
"zai",
142141
] as const
143142

@@ -354,6 +353,7 @@ const fakeAiSchema = baseProviderSettingsSchema.extend({
354353

355354
const xaiSchema = apiModelIdProviderModelSchema.extend({
356355
xaiApiKey: z.string().optional(),
356+
xaiModelContextWindow: z.number().optional(),
357357
})
358358

359359
const groqSchema = apiModelIdProviderModelSchema.extend({
@@ -709,7 +709,7 @@ export const MODELS_BY_PROVIDER: Record<
709709
label: "VS Code LM API",
710710
models: Object.keys(vscodeLlmModels),
711711
},
712-
xai: { id: "xai", label: "xAI (Grok)", models: Object.keys(xaiModels) },
712+
xai: { id: "xai", label: "xAI (Grok)", models: [] },
713713
zai: { id: "zai", label: "Zai", models: Object.keys(internationalZAiModels) },
714714

715715
// Dynamic providers; models pulled from remote APIs.

packages/types/src/providers/xai.ts

Lines changed: 36 additions & 66 deletions
Original file line numberDiff line numberDiff line change
@@ -3,93 +3,63 @@ import type { ModelInfo } from "../model.js"
33
// https://docs.x.ai/docs/api-reference
44
export type XAIModelId = keyof typeof xaiModels
55

6-
export const xaiDefaultModelId: XAIModelId = "grok-code-fast-1"
6+
export const xaiDefaultModelId: XAIModelId = "grok-4-fast-reasoning"
7+
8+
/**
9+
* Partial ModelInfo for xAI static registry.
10+
* Contains only fields not available from the xAI API:
11+
* - contextWindow: Not provided by API
12+
* - maxTokens: Not provided by API
13+
* - description: User-friendly descriptions
14+
* - supportsReasoningEffort: Special capability flag
15+
*
16+
* All other fields (pricing, supportsPromptCache, supportsImages) are fetched dynamically.
17+
*/
18+
type XAIStaticModelInfo = Pick<ModelInfo, "contextWindow" | "description"> & {
19+
maxTokens?: number | null
20+
supportsReasoningEffort?: boolean
21+
}
722

823
export const xaiModels = {
924
"grok-code-fast-1": {
1025
maxTokens: 16_384,
11-
contextWindow: 262_144,
12-
supportsImages: false,
13-
supportsPromptCache: true,
14-
inputPrice: 0.2,
15-
outputPrice: 1.5,
16-
cacheWritesPrice: 0.02,
17-
cacheReadsPrice: 0.02,
26+
contextWindow: 256_000,
1827
description: "xAI's Grok Code Fast model with 256K context window",
1928
},
20-
"grok-4": {
21-
maxTokens: 8192,
22-
contextWindow: 256000,
23-
supportsImages: true,
24-
supportsPromptCache: true,
25-
inputPrice: 3.0,
26-
outputPrice: 15.0,
27-
cacheWritesPrice: 0.75,
28-
cacheReadsPrice: 0.75,
29+
"grok-4-0709": {
30+
maxTokens: 16_384,
31+
contextWindow: 256_000,
2932
description: "xAI's Grok-4 model with 256K context window",
3033
},
34+
"grok-4-fast-non-reasoning": {
35+
maxTokens: 32_768,
36+
contextWindow: 2_000_000,
37+
description: "xAI's Grok-4 Fast Non-Reasoning model with 2M context window",
38+
},
39+
"grok-4-fast-reasoning": {
40+
maxTokens: 32_768,
41+
contextWindow: 2_000_000,
42+
description: "xAI's Grok-4 Fast Reasoning model with 2M context window",
43+
},
3144
"grok-3": {
3245
maxTokens: 8192,
33-
contextWindow: 131072,
34-
supportsImages: false,
35-
supportsPromptCache: true,
36-
inputPrice: 3.0,
37-
outputPrice: 15.0,
38-
cacheWritesPrice: 0.75,
39-
cacheReadsPrice: 0.75,
46+
contextWindow: 131_072,
4047
description: "xAI's Grok-3 model with 128K context window",
4148
},
42-
"grok-3-fast": {
43-
maxTokens: 8192,
44-
contextWindow: 131072,
45-
supportsImages: false,
46-
supportsPromptCache: true,
47-
inputPrice: 5.0,
48-
outputPrice: 25.0,
49-
cacheWritesPrice: 1.25,
50-
cacheReadsPrice: 1.25,
51-
description: "xAI's Grok-3 fast model with 128K context window",
52-
},
5349
"grok-3-mini": {
5450
maxTokens: 8192,
55-
contextWindow: 131072,
56-
supportsImages: false,
57-
supportsPromptCache: true,
58-
inputPrice: 0.3,
59-
outputPrice: 0.5,
60-
cacheWritesPrice: 0.07,
61-
cacheReadsPrice: 0.07,
51+
contextWindow: 131_072,
6252
description: "xAI's Grok-3 mini model with 128K context window",
6353
supportsReasoningEffort: true,
6454
},
65-
"grok-3-mini-fast": {
66-
maxTokens: 8192,
67-
contextWindow: 131072,
68-
supportsImages: false,
69-
supportsPromptCache: true,
70-
inputPrice: 0.6,
71-
outputPrice: 4.0,
72-
cacheWritesPrice: 0.15,
73-
cacheReadsPrice: 0.15,
74-
description: "xAI's Grok-3 mini fast model with 128K context window",
75-
supportsReasoningEffort: true,
76-
},
7755
"grok-2-1212": {
7856
maxTokens: 8192,
79-
contextWindow: 131072,
80-
supportsImages: false,
81-
supportsPromptCache: false,
82-
inputPrice: 2.0,
83-
outputPrice: 10.0,
84-
description: "xAI's Grok-2 model (version 1212) with 128K context window",
57+
contextWindow: 32_768,
58+
description: "xAI's Grok-2 model (version 1212) with 32K context window",
8559
},
8660
"grok-2-vision-1212": {
8761
maxTokens: 8192,
88-
contextWindow: 32768,
89-
supportsImages: true,
90-
supportsPromptCache: false,
91-
inputPrice: 2.0,
92-
outputPrice: 10.0,
62+
contextWindow: 32_768,
9363
description: "xAI's Grok-2 Vision model (version 1212) with image support and 32K context window",
9464
},
95-
} as const satisfies Record<string, ModelInfo>
65+
} as const satisfies Record<string, XAIStaticModelInfo>

src/api/providers/__tests__/xai.spec.ts

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,12 @@ describe("XAIHandler", () => {
5757
it("should return default model when no model is specified", () => {
5858
const model = handler.getModel()
5959
expect(model.id).toBe(xaiDefaultModelId)
60-
expect(model.info).toEqual(xaiModels[xaiDefaultModelId])
60+
expect(model.info).toMatchObject({
61+
contextWindow: xaiModels[xaiDefaultModelId].contextWindow,
62+
maxTokens: xaiModels[xaiDefaultModelId].maxTokens,
63+
description: xaiModels[xaiDefaultModelId].description,
64+
})
65+
expect(model.info.supportsPromptCache).toBe(false) // Placeholder until dynamic data loads
6166
})
6267

6368
test("should return specified model when valid model is provided", () => {
@@ -66,7 +71,12 @@ describe("XAIHandler", () => {
6671
const model = handlerWithModel.getModel()
6772

6873
expect(model.id).toBe(testModelId)
69-
expect(model.info).toEqual(xaiModels[testModelId])
74+
expect(model.info).toMatchObject({
75+
contextWindow: xaiModels[testModelId].contextWindow,
76+
maxTokens: xaiModels[testModelId].maxTokens,
77+
description: xaiModels[testModelId].description,
78+
})
79+
expect(model.info.supportsPromptCache).toBe(false) // Placeholder until dynamic data loads
7080
})
7181

7282
it("should include reasoning_effort parameter for mini models", async () => {
@@ -234,12 +244,13 @@ describe("XAIHandler", () => {
234244

235245
// Verify the usage data
236246
expect(firstChunk.done).toBe(false)
237-
expect(firstChunk.value).toEqual({
247+
expect(firstChunk.value).toMatchObject({
238248
type: "usage",
239249
inputTokens: 10,
240250
outputTokens: 20,
241251
cacheReadTokens: 5,
242252
cacheWriteTokens: 15,
253+
totalCost: expect.any(Number),
243254
})
244255
})
245256

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import { describe, it, expect, vi, beforeEach } from "vitest"
2+
import axios from "axios"
3+
4+
vi.mock("axios")
5+
6+
import { getXaiModels } from "../xai"
7+
import { xaiModels } from "@roo-code/types"
8+
9+
describe("getXaiModels", () => {
10+
const mockedAxios = axios as unknown as { get: ReturnType<typeof vi.fn> }
11+
12+
beforeEach(() => {
13+
vi.clearAllMocks()
14+
})
15+
16+
it("returns mapped models with pricing and modalities (augmenting static info when available)", async () => {
17+
mockedAxios.get = vi.fn().mockResolvedValue({
18+
data: {
19+
models: [
20+
{
21+
id: "grok-3",
22+
input_modalities: ["text"],
23+
output_modalities: ["text"],
24+
prompt_text_token_price: 2000, // 2000 fractional cents = $0.20 per 1M tokens
25+
cached_prompt_text_token_price: 500, // 500 fractional cents = $0.05 per 1M tokens
26+
completion_text_token_price: 10000, // 10000 fractional cents = $1.00 per 1M tokens
27+
aliases: ["grok-3-latest"],
28+
},
29+
],
30+
},
31+
})
32+
33+
const result = await getXaiModels("key", "https://api.x.ai/v1")
34+
expect(result["grok-3"]).toBeDefined()
35+
expect(result["grok-3"]?.supportsImages).toBe(false)
36+
expect(result["grok-3"]?.inputPrice).toBeCloseTo(0.2) // $0.20 per 1M tokens
37+
expect(result["grok-3"]?.outputPrice).toBeCloseTo(1.0) // $1.00 per 1M tokens
38+
expect(result["grok-3"]?.cacheReadsPrice).toBeCloseTo(0.05) // $0.05 per 1M tokens
39+
// aliases are not added to avoid UI duplication
40+
expect(result["grok-3-latest"]).toBeUndefined()
41+
})
42+
43+
it("returns empty object on schema mismatches (graceful degradation)", async () => {
44+
mockedAxios.get = vi.fn().mockResolvedValue({
45+
data: { data: [{ bogus: true }] },
46+
})
47+
const result = await getXaiModels("key")
48+
expect(result).toEqual({})
49+
})
50+
51+
it("includes Authorization header when apiKey provided", async () => {
52+
mockedAxios.get = vi.fn().mockResolvedValue({ data: { data: [] } })
53+
await getXaiModels("secret")
54+
expect((axios.get as any).mock.calls[0][1].headers.Authorization).toBe("Bearer secret")
55+
})
56+
})

src/api/providers/fetchers/modelCache.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import { getDeepInfraModels } from "./deepinfra"
2626
import { getHuggingFaceModels } from "./huggingface"
2727
import { getRooModels } from "./roo"
2828
import { getChutesModels } from "./chutes"
29+
import { getXaiModels } from "./xai"
2930

3031
const memoryCache = new NodeCache({ stdTTL: 5 * 60, checkperiod: 5 * 60 })
3132

@@ -101,6 +102,9 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
101102
case "huggingface":
102103
models = await getHuggingFaceModels()
103104
break
105+
case "xai":
106+
models = await getXaiModels(options.apiKey, options.baseUrl)
107+
break
104108
case "roo": {
105109
// Roo Code Cloud provider requires baseUrl and optional apiKey
106110
const rooBaseUrl =
@@ -121,7 +125,7 @@ export const getModels = async (options: GetModelsOptions): Promise<ModelRecord>
121125
// Cache the fetched models (even if empty, to signify a successful fetch with no models).
122126
memoryCache.set(provider, models)
123127

124-
await writeModels(provider, models).catch((err) =>
128+
await writeModels(provider, models || {}).catch((err) =>
125129
console.error(`[getModels] Error writing ${provider} models to file cache:`, err),
126130
)
127131

0 commit comments

Comments
 (0)