mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-08-28 05:27:24 +00:00
feat: migrate xAI provider to use dedicated @ai-sdk/xai package (#11158)
Co-authored-by: Roo Code <roomote@roocode.com> Co-authored-by: daniel-lxs <ricciodaniel98@gmail.com>
This commit is contained in:
roomote[bot]
GitHub
parent
67fb150727
commit
e90e6178e3
5 changed files with 829 additions and 634 deletions
|
|
@ -30,6 +30,8 @@ export const xaiModels = {
|
|||
cacheReadsPrice: 0.05,
|
||||
description:
|
||||
"xAI's Grok 4.1 Fast model with 2M context window, optimized for high-performance agentic tool calling with reasoning",
|
||||
supportsReasoningEffort: ["low", "high"],
|
||||
reasoningEffort: "low",
|
||||
includedTools: ["search_replace"],
|
||||
excludedTools: ["apply_diff"],
|
||||
},
|
||||
|
|
@ -58,6 +60,8 @@ export const xaiModels = {
|
|||
cacheReadsPrice: 0.05,
|
||||
description:
|
||||
"xAI's Grok 4 Fast model with 2M context window, optimized for high-performance agentic tool calling with reasoning",
|
||||
supportsReasoningEffort: ["low", "high"],
|
||||
reasoningEffort: "low",
|
||||
includedTools: ["search_replace"],
|
||||
excludedTools: ["apply_diff"],
|
||||
},
|
||||
|
|
|
|||
36
pnpm-lock.yaml
generated
36
pnpm-lock.yaml
generated
|
|
@ -761,6 +761,9 @@ importers:
|
|||
'@ai-sdk/mistral':
|
||||
specifier: ^3.0.0
|
||||
version: 3.0.18(zod@3.25.76)
|
||||
'@ai-sdk/xai':
|
||||
specifier: ^3.0.46
|
||||
version: 3.0.46(zod@3.25.76)
|
||||
'@anthropic-ai/bedrock-sdk':
|
||||
specifier: ^0.10.2
|
||||
version: 0.10.4
|
||||
|
|
@ -1462,6 +1465,12 @@ packages:
|
|||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/openai-compatible@2.0.26':
|
||||
resolution: {integrity: sha512-l6jdFjI1C2eDAEm7oo+dnRn0oG1EkcyqfbEZ7ozT0TnYrah6amX2JkftYMP1GRzNtAeCB3WNN8XspXdmi6ZNlQ==}
|
||||
engines: {node: '>=18'}
|
||||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@3.0.20':
|
||||
resolution: {integrity: sha512-iXHVe0apM2zUEzauqJwqmpC37A5rihrStAih5Ks+JE32iTe4LZ58y17UGBjpQQTCRw9YxMeo2UFLxLpBluyvLQ==}
|
||||
engines: {node: '>=18'}
|
||||
|
|
@ -1512,6 +1521,12 @@ packages:
|
|||
resolution: {integrity: sha512-VkPLrutM6VdA924/mG8OS+5frbVTcu6e046D2bgDo00tehBANR1QBJ/mPcZ9tXMFOsVcm6SQArOregxePzTFPw==}
|
||||
engines: {node: '>=18'}
|
||||
|
||||
'@ai-sdk/xai@3.0.46':
|
||||
resolution: {integrity: sha512-26qM/jYcFhF5krTM7bQT1CiZcdz22EQmA+r5me1hKYFM/yM20sSUMHnAcUzvzuuG9oQVKF0tziU2IcC0HX5huQ==}
|
||||
engines: {node: '>=18'}
|
||||
peerDependencies:
|
||||
zod: 3.25.76
|
||||
|
||||
'@alcalzone/ansi-tokenize@0.2.3':
|
||||
resolution: {integrity: sha512-jsElTJ0sQ4wHRz+C45tfect76BwbTbgkgKByOzpCN9xG61N5V6u/glvg1CsNJhq2xJIFpKHSwG3D2wPPuEYOrQ==}
|
||||
engines: {node: '>=18'}
|
||||
|
|
@ -6530,10 +6545,6 @@ packages:
|
|||
resolution: {integrity: sha512-mQw+2fkQbALzQ7V0MY0IqdnXNOeTtP4r0lN9z7AAawCXgqea7bDii20AYrIBrFd/Hx0M2Ocz6S111CaFkUcb0Q==}
|
||||
engines: {node: '>=0.8.x'}
|
||||
|
||||
eventsource-parser@3.0.2:
|
||||
resolution: {integrity: sha512-6RxOBZ/cYgd8usLwsEl+EC09Au/9BcmCKYF2/xbml6DNczf7nv0MQb+7BA2F+li6//I+28VNlQR37XfQtcAJuA==}
|
||||
engines: {node: '>=18.0.0'}
|
||||
|
||||
eventsource-parser@3.0.6:
|
||||
resolution: {integrity: sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==}
|
||||
engines: {node: '>=18.0.0'}
|
||||
|
|
@ -11146,6 +11157,12 @@ snapshots:
|
|||
'@ai-sdk/provider-utils': 4.0.11(zod@3.25.76)
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/openai-compatible@2.0.26(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 3.0.7
|
||||
'@ai-sdk/provider-utils': 4.0.13(zod@3.25.76)
|
||||
zod: 3.25.76
|
||||
|
||||
'@ai-sdk/provider-utils@3.0.20(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/provider': 2.0.1
|
||||
|
|
@ -11202,6 +11219,13 @@ snapshots:
|
|||
dependencies:
|
||||
json-schema: 0.4.0
|
||||
|
||||
'@ai-sdk/xai@3.0.46(zod@3.25.76)':
|
||||
dependencies:
|
||||
'@ai-sdk/openai-compatible': 2.0.26(zod@3.25.76)
|
||||
'@ai-sdk/provider': 3.0.7
|
||||
'@ai-sdk/provider-utils': 4.0.13(zod@3.25.76)
|
||||
zod: 3.25.76
|
||||
|
||||
'@alcalzone/ansi-tokenize@0.2.3':
|
||||
dependencies:
|
||||
ansi-styles: 6.2.3
|
||||
|
|
@ -17027,13 +17051,11 @@ snapshots:
|
|||
|
||||
events@3.3.0: {}
|
||||
|
||||
eventsource-parser@3.0.2: {}
|
||||
|
||||
eventsource-parser@3.0.6: {}
|
||||
|
||||
eventsource@3.0.7:
|
||||
dependencies:
|
||||
eventsource-parser: 3.0.2
|
||||
eventsource-parser: 3.0.6
|
||||
|
||||
exceljs@4.4.0:
|
||||
dependencies:
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -1,166 +1,190 @@
|
|||
import { Anthropic } from "@anthropic-ai/sdk"
|
||||
import OpenAI from "openai"
|
||||
import { createXai } from "@ai-sdk/xai"
|
||||
import { streamText, generateText, ToolSet } from "ai"
|
||||
|
||||
import { type XAIModelId, xaiDefaultModelId, xaiModels, ApiProviderError } from "@roo-code/types"
|
||||
import { TelemetryService } from "@roo-code/telemetry"
|
||||
import { type XAIModelId, xaiDefaultModelId, xaiModels, type ModelInfo } from "@roo-code/types"
|
||||
|
||||
import { NativeToolCallParser } from "../../core/assistant-message/NativeToolCallParser"
|
||||
import type { ApiHandlerOptions } from "../../shared/api"
|
||||
|
||||
import { ApiStream } from "../transform/stream"
|
||||
import { convertToOpenAiMessages } from "../transform/openai-format"
|
||||
import {
|
||||
convertToAiSdkMessages,
|
||||
convertToolsForAiSdk,
|
||||
processAiSdkStreamPart,
|
||||
mapToolChoice,
|
||||
handleAiSdkError,
|
||||
} from "../transform/ai-sdk"
|
||||
import { ApiStream, ApiStreamUsageChunk } from "../transform/stream"
|
||||
import { getModelParams } from "../transform/model-params"
|
||||
|
||||
import { DEFAULT_HEADERS } from "./constants"
|
||||
import { BaseProvider } from "./base-provider"
|
||||
import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
|
||||
import { handleOpenAIError } from "./utils/openai-error-handler"
|
||||
|
||||
const XAI_DEFAULT_TEMPERATURE = 0
|
||||
|
||||
/**
|
||||
* xAI provider using the dedicated @ai-sdk/xai package.
|
||||
* Provides native support for Grok models including reasoning models.
|
||||
*/
|
||||
export class XAIHandler extends BaseProvider implements SingleCompletionHandler {
|
||||
protected options: ApiHandlerOptions
|
||||
private client: OpenAI
|
||||
private readonly providerName = "xAI"
|
||||
protected provider: ReturnType<typeof createXai>
|
||||
|
||||
constructor(options: ApiHandlerOptions) {
|
||||
super()
|
||||
this.options = options
|
||||
|
||||
const apiKey = this.options.xaiApiKey ?? "not-provided"
|
||||
|
||||
this.client = new OpenAI({
|
||||
// Create the xAI provider using AI SDK
|
||||
this.provider = createXai({
|
||||
baseURL: "https://api.x.ai/v1",
|
||||
apiKey: apiKey,
|
||||
defaultHeaders: DEFAULT_HEADERS,
|
||||
apiKey: options.xaiApiKey ?? "not-provided",
|
||||
headers: DEFAULT_HEADERS,
|
||||
})
|
||||
}
|
||||
|
||||
override getModel() {
|
||||
override getModel(): {
|
||||
id: XAIModelId
|
||||
info: ModelInfo
|
||||
maxTokens?: number
|
||||
temperature?: number
|
||||
reasoning?: any
|
||||
} {
|
||||
const id =
|
||||
this.options.apiModelId && this.options.apiModelId in xaiModels
|
||||
? (this.options.apiModelId as XAIModelId)
|
||||
: xaiDefaultModelId
|
||||
|
||||
const info = xaiModels[id]
|
||||
const params = getModelParams({ format: "openai", modelId: id, model: info, settings: this.options })
|
||||
const params = getModelParams({
|
||||
format: "openai",
|
||||
modelId: id,
|
||||
model: info,
|
||||
settings: this.options,
|
||||
defaultTemperature: XAI_DEFAULT_TEMPERATURE,
|
||||
})
|
||||
return { id, info, ...params }
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the language model for the configured model ID.
|
||||
*/
|
||||
protected getLanguageModel() {
|
||||
const { id } = this.getModel()
|
||||
return this.provider(id)
|
||||
}
|
||||
|
||||
/**
|
||||
* Process usage metrics from the AI SDK response.
|
||||
*/
|
||||
protected processUsageMetrics(
|
||||
usage: {
|
||||
inputTokens?: number
|
||||
outputTokens?: number
|
||||
details?: {
|
||||
cachedInputTokens?: number
|
||||
reasoningTokens?: number
|
||||
}
|
||||
},
|
||||
providerMetadata?: {
|
||||
xai?: {
|
||||
cachedPromptTokens?: number
|
||||
}
|
||||
},
|
||||
): ApiStreamUsageChunk {
|
||||
// Extract cache metrics from xAI's providerMetadata if available
|
||||
// xAI supports prompt caching through prompt_tokens_details.cached_tokens
|
||||
const cacheReadTokens = providerMetadata?.xai?.cachedPromptTokens ?? usage.details?.cachedInputTokens
|
||||
|
||||
return {
|
||||
type: "usage",
|
||||
inputTokens: usage.inputTokens || 0,
|
||||
outputTokens: usage.outputTokens || 0,
|
||||
cacheReadTokens,
|
||||
cacheWriteTokens: undefined, // xAI doesn't report cache write tokens separately
|
||||
reasoningTokens: usage.details?.reasoningTokens,
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the max tokens parameter to include in the request.
|
||||
*/
|
||||
protected getMaxOutputTokens(): number | undefined {
|
||||
const { info } = this.getModel()
|
||||
return this.options.modelMaxTokens || info.maxTokens || undefined
|
||||
}
|
||||
|
||||
/**
|
||||
* Create a message stream using the AI SDK.
|
||||
*/
|
||||
override async *createMessage(
|
||||
systemPrompt: string,
|
||||
messages: Anthropic.Messages.MessageParam[],
|
||||
metadata?: ApiHandlerCreateMessageMetadata,
|
||||
): ApiStream {
|
||||
const { id: modelId, info: modelInfo, reasoning } = this.getModel()
|
||||
const { temperature, reasoning } = this.getModel()
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
// Use the OpenAI-compatible API.
|
||||
const requestOptions = {
|
||||
model: modelId,
|
||||
max_tokens: modelInfo.maxTokens,
|
||||
temperature: this.options.modelTemperature ?? XAI_DEFAULT_TEMPERATURE,
|
||||
messages: [
|
||||
{ role: "system", content: systemPrompt },
|
||||
...convertToOpenAiMessages(messages),
|
||||
] as OpenAI.Chat.ChatCompletionMessageParam[],
|
||||
stream: true as const,
|
||||
stream_options: { include_usage: true },
|
||||
...(reasoning && reasoning),
|
||||
tools: this.convertToolsForOpenAI(metadata?.tools),
|
||||
tool_choice: metadata?.tool_choice,
|
||||
parallel_tool_calls: metadata?.parallelToolCalls ?? true,
|
||||
// Convert messages to AI SDK format
|
||||
const aiSdkMessages = convertToAiSdkMessages(messages)
|
||||
|
||||
// Convert tools to OpenAI format first, then to AI SDK format
|
||||
const openAiTools = this.convertToolsForOpenAI(metadata?.tools)
|
||||
const aiSdkTools = convertToolsForAiSdk(openAiTools) as ToolSet | undefined
|
||||
|
||||
// Build the request options
|
||||
const requestOptions: Parameters<typeof streamText>[0] = {
|
||||
model: languageModel,
|
||||
system: systemPrompt,
|
||||
messages: aiSdkMessages,
|
||||
temperature: this.options.modelTemperature ?? temperature ?? XAI_DEFAULT_TEMPERATURE,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
tools: aiSdkTools,
|
||||
toolChoice: mapToolChoice(metadata?.tool_choice),
|
||||
...(reasoning && { providerOptions: { xai: reasoning } }),
|
||||
}
|
||||
|
||||
let stream
|
||||
// Use streamText for streaming responses
|
||||
const result = streamText(requestOptions)
|
||||
|
||||
try {
|
||||
stream = await this.client.chat.completions.create(requestOptions)
|
||||
// Process the full stream to get all events including reasoning
|
||||
for await (const part of result.fullStream) {
|
||||
for (const chunk of processAiSdkStreamPart(part)) {
|
||||
yield chunk
|
||||
}
|
||||
}
|
||||
|
||||
// Yield usage metrics at the end, including cache metrics from providerMetadata
|
||||
const usage = await result.usage
|
||||
const providerMetadata = await result.providerMetadata
|
||||
if (usage) {
|
||||
yield this.processUsageMetrics(usage, providerMetadata as any)
|
||||
}
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error)
|
||||
const apiError = new ApiProviderError(errorMessage, this.providerName, modelId, "createMessage")
|
||||
TelemetryService.instance.captureException(apiError)
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
}
|
||||
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk.choices[0]?.delta
|
||||
const finishReason = chunk.choices[0]?.finish_reason
|
||||
|
||||
if (delta?.content) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: delta.content,
|
||||
}
|
||||
}
|
||||
|
||||
if (delta && "reasoning_content" in delta && delta.reasoning_content) {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: delta.reasoning_content as string,
|
||||
}
|
||||
}
|
||||
|
||||
// Handle tool calls in stream - emit partial chunks for NativeToolCallParser
|
||||
if (delta?.tool_calls) {
|
||||
for (const toolCall of delta.tool_calls) {
|
||||
yield {
|
||||
type: "tool_call_partial",
|
||||
index: toolCall.index,
|
||||
id: toolCall.id,
|
||||
name: toolCall.function?.name,
|
||||
arguments: toolCall.function?.arguments,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process finish_reason to emit tool_call_end events
|
||||
// This ensures tool calls are finalized even if the stream doesn't properly close
|
||||
if (finishReason) {
|
||||
const endEvents = NativeToolCallParser.processFinishReason(finishReason)
|
||||
for (const event of endEvents) {
|
||||
yield event
|
||||
}
|
||||
}
|
||||
|
||||
if (chunk.usage) {
|
||||
// Extract detailed token information if available
|
||||
// First check for prompt_tokens_details structure (real API response)
|
||||
const promptDetails = "prompt_tokens_details" in chunk.usage ? chunk.usage.prompt_tokens_details : null
|
||||
const cachedTokens = promptDetails && "cached_tokens" in promptDetails ? promptDetails.cached_tokens : 0
|
||||
|
||||
// Fall back to direct fields in usage (used in test mocks)
|
||||
const readTokens =
|
||||
cachedTokens ||
|
||||
("cache_read_input_tokens" in chunk.usage ? (chunk.usage as any).cache_read_input_tokens : 0)
|
||||
const writeTokens =
|
||||
"cache_creation_input_tokens" in chunk.usage ? (chunk.usage as any).cache_creation_input_tokens : 0
|
||||
|
||||
yield {
|
||||
type: "usage",
|
||||
inputTokens: chunk.usage.prompt_tokens || 0,
|
||||
outputTokens: chunk.usage.completion_tokens || 0,
|
||||
cacheReadTokens: readTokens,
|
||||
cacheWriteTokens: writeTokens,
|
||||
}
|
||||
}
|
||||
// Handle AI SDK errors (AI_RetryError, AI_APICallError, etc.)
|
||||
throw handleAiSdkError(error, "xAI")
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Complete a prompt using the AI SDK generateText.
|
||||
*/
|
||||
async completePrompt(prompt: string): Promise<string> {
|
||||
const { id: modelId, reasoning } = this.getModel()
|
||||
const { temperature, reasoning } = this.getModel()
|
||||
const languageModel = this.getLanguageModel()
|
||||
|
||||
try {
|
||||
const response = await this.client.chat.completions.create({
|
||||
model: modelId,
|
||||
messages: [{ role: "user", content: prompt }],
|
||||
...(reasoning && reasoning),
|
||||
const { text } = await generateText({
|
||||
model: languageModel,
|
||||
prompt,
|
||||
maxOutputTokens: this.getMaxOutputTokens(),
|
||||
temperature: this.options.modelTemperature ?? temperature ?? XAI_DEFAULT_TEMPERATURE,
|
||||
...(reasoning && { providerOptions: { xai: reasoning } }),
|
||||
})
|
||||
|
||||
return response.choices[0]?.message.content || ""
|
||||
return text
|
||||
} catch (error) {
|
||||
const errorMessage = error instanceof Error ? error.message : String(error)
|
||||
const apiError = new ApiProviderError(errorMessage, this.providerName, modelId, "completePrompt")
|
||||
TelemetryService.instance.captureException(apiError)
|
||||
throw handleOpenAIError(error, this.providerName)
|
||||
throw handleAiSdkError(error, "xAI")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -455,6 +455,7 @@
|
|||
"@ai-sdk/fireworks": "^2.0.26",
|
||||
"@ai-sdk/groq": "^3.0.19",
|
||||
"@ai-sdk/mistral": "^3.0.0",
|
||||
"@ai-sdk/xai": "^3.0.46",
|
||||
"sambanova-ai-provider": "^1.2.2",
|
||||
"@anthropic-ai/bedrock-sdk": "^0.10.2",
|
||||
"@anthropic-ai/sdk": "^0.37.0",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue