Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions packages/types/src/model.ts
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,8 @@ export const modelInfoSchema = z.object({
// Capability flag to indicate whether the model supports an output verbosity parameter
supportsVerbosity: z.boolean().optional(),
supportsReasoningBudget: z.boolean().optional(),
// Capability flag to indicate whether the model supports simple on/off reasoning
supportsReasoning: z.boolean().optional(),
// Capability flag to indicate whether the model supports temperature parameter
supportsTemperature: z.boolean().optional(),
requiredReasoningBudget: z.boolean().optional(),
Expand Down
4 changes: 4 additions & 0 deletions packages/types/src/providers/zai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ export const internationalZAiModels = {
contextWindow: 131_072,
supportsImages: false,
supportsPromptCache: true,
supportsReasoning: true,
inputPrice: 0.6,
outputPrice: 2.2,
cacheWritesPrice: 0,
Expand Down Expand Up @@ -86,6 +87,7 @@ export const internationalZAiModels = {
contextWindow: 200_000,
supportsImages: false,
supportsPromptCache: true,
supportsReasoning: true,
inputPrice: 0.6,
outputPrice: 2.2,
cacheWritesPrice: 0,
Expand Down Expand Up @@ -114,6 +116,7 @@ export const mainlandZAiModels = {
contextWindow: 131_072,
supportsImages: false,
supportsPromptCache: true,
supportsReasoning: true,
inputPrice: 0.29,
outputPrice: 1.14,
cacheWritesPrice: 0,
Expand Down Expand Up @@ -184,6 +187,7 @@ export const mainlandZAiModels = {
contextWindow: 204_800,
supportsImages: false,
supportsPromptCache: true,
supportsReasoning: true,
inputPrice: 0.29,
outputPrice: 1.14,
cacheWritesPrice: 0,
Expand Down
138 changes: 138 additions & 0 deletions src/api/providers/__tests__/zai.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -295,5 +295,143 @@ describe("ZAiHandler", () => {
undefined,
)
})

describe("Reasoning functionality", () => {
it("should include thinking parameter when enableReasoningEffort is true and model supports reasoning in createMessage", async () => {
const handlerWithReasoning = new ZAiHandler({
apiModelId: "glm-4.6", // GLM-4.6 has supportsReasoning: true
zaiApiKey: "test-zai-api-key",
zaiApiLine: "international_coding",
enableReasoningEffort: true,
})

mockCreate.mockImplementationOnce(() => {
return {
[Symbol.asyncIterator]: () => ({
async next() {
return { done: true }
},
}),
}
})

const systemPrompt = "Test system prompt"
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]

const messageGenerator = handlerWithReasoning.createMessage(systemPrompt, messages)
await messageGenerator.next()

expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
thinking: { type: "enabled" },
}),
undefined,
)
})

it("should not include thinking parameter when enableReasoningEffort is false in createMessage", async () => {
const handlerWithoutReasoning = new ZAiHandler({
apiModelId: "glm-4.6", // GLM-4.6 has supportsReasoning: true
zaiApiKey: "test-zai-api-key",
zaiApiLine: "international_coding",
enableReasoningEffort: false,
})

mockCreate.mockImplementationOnce(() => {
return {
[Symbol.asyncIterator]: () => ({
async next() {
return { done: true }
},
}),
}
})

const systemPrompt = "Test system prompt"
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]

const messageGenerator = handlerWithoutReasoning.createMessage(systemPrompt, messages)
await messageGenerator.next()

expect(mockCreate).toHaveBeenCalledWith(
expect.not.objectContaining({
thinking: expect.anything(),
}),
undefined,
)
})

it("should not include thinking parameter when model does not support reasoning in createMessage", async () => {
const handlerWithNonReasoningModel = new ZAiHandler({
apiModelId: "glm-4-32b-0414-128k", // This model doesn't have supportsReasoning: true
zaiApiKey: "test-zai-api-key",
zaiApiLine: "international_coding",
enableReasoningEffort: true,
})

mockCreate.mockImplementationOnce(() => {
return {
[Symbol.asyncIterator]: () => ({
async next() {
return { done: true }
},
}),
}
})

const systemPrompt = "Test system prompt"
const messages: Anthropic.Messages.MessageParam[] = [{ role: "user", content: "Test message" }]

const messageGenerator = handlerWithNonReasoningModel.createMessage(systemPrompt, messages)
await messageGenerator.next()

expect(mockCreate).toHaveBeenCalledWith(
expect.not.objectContaining({
thinking: expect.anything(),
}),
undefined,
)
})

it("should include thinking parameter when enableReasoningEffort is true and model supports reasoning in completePrompt", async () => {
const handlerWithReasoning = new ZAiHandler({
apiModelId: "glm-4.5", // GLM-4.5 has supportsReasoning: true
zaiApiKey: "test-zai-api-key",
zaiApiLine: "international_coding",
enableReasoningEffort: true,
})

const expectedResponse = "This is a test response"
mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: expectedResponse } }] })

await handlerWithReasoning.completePrompt("test prompt")

expect(mockCreate).toHaveBeenCalledWith(
expect.objectContaining({
thinking: { type: "enabled" },
}),
)
})

it("should not include thinking parameter when enableReasoningEffort is false in completePrompt", async () => {
const handlerWithoutReasoning = new ZAiHandler({
apiModelId: "glm-4.5", // GLM-4.5 has supportsReasoning: true
zaiApiKey: "test-zai-api-key",
zaiApiLine: "international_coding",
enableReasoningEffort: false,
})

const expectedResponse = "This is a test response"
mockCreate.mockResolvedValueOnce({ choices: [{ message: { content: expectedResponse } }] })

await handlerWithoutReasoning.completePrompt("test prompt")

expect(mockCreate).toHaveBeenCalledWith(
expect.not.objectContaining({
thinking: expect.anything(),
}),
)
})
})
})
})
70 changes: 70 additions & 0 deletions src/api/providers/zai.ts
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,14 @@ import {
zaiApiLineConfigs,
} from "@roo-code/types"

import { Anthropic } from "@anthropic-ai/sdk"
import OpenAI from "openai"

import type { ApiHandlerOptions } from "../../shared/api"
import { getModelMaxOutputTokens } from "../../shared/api"
import { convertToOpenAiMessages } from "../transform/openai-format"
import type { ApiHandlerCreateMessageMetadata } from "../index"
import { handleOpenAIError } from "./utils/openai-error-handler"

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

Expand All @@ -30,4 +37,67 @@ export class ZAiHandler extends BaseOpenAiCompatibleProvider<string> {
defaultTemperature: ZAI_DEFAULT_TEMPERATURE,
})
}

protected override createStream(
systemPrompt: string,
messages: Anthropic.Messages.MessageParam[],
metadata?: ApiHandlerCreateMessageMetadata,
requestOptions?: OpenAI.RequestOptions,
) {
const { id: model, info } = this.getModel()

// Centralized cap: clamp to 20% of the context window (unless provider-specific exceptions apply)
const max_tokens =
getModelMaxOutputTokens({
modelId: model,
model: info,
settings: this.options,
format: "openai",
}) ?? undefined

const temperature = this.options.modelTemperature ?? this.defaultTemperature

const params: OpenAI.Chat.Completions.ChatCompletionCreateParamsStreaming = {
model,
max_tokens,
temperature,
messages: [{ role: "system", content: systemPrompt }, ...convertToOpenAiMessages(messages)],
stream: true,
stream_options: { include_usage: true },
}

// Add thinking parameter if reasoning is enabled and model supports it
const { id: modelId, info: modelInfo } = this.getModel()
if (this.options.enableReasoningEffort && modelInfo.supportsReasoning) {
;(params as any).thinking = { type: "enabled" }
}

try {
return this.client.chat.completions.create(params, requestOptions)
} catch (error) {
throw handleOpenAIError(error, this.providerName)
}
}

override async completePrompt(prompt: string): Promise<string> {
const { id: modelId } = this.getModel()

const params: OpenAI.Chat.Completions.ChatCompletionCreateParams = {
model: modelId,
messages: [{ role: "user", content: prompt }],
}

// Add thinking parameter if reasoning is enabled and model supports it
const { info: modelInfo } = this.getModel()
if (this.options.enableReasoningEffort && modelInfo.supportsReasoning) {
;(params as any).thinking = { type: "enabled" }
}

try {
const response = await this.client.chat.completions.create(params)
return response.choices[0]?.message.content || ""
} catch (error) {
throw handleOpenAIError(error, this.providerName)
}
}
}
16 changes: 16 additions & 0 deletions webview-ui/src/components/settings/ThinkingBudget.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ export const ThinkingBudget = ({ apiConfiguration, setApiConfigurationField, mod
const minThinkingTokens = isGemini25Pro ? GEMINI_25_PRO_MIN_THINKING_TOKENS : 1024

// Check model capabilities
const isReasoningSupported = !!modelInfo && modelInfo.supportsReasoning
const isReasoningBudgetSupported = !!modelInfo && modelInfo.supportsReasoningBudget
const isReasoningBudgetRequired = !!modelInfo && modelInfo.requiredReasoningBudget
const isReasoningEffortSupported = !!modelInfo && modelInfo.supportsReasoningEffort
Expand Down Expand Up @@ -103,6 +104,21 @@ export const ThinkingBudget = ({ apiConfiguration, setApiConfigurationField, mod
return null
}

// Models with supportsReasoning (binary reasoning) show a simple on/off toggle
if (isReasoningSupported) {
return (
<div className="flex flex-col gap-1">
<Checkbox
checked={enableReasoningEffort}
onChange={(checked: boolean) =>
setApiConfigurationField("enableReasoningEffort", checked === true)
}>
{t("settings:providers.useReasoning")}
</Checkbox>
</div>
)
}

return isReasoningBudgetSupported && !!modelInfo.maxTokens ? (
<>
{!isReasoningBudgetRequired && (
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,27 @@ describe("ThinkingBudget", () => {
expect(container.firstChild).toBeNull()
})

it("should render simple reasoning toggle when model has supportsReasoning (binary reasoning)", () => {
render(
<ThinkingBudget
{...defaultProps}
modelInfo={{
...mockModelInfo,
supportsReasoning: true,
supportsReasoningBudget: false,
supportsReasoningEffort: false,
}}
/>,
)

// Should show the reasoning checkbox (translation key)
expect(screen.getByText("settings:providers.useReasoning")).toBeInTheDocument()

// Should NOT show sliders or other complex reasoning controls
expect(screen.queryByTestId("reasoning-budget")).not.toBeInTheDocument()
expect(screen.queryByTestId("reasoning-effort")).not.toBeInTheDocument()
})

it("should render sliders when model supports thinking", () => {
render(<ThinkingBudget {...defaultProps} />)

Expand Down