feat: add X AI provider integration

This commit is contained in:
Andrew Monostate 2025-02-27 11:52:40 -03:00 committed by ocasta181
parent d6184e9dac
commit d7e9ead730
8 changed files with 212 additions and 0 deletions

View file

@ -0,0 +1,5 @@
---
"claude-dev": minor
---
Added X AI as a new provider with support for all current models including Grok-2 and Grok Vision. This integration enables users to connect to X AI's API using their API key and access models with context windows up to 131K tokens. The implementation includes proper handling for vision models and accurate pricing information.

View file

@ -17,6 +17,7 @@ import { QwenHandler } from "./providers/qwen"
import { MistralHandler } from "./providers/mistral"
import { VsCodeLmHandler } from "./providers/vscode-lm"
import { LiteLlmHandler } from "./providers/litellm"
import { XAIHandler } from "./providers/xai"
export interface ApiHandler {
createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream
@ -62,6 +63,8 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
return new VsCodeLmHandler(options)
case "litellm":
return new LiteLlmHandler(options)
case "xai":
return new XAIHandler(options)
default:
return new AnthropicHandler(options)
}

64
src/api/providers/xai.ts Normal file
View file

@ -0,0 +1,64 @@
import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"
import { ApiHandler } from "../"
import { ApiHandlerOptions, XAIModelId, ModelInfo, xaiDefaultModelId, xaiModels } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import { ApiStream } from "../transform/stream"
export class XAIHandler implements ApiHandler {
private options: ApiHandlerOptions
private client: OpenAI
constructor(options: ApiHandlerOptions) {
this.options = options
this.client = new OpenAI({
baseURL: "https://api.x.ai/v1",
apiKey: this.options.xaiApiKey,
})
}
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
const stream = await this.client.chat.completions.create({
model: this.getModel().id,
max_completion_tokens: this.getModel().info.maxTokens,
temperature: 0,
messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
stream: true,
stream_options: { include_usage: true },
})
for await (const chunk of stream) {
const delta = chunk.choices[0]?.delta
if (delta?.content) {
yield {
type: "text",
text: delta.content,
}
}
if (chunk.usage) {
yield {
type: "usage",
inputTokens: 0,
outputTokens: chunk.usage.completion_tokens || 0,
// @ts-ignore-next-line
cacheReadTokens: chunk.usage.prompt_cache_hit_tokens || 0,
// @ts-ignore-next-line
cacheWriteTokens: chunk.usage.prompt_cache_miss_tokens || 0,
}
}
}
}
getModel(): { id: XAIModelId; info: ModelInfo } {
const modelId = this.options.apiModelId
if (modelId && modelId in xaiModels) {
const id = modelId as XAIModelId
return { id, info: xaiModels[id] }
}
return {
id: xaiDefaultModelId,
info: xaiModels[xaiDefaultModelId],
}
}
}

View file

@ -59,6 +59,7 @@ type SecretKey =
| "liteLlmApiKey"
| "authToken"
| "authNonce"
| "xaiApiKey"
type GlobalStateKey =
| "apiProvider"
| "apiModelId"
@ -593,6 +594,7 @@ export class ClineProvider implements vscode.WebviewViewProvider {
liteLlmModelId,
liteLlmApiKey,
qwenApiLine,
xaiApiKey,
} = message.apiConfiguration
await this.updateGlobalState("apiProvider", apiProvider)
await this.updateGlobalState("apiModelId", apiModelId)
@ -624,6 +626,7 @@ export class ClineProvider implements vscode.WebviewViewProvider {
await this.storeSecret("qwenApiKey", qwenApiKey)
await this.storeSecret("mistralApiKey", mistralApiKey)
await this.storeSecret("liteLlmApiKey", liteLlmApiKey)
await this.storeSecret("xaiApiKey", xaiApiKey)
await this.updateGlobalState("azureApiVersion", azureApiVersion)
await this.updateGlobalState("openRouterModelId", openRouterModelId)
await this.updateGlobalState("openRouterModelInfo", openRouterModelInfo)
@ -1879,6 +1882,7 @@ Here is the project's README to help you get started:\n\n${mcpDetails.readmeCont
qwenApiLine,
liteLlmApiKey,
telemetrySetting,
xaiApiKey,
] = await Promise.all([
this.getGlobalState("apiProvider") as Promise<ApiProvider | undefined>,
this.getGlobalState("apiModelId") as Promise<string | undefined>,
@ -1931,6 +1935,7 @@ Here is the project's README to help you get started:\n\n${mcpDetails.readmeCont
this.getGlobalState("qwenApiLine") as Promise<string | undefined>,
this.getSecret("liteLlmApiKey") as Promise<string | undefined>,
this.getGlobalState("telemetrySetting") as Promise<TelemetrySetting | undefined>,
this.getSecret("xaiApiKey") as Promise<string | undefined>,
])
let apiProvider: ApiProvider
@ -1995,6 +2000,7 @@ Here is the project's README to help you get started:\n\n${mcpDetails.readmeCont
liteLlmBaseUrl,
liteLlmModelId,
liteLlmApiKey,
xaiApiKey,
},
lastShownAnnouncementId,
customInstructions,
@ -2138,6 +2144,7 @@ Here is the project's README to help you get started:\n\n${mcpDetails.readmeCont
"mistralApiKey",
"liteLlmApiKey",
"authToken",
"xaiApiKey",
]
for (const key of secretKeys) {
await this.storeSecret(key, undefined)

View file

@ -15,6 +15,7 @@ export type ApiProvider =
| "mistral"
| "vscode-lm"
| "litellm"
| "xai"
export interface ApiHandlerOptions {
apiModelId?: string
@ -56,6 +57,7 @@ export interface ApiHandlerOptions {
vsCodeLmModelSelector?: any
o3MiniReasoningEffort?: string
qwenApiLine?: string
xaiApiKey?: string
}
export type ApiConfiguration = ApiHandlerOptions & {
@ -799,3 +801,82 @@ export const liteLlmModelInfoSaneDefaults: ModelInfo = {
inputPrice: 0,
outputPrice: 0,
}
// X AI
// https://docs.x.ai/docs/api-reference
export type XAIModelId = keyof typeof xaiModels
export const xaiDefaultModelId: XAIModelId = "grok-2-latest"
export const xaiModels = {
"grok-2-latest": {
maxTokens: 8192,
contextWindow: 131072,
supportsImages: false,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 model - latest version with 131K context window",
},
"grok-2": {
maxTokens: 8192,
contextWindow: 131072,
supportsImages: false,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 model with 131K context window",
},
"grok-2-1212": {
maxTokens: 8192,
contextWindow: 131072,
supportsImages: false,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 model (version 1212) with 131K context window",
},
"grok-2-vision-latest": {
maxTokens: 8192,
contextWindow: 32768,
supportsImages: true,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 Vision model - latest version with image support and 32K context window",
},
"grok-2-vision": {
maxTokens: 8192,
contextWindow: 32768,
supportsImages: true,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 Vision model with image support and 32K context window",
},
"grok-2-vision-1212": {
maxTokens: 8192,
contextWindow: 32768,
supportsImages: true,
supportsPromptCache: false,
inputPrice: 2.0,
outputPrice: 10.0,
description: "X AI's Grok-2 Vision model (version 1212) with image support and 32K context window",
},
"grok-vision-beta": {
maxTokens: 8192,
contextWindow: 8192,
supportsImages: true,
supportsPromptCache: false,
inputPrice: 5.0,
outputPrice: 15.0,
description: "X AI's Grok Vision Beta model with image support and 8K context window",
},
"grok-beta": {
maxTokens: 8192,
contextWindow: 131072,
supportsImages: false,
supportsPromptCache: false,
inputPrice: 5.0,
outputPrice: 15.0,
description: "X AI's Grok Beta model (legacy) with 131K context window",
},
} as const satisfies Record<string, ModelInfo>

View file

@ -33,6 +33,8 @@ import {
openRouterDefaultModelInfo,
vertexDefaultModelId,
vertexModels,
xaiDefaultModelId,
xaiModels,
} from "../../../../src/shared/api"
import { ExtensionMessage } from "../../../../src/shared/ExtensionMessage"
import { useExtensionState } from "../../context/ExtensionStateContext"
@ -195,6 +197,7 @@ const ApiOptions = ({ showModelOptions, apiErrorMessage, modelIdErrorMessage, is
<VSCodeOption value="lmstudio">LM Studio</VSCodeOption>
<VSCodeOption value="ollama">Ollama</VSCodeOption>
<VSCodeOption value="litellm">LiteLLM</VSCodeOption>
<VSCodeOption value="xai">X AI</VSCodeOption>
</VSCodeDropdown>
</DropdownContainer>
@ -1122,6 +1125,46 @@ const ApiOptions = ({ showModelOptions, apiErrorMessage, modelIdErrorMessage, is
</div>
)}
{selectedProvider === "xai" && (
<div>
<VSCodeTextField
value={apiConfiguration?.xaiApiKey || ""}
style={{ width: "100%" }}
type="password"
onInput={handleInputChange("xaiApiKey")}
placeholder="Enter API Key...">
<span style={{ fontWeight: 500 }}>X AI API Key</span>
</VSCodeTextField>
<p
style={{
fontSize: "12px",
marginTop: 3,
color: "var(--vscode-descriptionForeground)",
}}>
This key is stored locally and only used to make API requests from this extension.
{!apiConfiguration?.xaiApiKey && (
<VSCodeLink href="https://x.ai" style={{ display: "inline", fontSize: "inherit" }}>
You can get an X AI API key by signing up here.
</VSCodeLink>
)}
</p>
{/* Note: To fully implement this, you would need to add a handler in ClineProvider.ts */}
{/* {apiConfiguration?.xaiApiKey && (
<button
onClick={() => {
vscode.postMessage({
type: "requestXAIModels",
text: apiConfiguration?.xaiApiKey,
})
}}
style={{ margin: "5px 0 0 0" }}
className="vscode-button">
Fetch Available Models
</button>
)} */}
</div>
)}
{apiErrorMessage && (
<p
style={{
@ -1152,6 +1195,7 @@ const ApiOptions = ({ showModelOptions, apiErrorMessage, modelIdErrorMessage, is
{selectedProvider === "deepseek" && createDropdown(deepSeekModels)}
{selectedProvider === "qwen" && createDropdown(qwenModels)}
{selectedProvider === "mistral" && createDropdown(mistralModels)}
{selectedProvider === "xai" && createDropdown(xaiModels)}
</DropdownContainer>
<ModelInfoView
@ -1403,6 +1447,8 @@ export function normalizeApiConfiguration(apiConfiguration?: ApiConfiguration):
selectedModelId: apiConfiguration?.liteLlmModelId || "",
selectedModelInfo: openAiModelInfoSaneDefaults,
}
case "xai":
return getProviderData(xaiModels, xaiDefaultModelId)
default:
return getProviderData(anthropicModels, anthropicDefaultModelId)
}

View file

@ -79,6 +79,7 @@ export const ExtensionStateContextProvider: React.FC<{
config.qwenApiKey,
config.mistralApiKey,
config.vsCodeLmModelSelector,
config.xaiApiKey,
].some((key) => key !== undefined)
: false
setShowWelcome(!hasKey)

View file

@ -38,6 +38,11 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
return "You must provide a valid API key or choose a different provider."
}
break
case "xai":
if (!apiConfiguration.xaiApiKey) {
return "You must provide a valid API key or choose a different provider."
}
break
case "qwen":
if (!apiConfiguration.qwenApiKey) {
return "You must provide a valid API key or choose a different provider."