Skip to content

Commit b97e57f

Browse files
AktsvigunAkim Tsvigunkitcat-dev
authored
Integration with Nebius AI Studio (RooCodeInc#2789)
* Integration with Nebius AI Studio added * changeset added * tests fixed * Nebius naming changed * bugs fixed * styling fixed * minor bug fixed * Remove obsolete ClineProvider.ts file that was causing build errors * redundant 'Model' section removed for Nebius * feat: add Nebius AI Studio to the list of inference providers (RooCodeInc#2789) - replace `nebiusModelId` to `apiModelId` - add our latest models - fix `getModel` method - fix spaces and the link to Nebius AI Studio api keys - delete extra code - add a couple of tests --------- Co-authored-by: Akim Tsvigun <[email protected]> Co-authored-by: Albert Abdulmanov <[email protected]>
1 parent cd927ae commit b97e57f

File tree

11 files changed

+272
-0
lines changed

11 files changed

+272
-0
lines changed

.changeset/poor-fireants-move.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"claude-dev": minor
3+
---
4+
5+
integration with nebius ai studio

src/api/index.ts

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ import { ApiStream, ApiStreamUsageChunk } from "./transform/stream"
1313
import { DeepSeekHandler } from "./providers/deepseek"
1414
import { RequestyHandler } from "./providers/requesty"
1515
import { TogetherHandler } from "./providers/together"
16+
import { NebiusHandler } from "./providers/nebius"
1617
import { QwenHandler } from "./providers/qwen"
1718
import { MistralHandler } from "./providers/mistral"
1819
import { DoubaoHandler } from "./providers/doubao"
@@ -75,6 +76,8 @@ export function buildApiHandler(configuration: ApiConfiguration): ApiHandler {
7576
return new ClineHandler(options)
7677
case "litellm":
7778
return new LiteLlmHandler(options)
79+
case "nebius":
80+
return new NebiusHandler(options)
7881
case "asksage":
7982
return new AskSageHandler(options)
8083
case "xai":

src/api/providers/nebius.ts

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
import { Anthropic } from "@anthropic-ai/sdk"
2+
import OpenAI from "openai"
3+
import { withRetry } from "../retry"
4+
import { ApiHandler } from "../index"
5+
import { convertToOpenAiMessages } from "../transform/openai-format"
6+
import { ApiStream } from "../transform/stream"
7+
import { convertToR1Format } from "../transform/r1-format"
8+
import { nebiusDefaultModelId, nebiusModels, type ModelInfo, type ApiHandlerOptions, type NebiusModelId } from "../../shared/api"
9+
10+
export class NebiusHandler implements ApiHandler {
11+
private client: OpenAI
12+
13+
constructor(private readonly options: ApiHandlerOptions) {
14+
this.client = new OpenAI({
15+
baseURL: "https://api.studio.nebius.ai/v1",
16+
apiKey: this.options.nebiusApiKey,
17+
})
18+
}
19+
20+
@withRetry()
21+
async *createMessage(systemPrompt: string, messages: Anthropic.Messages.MessageParam[]): ApiStream {
22+
const model = this.getModel()
23+
24+
const openAiMessages: OpenAI.Chat.ChatCompletionMessageParam[] = model.id.includes("DeepSeek-R1")
25+
? convertToR1Format([{ role: "user", content: systemPrompt }, ...messages])
26+
: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)]
27+
28+
const stream = await this.client.chat.completions.create({
29+
model: model.id,
30+
messages: openAiMessages,
31+
temperature: 0,
32+
stream: true,
33+
stream_options: { include_usage: true },
34+
})
35+
for await (const chunk of stream) {
36+
const delta = chunk.choices[0]?.delta
37+
if (delta?.content) {
38+
yield {
39+
type: "text",
40+
text: delta.content,
41+
}
42+
}
43+
44+
if (delta && "reasoning_content" in delta && delta.reasoning_content) {
45+
yield {
46+
type: "reasoning",
47+
reasoning: (delta.reasoning_content as string | undefined) || "",
48+
}
49+
}
50+
51+
if (chunk.usage) {
52+
yield {
53+
type: "usage",
54+
inputTokens: chunk.usage.prompt_tokens || 0,
55+
outputTokens: chunk.usage.completion_tokens || 0,
56+
}
57+
}
58+
}
59+
}
60+
61+
getModel(): { id: string; info: ModelInfo } {
62+
const modelId = this.options.apiModelId
63+
64+
if (modelId !== undefined && modelId in nebiusModels) {
65+
return { id: modelId, info: nebiusModels[modelId as NebiusModelId] }
66+
}
67+
return { id: nebiusDefaultModelId, info: nebiusModels[nebiusDefaultModelId] }
68+
}
69+
}

src/core/storage/state-keys.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export type SecretKey =
1919
| "authNonce"
2020
| "asksageApiKey"
2121
| "xaiApiKey"
22+
| "nebiusApiKey"
2223
| "sambanovaApiKey"
2324

2425
export type GlobalStateKey =

src/core/storage/state.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,7 @@ export async function getAllExtensionState(context: vscode.ExtensionContext) {
155155
thinkingBudgetTokens,
156156
reasoningEffort,
157157
sambanovaApiKey,
158+
nebiusApiKey,
158159
planActSeparateModelsSettingRaw,
159160
favoritedModelIds,
160161
globalClineRulesToggles,
@@ -242,6 +243,7 @@ export async function getAllExtensionState(context: vscode.ExtensionContext) {
242243
getGlobalState(context, "thinkingBudgetTokens") as Promise<number | undefined>,
243244
getGlobalState(context, "reasoningEffort") as Promise<string | undefined>,
244245
getSecret(context, "sambanovaApiKey") as Promise<string | undefined>,
246+
getSecret(context, "nebiusApiKey") as Promise<string | undefined>,
245247
getGlobalState(context, "planActSeparateModelsSetting") as Promise<boolean | undefined>,
246248
getGlobalState(context, "favoritedModelIds") as Promise<string[] | undefined>,
247249
getGlobalState(context, "globalClineRulesToggles") as Promise<ClineRulesToggles | undefined>,
@@ -353,6 +355,7 @@ export async function getAllExtensionState(context: vscode.ExtensionContext) {
353355
asksageApiUrl,
354356
xaiApiKey,
355357
sambanovaApiKey,
358+
nebiusApiKey,
356359
favoritedModelIds,
357360
requestTimeoutMs,
358361
},
@@ -445,6 +448,7 @@ export async function updateApiConfiguration(context: vscode.ExtensionContext, a
445448
reasoningEffort,
446449
clineApiKey,
447450
sambanovaApiKey,
451+
nebiusApiKey,
448452
favoritedModelIds,
449453
} = apiConfiguration
450454
await updateGlobalState(context, "apiProvider", apiProvider)
@@ -505,6 +509,7 @@ export async function updateApiConfiguration(context: vscode.ExtensionContext, a
505509
await updateGlobalState(context, "reasoningEffort", reasoningEffort)
506510
await storeSecret(context, "clineApiKey", clineApiKey)
507511
await storeSecret(context, "sambanovaApiKey", sambanovaApiKey)
512+
await storeSecret(context, "nebiusApiKey", nebiusApiKey)
508513
await updateGlobalState(context, "favoritedModelIds", favoritedModelIds)
509514
await updateGlobalState(context, "requestTimeoutMs", apiConfiguration.requestTimeoutMs)
510515
}
@@ -534,6 +539,7 @@ export async function resetExtensionState(context: vscode.ExtensionContext) {
534539
"asksageApiKey",
535540
"xaiApiKey",
536541
"sambanovaApiKey",
542+
"nebiusApiKey",
537543
]
538544
for (const key of secretKeys) {
539545
await storeSecret(context, key, undefined)

src/shared/api.ts

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ export type ApiProvider =
1919
| "vscode-lm"
2020
| "cline"
2121
| "litellm"
22+
| "nebius"
2223
| "fireworks"
2324
| "asksage"
2425
| "xai"
@@ -81,6 +82,7 @@ export interface ApiHandlerOptions {
8182
azureApiVersion?: string
8283
vsCodeLmModelSelector?: LanguageModelChatSelector
8384
qwenApiLine?: string
85+
nebiusApiKey?: string
8486
asksageApiUrl?: string
8587
asksageApiKey?: string
8688
xaiApiKey?: string
@@ -1560,6 +1562,93 @@ export const askSageModels = {
15601562
},
15611563
}
15621564

1565+
// Nebius AI Studio
1566+
// https://docs.nebius.com/studio/inference/models
1567+
export const nebiusModels = {
1568+
"deepseek-ai/DeepSeek-V3": {
1569+
maxTokens: 32_000,
1570+
contextWindow: 96_000,
1571+
supportsImages: false,
1572+
supportsPromptCache: false,
1573+
inputPrice: 0.5,
1574+
outputPrice: 1.5,
1575+
},
1576+
"deepseek-ai/DeepSeek-V3-0324-fast": {
1577+
maxTokens: 128_000,
1578+
contextWindow: 128_000,
1579+
supportsImages: false,
1580+
supportsPromptCache: false,
1581+
inputPrice: 2,
1582+
outputPrice: 6,
1583+
},
1584+
"deepseek-ai/DeepSeek-R1": {
1585+
maxTokens: 32_000,
1586+
contextWindow: 96_000,
1587+
supportsImages: false,
1588+
supportsPromptCache: false,
1589+
inputPrice: 0.8,
1590+
outputPrice: 2.4,
1591+
},
1592+
"deepseek-ai/DeepSeek-R1-fast": {
1593+
maxTokens: 32_000,
1594+
contextWindow: 96_000,
1595+
supportsImages: false,
1596+
supportsPromptCache: false,
1597+
inputPrice: 2,
1598+
outputPrice: 6,
1599+
},
1600+
"meta-llama/Llama-3.3-70B-Instruct-fast": {
1601+
maxTokens: 32_000,
1602+
contextWindow: 96_000,
1603+
supportsImages: false,
1604+
supportsPromptCache: false,
1605+
inputPrice: 0.25,
1606+
outputPrice: 0.75,
1607+
},
1608+
"Qwen/Qwen2.5-32B-Instruct-fast": {
1609+
maxTokens: 8_192,
1610+
contextWindow: 32_768,
1611+
supportsImages: false,
1612+
supportsPromptCache: false,
1613+
inputPrice: 0.13,
1614+
outputPrice: 0.4,
1615+
},
1616+
"Qwen/Qwen2.5-Coder-32B-Instruct-fast": {
1617+
maxTokens: 128_000,
1618+
contextWindow: 128_000,
1619+
supportsImages: false,
1620+
supportsPromptCache: false,
1621+
inputPrice: 0.1,
1622+
outputPrice: 0.3,
1623+
},
1624+
"Qwen/Qwen3-4B-fast": {
1625+
maxTokens: 32_000,
1626+
contextWindow: 41_000,
1627+
supportsImages: false,
1628+
supportsPromptCache: false,
1629+
inputPrice: 0.08,
1630+
outputPrice: 0.24,
1631+
},
1632+
"Qwen/Qwen3-30B-A3B-fast": {
1633+
maxTokens: 32_000,
1634+
contextWindow: 41_000,
1635+
supportsImages: false,
1636+
supportsPromptCache: false,
1637+
inputPrice: 0.3,
1638+
outputPrice: 0.9,
1639+
},
1640+
"Qwen/Qwen3-235B-A22B": {
1641+
maxTokens: 32_000,
1642+
contextWindow: 41_000,
1643+
supportsImages: false,
1644+
supportsPromptCache: false,
1645+
inputPrice: 0.2,
1646+
outputPrice: 0.6,
1647+
},
1648+
} as const satisfies Record<string, ModelInfo>
1649+
export type NebiusModelId = keyof typeof nebiusModels
1650+
export const nebiusDefaultModelId = "Qwen/Qwen2.5-32B-Instruct-fast" satisfies NebiusModelId
1651+
15631652
// X AI
15641653
// https://docs.x.ai/docs/api-reference
15651654
export type XAIModelId = keyof typeof xaiModels

webview-ui/src/components/chat/ChatTextArea.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1071,6 +1071,7 @@ const ChatTextArea = forwardRef<HTMLTextAreaElement, ChatTextAreaProps>(
10711071
case "litellm":
10721072
return `${selectedProvider}:${apiConfiguration.liteLlmModelId}`
10731073
case "requesty":
1074+
return `${selectedProvider}:${apiConfiguration.requestyModelId}`
10741075
case "anthropic":
10751076
case "openrouter":
10761077
default:

webview-ui/src/components/settings/ApiOptions.tsx

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,8 @@ import {
4747
askSageDefaultURL,
4848
xaiDefaultModelId,
4949
xaiModels,
50+
nebiusModels,
51+
nebiusDefaultModelId,
5052
sambanovaModels,
5153
sambanovaDefaultModelId,
5254
doubaoModels,
@@ -324,6 +326,7 @@ const ApiOptions = ({
324326
<VSCodeOption value="lmstudio">LM Studio</VSCodeOption>
325327
<VSCodeOption value="ollama">Ollama</VSCodeOption>
326328
<VSCodeOption value="litellm">LiteLLM</VSCodeOption>
329+
<VSCodeOption value="nebius">Nebius AI Studio</VSCodeOption>
327330
<VSCodeOption value="asksage">AskSage</VSCodeOption>
328331
<VSCodeOption value="xai">xAI</VSCodeOption>
329332
<VSCodeOption value="sambanova">SambaNova</VSCodeOption>
@@ -1884,6 +1887,52 @@ const ApiOptions = ({
18841887
</div>
18851888
)}
18861889

1890+
{selectedProvider === "nebius" && (
1891+
<div>
1892+
<VSCodeTextField
1893+
value={apiConfiguration?.nebiusApiKey || ""}
1894+
style={{ width: "100%" }}
1895+
type="password"
1896+
onInput={handleInputChange("nebiusApiKey")}
1897+
placeholder="Enter API Key...">
1898+
<span style={{ fontWeight: 500 }}>Nebius API Key</span>
1899+
</VSCodeTextField>
1900+
<p
1901+
style={{
1902+
fontSize: "12px",
1903+
marginTop: 3,
1904+
color: "var(--vscode-descriptionForeground)",
1905+
}}>
1906+
This key is stored locally and only used to make API requests from this extension.{" "}
1907+
{!apiConfiguration?.nebiusApiKey && (
1908+
<VSCodeLink
1909+
href="https://studio.nebius.com/settings/api-keys"
1910+
style={{
1911+
display: "inline",
1912+
fontSize: "inherit",
1913+
}}>
1914+
You can get a Nebius API key by signing up here.{" "}
1915+
</VSCodeLink>
1916+
)}
1917+
<span style={{ color: "var(--vscode-errorForeground)" }}>
1918+
(<span style={{ fontWeight: 500 }}>Note:</span> Cline uses complex prompts and works best with Claude
1919+
models. Less capable models may not work as expected.)
1920+
</span>
1921+
</p>
1922+
</div>
1923+
)}
1924+
1925+
{apiErrorMessage && (
1926+
<p
1927+
style={{
1928+
margin: "-10px 0 4px 0",
1929+
fontSize: 12,
1930+
color: "var(--vscode-errorForeground)",
1931+
}}>
1932+
{apiErrorMessage}
1933+
</p>
1934+
)}
1935+
18871936
{selectedProvider === "xai" && (
18881937
<div>
18891938
<VSCodeTextField
@@ -1900,6 +1949,10 @@ const ApiOptions = ({
19001949
marginTop: 3,
19011950
color: "var(--vscode-descriptionForeground)",
19021951
}}>
1952+
<span style={{ color: "var(--vscode-errorForeground)" }}>
1953+
(<span style={{ fontWeight: 500 }}>Note:</span> Cline uses complex prompts and works best with Claude
1954+
models. Less capable models may not work as expected.)
1955+
</span>
19031956
This key is stored locally and only used to make API requests from this extension.
19041957
{!apiConfiguration?.xaiApiKey && (
19051958
<VSCodeLink href="https://x.ai" style={{ display: "inline", fontSize: "inherit" }}>
@@ -2072,6 +2125,7 @@ const ApiOptions = ({
20722125
{selectedProvider === "asksage" && createDropdown(askSageModels)}
20732126
{selectedProvider === "xai" && createDropdown(xaiModels)}
20742127
{selectedProvider === "sambanova" && createDropdown(sambanovaModels)}
2128+
{selectedProvider === "nebius" && createDropdown(nebiusModels)}
20752129
</DropdownContainer>
20762130

20772131
{((selectedProvider === "anthropic" && selectedModelId === "claude-3-7-sonnet-20250219") ||
@@ -2488,6 +2542,8 @@ export function normalizeApiConfiguration(apiConfiguration?: ApiConfiguration):
24882542
}
24892543
case "xai":
24902544
return getProviderData(xaiModels, xaiDefaultModelId)
2545+
case "nebius":
2546+
return getProviderData(nebiusModels, nebiusDefaultModelId)
24912547
case "sambanova":
24922548
return getProviderData(sambanovaModels, sambanovaDefaultModelId)
24932549
default:

0 commit comments

Comments
 (0)