Skip to content

Commit d02e9f9

Browse files
committed
feat: add thinking mode toggle for Z AI provider
- Add zaiEnableThinking field to provider settings schema - Update Z AI provider to conditionally enable thinking mode - Add UI checkbox for thinking mode in Z AI settings component - Add translation keys for the new UI elements - Add comprehensive tests for thinking mode functionality Addresses #8465
1 parent 13534cc commit d02e9f9

File tree

5 files changed

+157
-1
lines changed

5 files changed

+157
-1
lines changed

packages/types/src/provider-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -384,6 +384,7 @@ export type ZaiApiLine = z.infer<typeof zaiApiLineSchema>
384384
const zaiSchema = apiModelIdProviderModelSchema.extend({
385385
zaiApiKey: z.string().optional(),
386386
zaiApiLine: zaiApiLineSchema.optional(),
387+
zaiEnableThinking: z.boolean().optional(),
387388
})
388389

389390
const fireworksSchema = apiModelIdProviderModelSchema.extend({

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

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -262,5 +262,97 @@ describe("ZAiHandler", () => {
262262
undefined,
263263
)
264264
})
265+
266+
it("should include enable_thinking parameter when zaiEnableThinking is true", async () => {
267+
const handlerWithThinking = new ZAiHandler({
268+
zaiApiKey: "test-zai-api-key",
269+
zaiApiLine: "international",
270+
zaiEnableThinking: true,
271+
})
272+
273+
mockCreate.mockImplementationOnce(() => {
274+
return {
275+
[Symbol.asyncIterator]: () => ({
276+
async next() {
277+
return { done: true }
278+
},
279+
}),
280+
}
281+
})
282+
283+
const systemPrompt = "Test system prompt"
284+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]
285+
286+
const messageGenerator = handlerWithThinking.createMessage(systemPrompt, messages)
287+
await messageGenerator.next()
288+
289+
expect(mockCreate).toHaveBeenCalledWith(
290+
expect.objectContaining({
291+
enable_thinking: true,
292+
}),
293+
undefined,
294+
)
295+
})
296+
297+
it("should not include enable_thinking parameter when zaiEnableThinking is false", async () => {
298+
const handlerWithoutThinking = new ZAiHandler({
299+
zaiApiKey: "test-zai-api-key",
300+
zaiApiLine: "international",
301+
zaiEnableThinking: false,
302+
})
303+
304+
mockCreate.mockImplementationOnce(() => {
305+
return {
306+
[Symbol.asyncIterator]: () => ({
307+
async next() {
308+
return { done: true }
309+
},
310+
}),
311+
}
312+
})
313+
314+
const systemPrompt = "Test system prompt"
315+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]
316+
317+
const messageGenerator = handlerWithoutThinking.createMessage(systemPrompt, messages)
318+
await messageGenerator.next()
319+
320+
expect(mockCreate).toHaveBeenCalledWith(
321+
expect.not.objectContaining({
322+
enable_thinking: true,
323+
}),
324+
undefined,
325+
)
326+
})
327+
328+
it("should not include enable_thinking parameter when zaiEnableThinking is undefined", async () => {
329+
const handlerDefault = new ZAiHandler({
330+
zaiApiKey: "test-zai-api-key",
331+
zaiApiLine: "international",
332+
})
333+
334+
mockCreate.mockImplementationOnce(() => {
335+
return {
336+
[Symbol.asyncIterator]: () => ({
337+
async next() {
338+
return { done: true }
339+
},
340+
}),
341+
}
342+
})
343+
344+
const systemPrompt = "Test system prompt"
345+
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]
346+
347+
const messageGenerator = handlerDefault.createMessage(systemPrompt, messages)
348+
await messageGenerator.next()
349+
350+
expect(mockCreate).toHaveBeenCalledWith(
351+
expect.not.objectContaining({
352+
enable_thinking: true,
353+
}),
354+
undefined,
355+
)
356+
})
265357
})
266358
})

src/api/providers/zai.ts

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,12 @@ import {
88
ZAI_DEFAULT_TEMPERATURE,
99
zaiApiLineConfigs,
1010
} from "@roo-code/types"
11+
import { Anthropic } from "@anthropic-ai/sdk"
12+
import OpenAI from "openai"
1113

1214
import type { ApiHandlerOptions } from "../../shared/api"
15+
import type { ApiHandlerCreateMessageMetadata } from "../index"
16+
import { convertToOpenAiMessages } from "../transform/openai-format"
1317

1418
import { BaseOpenAiCompatibleProvider } from "./base-openai-compatible-provider"
1519

@@ -29,4 +33,43 @@ export class ZAiHandler extends BaseOpenAiCompatibleProvider<InternationalZAiMod
2933
defaultTemperature: ZAI_DEFAULT_TEMPERATURE,
3034
})
3135
}
36+
37+
protected override createStream(
38+
systemPrompt: string,
39+
messages: Anthropic.Messages.MessageParam[],
40+
metadata?: ApiHandlerCreateMessageMetadata,
41+
requestOptions?: OpenAI.RequestOptions,
42+
) {
43+
const {
44+
id: model,
45+
info: { maxTokens: max_tokens },
46+
} = this.getModel()
47+
48+
const temperature = this.options.modelTemperature ?? ZAI_DEFAULT_TEMPERATURE
49+
50+
// Build base parameters
51+
const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
52+
model,
53+
max_tokens,
54+
temperature,
55+
messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
56+
stream: true,
57+
stream_options: { include_usage: true },
58+
}
59+
60+
// Add thinking parameter for models that support it (GLM-4.6, etc.)
61+
// Only add if explicitly enabled via the zaiEnableThinking setting
62+
if (this.options.zaiEnableThinking === true) {
63+
// Z AI uses a custom parameter for thinking mode
64+
// This follows the pattern used by other providers with thinking support
65+
;(params as any).enable_thinking = true
66+
}
67+
68+
try {
69+
return this.client.chat.completions.create(params, requestOptions)
70+
} catch (error) {
71+
const errorMessage = error instanceof Error ? error.message : String(error)
72+
throw new Error(`Z AI completion error: ${errorMessage}`)
73+
}
74+
}
3275
}

webview-ui/src/components/settings/providers/ZAi.tsx

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { useCallback } from "react"
2-
import { VSCodeTextField, VSCodeDropdown, VSCodeOption } from "@vscode/webview-ui-toolkit/react"
2+
import { VSCodeTextField, VSCodeDropdown, VSCodeOption, VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react"
33

44
import { zaiApiLineConfigs, zaiApiLineSchema, type ProviderSettings } from "@roo-code/types"
55

@@ -28,6 +28,14 @@ export const ZAi = ({ apiConfiguration, setApiConfigurationField }: ZAiProps) =>
2828
[setApiConfigurationField],
2929
)
3030

31+
const handleCheckboxChange = useCallback(
32+
(field: keyof ProviderSettings) => (event: Event) => {
33+
const target = event.target as HTMLInputElement
34+
setApiConfigurationField(field, target.checked)
35+
},
36+
[setApiConfigurationField],
37+
)
38+
3139
return (
3240
<>
3341
<div>
@@ -73,6 +81,16 @@ export const ZAi = ({ apiConfiguration, setApiConfigurationField }: ZAiProps) =>
7381
</VSCodeButtonLink>
7482
)}
7583
</div>
84+
<div>
85+
<VSCodeCheckbox
86+
checked={apiConfiguration?.zaiEnableThinking || false}
87+
onChange={handleCheckboxChange("zaiEnableThinking")}>
88+
<label className="font-medium">{t("settings:providers.zaiEnableThinking")}</label>
89+
</VSCodeCheckbox>
90+
<div className="text-xs text-vscode-descriptionForeground mt-1 ml-6">
91+
{t("settings:providers.zaiEnableThinkingDescription")}
92+
</div>
93+
</div>
7694
</>
7795
)
7896
}

webview-ui/src/i18n/locales/en/settings.json

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,8 @@
302302
"getZaiApiKey": "Get Z AI API Key",
303303
"zaiEntrypoint": "Z AI Entrypoint",
304304
"zaiEntrypointDescription": "Please select the appropriate API entrypoint based on your location. If you are in China, choose open.bigmodel.cn. Otherwise, choose api.z.ai.",
305+
"zaiEnableThinking": "Enable Thinking Mode",
306+
"zaiEnableThinkingDescription": "Enable thinking/reasoning mode for models that support it (e.g., GLM-4.6, DeepSeek v3.2). This allows models to show their reasoning process before providing answers.",
305307
"geminiApiKey": "Gemini API Key",
306308
"getGroqApiKey": "Get Groq API Key",
307309
"groqApiKey": "Groq API Key",

0 commit comments

Comments
 (0)