|
| 1 | +import axios from "axios" |
| 2 | +import { z } from "zod" |
| 3 | + |
| 4 | +import { type ModelInfo, xaiModels } from "@roo-code/types" |
| 5 | +import { DEFAULT_HEADERS } from "../../providers/constants" |
| 6 | + |
| 7 | +/** |
| 8 | + * Schema for GET https://api.x.ai/v1/language-models |
| 9 | + * This endpoint returns rich metadata including modalities and pricing. |
| 10 | + */ |
| 11 | +const xaiLanguageModelSchema = z.object({ |
| 12 | + id: z.string(), |
| 13 | + input_modalities: z.array(z.string()).optional(), |
| 14 | + output_modalities: z.array(z.string()).optional(), |
| 15 | + prompt_text_token_price: z.number().optional(), // cents per 1M tokens |
| 16 | + cached_prompt_text_token_price: z.number().optional(), // cents per 1M tokens |
| 17 | + prompt_image_token_price: z.number().optional(), // cents per 1M tokens |
| 18 | + completion_text_token_price: z.number().optional(), // cents per 1M tokens |
| 19 | + search_price: z.number().optional(), |
| 20 | + aliases: z.array(z.string()).optional(), |
| 21 | +}) |
| 22 | + |
| 23 | +const xaiLanguageModelsResponseSchema = z.object({ |
| 24 | + models: z.array(xaiLanguageModelSchema), |
| 25 | +}) |
| 26 | + |
| 27 | +/** |
| 28 | + * Fetch available xAI models for the authenticated account. |
| 29 | + * - Uses Bearer Authorization header when apiKey is provided |
| 30 | + * - Maps discovered IDs to ModelInfo using static catalog (xaiModels) when possible |
| 31 | + * - For models not in static catalog, contextWindow and maxTokens remain undefined |
| 32 | + */ |
| 33 | +export async function getXaiModels(apiKey?: string, baseUrl?: string): Promise<Record<string, ModelInfo>> { |
| 34 | + const models: Record<string, ModelInfo> = {} |
| 35 | + // Build proper endpoint whether user passes https://api.x.ai or https://api.x.ai/v1 |
| 36 | + const base = baseUrl ? baseUrl.replace(/\/+$/, "") : "https://api.x.ai" |
| 37 | + const url = base.endsWith("/v1") ? `${base}/language-models` : `${base}/v1/language-models` |
| 38 | + |
| 39 | + try { |
| 40 | + const resp = await axios.get(url, { |
| 41 | + headers: { |
| 42 | + ...DEFAULT_HEADERS, |
| 43 | + Accept: "application/json", |
| 44 | + ...(apiKey ? { Authorization: `Bearer ${apiKey}` } : {}), |
| 45 | + }, |
| 46 | + }) |
| 47 | + |
| 48 | + const parsed = xaiLanguageModelsResponseSchema.safeParse(resp.data) |
| 49 | + const items = parsed.success |
| 50 | + ? parsed.data.models |
| 51 | + : Array.isArray((resp.data as any)?.models) |
| 52 | + ? (resp.data as any)?.models |
| 53 | + : [] |
| 54 | + |
| 55 | + if (!parsed.success) { |
| 56 | + console.error("xAI language models response validation failed", parsed.error?.format?.() ?? parsed.error) |
| 57 | + } |
| 58 | + |
| 59 | + // Helper to convert cents-per-1M to dollars-per-1M (assumption per API examples) |
| 60 | + const centsToDollars = (v?: number) => (typeof v === "number" ? v / 100 : undefined) |
| 61 | + |
| 62 | + for (const m of items) { |
| 63 | + const id = m.id |
| 64 | + const staticInfo = xaiModels[id as keyof typeof xaiModels] |
| 65 | + const supportsImages = Array.isArray(m.input_modalities) ? m.input_modalities.includes("image") : false |
| 66 | + |
| 67 | + // Cache support is indicated by presence of cached_prompt_text_token_price field (even if 0) |
| 68 | + const supportsPromptCache = typeof m.cached_prompt_text_token_price === "number" |
| 69 | + const cacheReadsPrice = supportsPromptCache ? centsToDollars(m.cached_prompt_text_token_price) : undefined |
| 70 | + |
| 71 | + const info: ModelInfo = { |
| 72 | + maxTokens: staticInfo?.maxTokens ?? undefined, |
| 73 | + contextWindow: staticInfo?.contextWindow ?? undefined, |
| 74 | + supportsImages, |
| 75 | + supportsPromptCache, |
| 76 | + inputPrice: centsToDollars(m.prompt_text_token_price), |
| 77 | + outputPrice: centsToDollars(m.completion_text_token_price), |
| 78 | + cacheReadsPrice, |
| 79 | + cacheWritesPrice: cacheReadsPrice, // xAI uses same price for reads and writes |
| 80 | + description: staticInfo?.description, |
| 81 | + supportsReasoningEffort: |
| 82 | + staticInfo && "supportsReasoningEffort" in staticInfo |
| 83 | + ? staticInfo.supportsReasoningEffort |
| 84 | + : undefined, |
| 85 | + // leave other optional fields undefined unless available via static definitions |
| 86 | + } |
| 87 | + |
| 88 | + models[id] = info |
| 89 | + // Aliases are not added to the model list to avoid duplication in UI |
| 90 | + // Users should use the primary model ID; xAI API will handle alias resolution |
| 91 | + } |
| 92 | + } catch (error) { |
| 93 | + try { |
| 94 | + const err = JSON.stringify(error, Object.getOwnPropertyNames(error), 2) |
| 95 | + console.error(`[xAI] models fetch failed: ${err}`) |
| 96 | + } catch { |
| 97 | + console.error("[xAI] models fetch failed.") |
| 98 | + } |
| 99 | + throw error |
| 100 | + } |
| 101 | + |
| 102 | + return models |
| 103 | +} |
0 commit comments