Skip to content

Commit 8e233b3

Browse files
committed
Fix #5141: Add option to disable Gemini intermediate reasoning streaming
- Add geminiDisableIntermediateReasoning setting to Gemini provider schema - Modify Gemini provider to suppress reasoning chunks when setting is enabled - Add comprehensive tests for both enabled and disabled reasoning scenarios - Reasoning still works but intermediate results aren't streamed when disabled
1 parent 3a8ba27 commit 8e233b3

File tree

3 files changed

+116
-1
lines changed

3 files changed

+116
-1
lines changed

packages/types/src/provider-settings.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,7 @@ const lmStudioSchema = baseProviderSettingsSchema.extend({
157157
const geminiSchema = apiModelIdProviderModelSchema.extend({
158158
geminiApiKey: z.string().optional(),
159159
googleGeminiBaseUrl: z.string().optional(),
160+
geminiDisableIntermediateReasoning: z.boolean().optional(),
160161
})
161162

162163
const geminiCliSchema = apiModelIdProviderModelSchema.extend({

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,119 @@ describe("GeminiHandler", () => {
8989
)
9090
})
9191

92+
it("should handle reasoning chunks correctly when intermediate reasoning is enabled", async () => {
93+
// Setup the mock implementation to return an async generator with reasoning chunks
94+
;(handler["client"].models.generateContentStream as any).mockResolvedValue({
95+
[Symbol.asyncIterator]: async function* () {
96+
yield {
97+
candidates: [
98+
{
99+
content: {
100+
parts: [{ thought: true, text: "Let me think about this..." }, { text: "Hello" }],
101+
},
102+
},
103+
],
104+
}
105+
yield {
106+
candidates: [
107+
{
108+
content: {
109+
parts: [{ thought: true, text: "I need to consider..." }, { text: " world!" }],
110+
},
111+
},
112+
],
113+
}
114+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, thoughtsTokenCount: 20 } }
115+
},
116+
})
117+
118+
const stream = handler.createMessage(systemPrompt, mockMessages)
119+
const chunks = []
120+
121+
for await (const chunk of stream) {
122+
chunks.push(chunk)
123+
}
124+
125+
// Should have 6 chunks: 2 reasoning + 2 text + 2 reasoning + 2 text + usage
126+
expect(chunks.length).toBe(5)
127+
expect(chunks[0]).toEqual({ type: "reasoning", text: "Let me think about this..." })
128+
expect(chunks[1]).toEqual({ type: "text", text: "Hello" })
129+
expect(chunks[2]).toEqual({ type: "reasoning", text: "I need to consider..." })
130+
expect(chunks[3]).toEqual({ type: "text", text: " world!" })
131+
expect(chunks[4]).toEqual({
132+
type: "usage",
133+
inputTokens: 10,
134+
outputTokens: 5,
135+
reasoningTokens: 20,
136+
})
137+
})
138+
139+
it("should suppress reasoning chunks when geminiDisableIntermediateReasoning is enabled", async () => {
140+
// Create a new handler with the setting enabled
141+
const handlerWithDisabledReasoning = new GeminiHandler({
142+
apiKey: "test-key",
143+
apiModelId: GEMINI_20_FLASH_THINKING_NAME,
144+
geminiApiKey: "test-key",
145+
geminiDisableIntermediateReasoning: true,
146+
})
147+
148+
// Replace the client with our mock
149+
handlerWithDisabledReasoning["client"] = {
150+
models: {
151+
generateContentStream: vitest.fn(),
152+
generateContent: vitest.fn(),
153+
getGenerativeModel: vitest.fn(),
154+
},
155+
} as any
156+
157+
// Setup the mock implementation to return an async generator with reasoning chunks
158+
;(handlerWithDisabledReasoning["client"].models.generateContentStream as any).mockResolvedValue({
159+
[Symbol.asyncIterator]: async function* () {
160+
yield {
161+
candidates: [
162+
{
163+
content: {
164+
parts: [{ thought: true, text: "Let me think about this..." }, { text: "Hello" }],
165+
},
166+
},
167+
],
168+
}
169+
yield {
170+
candidates: [
171+
{
172+
content: {
173+
parts: [{ thought: true, text: "I need to consider..." }, { text: " world!" }],
174+
},
175+
},
176+
],
177+
}
178+
yield { usageMetadata: { promptTokenCount: 10, candidatesTokenCount: 5, thoughtsTokenCount: 20 } }
179+
},
180+
})
181+
182+
const stream = handlerWithDisabledReasoning.createMessage(systemPrompt, mockMessages)
183+
const chunks = []
184+
185+
for await (const chunk of stream) {
186+
chunks.push(chunk)
187+
}
188+
189+
// Should have only 3 chunks: 2 text + usage (reasoning chunks should be suppressed)
190+
expect(chunks.length).toBe(3)
191+
expect(chunks[0]).toEqual({ type: "text", text: "Hello" })
192+
expect(chunks[1]).toEqual({ type: "text", text: " world!" })
193+
expect(chunks[2]).toEqual({
194+
type: "usage",
195+
inputTokens: 10,
196+
outputTokens: 5,
197+
reasoningTokens: 20,
198+
})
199+
200+
// Verify no reasoning chunks were yielded
201+
const reasoningChunks = chunks.filter((chunk) => chunk.type === "reasoning")
202+
expect(reasoningChunks.length).toBe(0)
203+
})
204+
92205
it("should handle API errors", async () => {
93206
const mockError = new Error("Gemini API error")
94207
;(handler["client"].models.generateContentStream as any).mockRejectedValue(mockError)

src/api/providers/gemini.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,8 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
8989
for (const part of candidate.content.parts) {
9090
if (part.thought) {
9191
// This is a thinking/reasoning part
92-
if (part.text) {
92+
// Only yield reasoning chunks if intermediate reasoning is not disabled
93+
if (part.text && !this.options.geminiDisableIntermediateReasoning) {
9394
yield { type: "reasoning", text: part.text }
9495
}
9596
} else {

0 commit comments

Comments
 (0)