Skip to content

Commit 3889788

Browse files
committed
fix: Implement review comments for PR #4447 - Bedrock Extended Thinking
- Refactor Bedrock provider for clarity and type safety. - Extract constants and helper methods. - Update documentation. - Rename and update reasoning tests to Vitest. - Ensure all tests pass.
1 parent 3d794b9 commit 3889788

File tree

3 files changed

+394
-326
lines changed

3 files changed

+394
-326
lines changed
Lines changed: 286 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,286 @@
1+
import { vi, describe, it, expect, beforeEach } from "vitest"
2+
3+
// Mock AWS SDK modules before importing the handler
4+
vi.mock("@aws-sdk/credential-providers", () => ({
5+
fromIni: vi.fn(),
6+
}))
7+
8+
// Define a shared mock for the send function that will be used by all instances
9+
const sharedMockSend = vi.fn()
10+
11+
vi.mock("@aws-sdk/client-bedrock-runtime", () => ({
12+
BedrockRuntimeClient: vi.fn().mockImplementation(() => ({
13+
// Ensure all instances of BedrockRuntimeClient use the sharedMockSend
14+
send: sharedMockSend,
15+
config: { region: "us-east-1" },
16+
})),
17+
ConverseStreamCommand: vi.fn(), // This will be the mock constructor for ConverseStreamCommand
18+
ConverseCommand: vi.fn(),
19+
}))
20+
21+
// Import after mocks are set up
22+
import { AwsBedrockHandler } from "../bedrock"
23+
// Import ConverseStreamCommand to check its mock constructor (which is vi.fn() from the mock factory)
24+
import { ConverseStreamCommand } from "@aws-sdk/client-bedrock-runtime"
25+
26+
describe("AwsBedrockHandler - Extended Thinking", () => {
27+
let handler: AwsBedrockHandler
28+
// This will hold the reference to sharedMockSend for use in tests
29+
let mockSend: typeof sharedMockSend
30+
31+
const mockOptions = {
32+
awsRegion: "us-east-1",
33+
apiModelId: "anthropic.claude-3-7-sonnet-20241029-v1:0",
34+
enableReasoningEffort: false, // Default to false
35+
modelTemperature: 0.7,
36+
}
37+
38+
beforeEach(() => {
39+
// Clear all mocks. This will clear sharedMockSend and the ConverseStreamCommand mock constructor.
40+
vi.clearAllMocks()
41+
// Assign the shared mock to mockSend so tests can configure it.
42+
mockSend = sharedMockSend
43+
44+
// AwsBedrockHandler will instantiate BedrockRuntimeClient, which will get the sharedMockSend.
45+
handler = new AwsBedrockHandler(mockOptions)
46+
})
47+
48+
describe("Extended Thinking Configuration", () => {
49+
it("should NOT enable extended thinking by default", async () => {
50+
// Setup mock response
51+
mockSend.mockResolvedValue({
52+
stream: (async function* () {
53+
yield { messageStart: { role: "assistant" } }
54+
yield {
55+
contentBlockStart: {
56+
start: { text: "Hello" },
57+
contentBlockIndex: 0,
58+
},
59+
}
60+
yield { messageStop: { stopReason: "end_turn" } }
61+
})(),
62+
})
63+
64+
// Create message
65+
const messages = [{ role: "user" as const, content: "Test message" }]
66+
const stream = handler.createMessage("", messages)
67+
68+
// Consume stream
69+
const chunks = []
70+
for await (const chunk of stream) {
71+
chunks.push(chunk)
72+
}
73+
74+
// Verify the command was called
75+
expect(ConverseStreamCommand).toHaveBeenCalled()
76+
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
77+
78+
// Extended thinking should NOT be enabled by default
79+
expect(payload.anthropic_version).toBeUndefined()
80+
expect(payload.additionalModelRequestFields).toBeUndefined()
81+
expect(payload.inferenceConfig.temperature).toBeDefined()
82+
expect(payload.inferenceConfig.topP).toBeDefined()
83+
})
84+
85+
it("should enable extended thinking when explicitly enabled with reasoning budget", async () => {
86+
// Enable reasoning mode with thinking tokens
87+
handler = new AwsBedrockHandler({
88+
...mockOptions,
89+
enableReasoningEffort: true,
90+
modelMaxThinkingTokens: 5000,
91+
})
92+
93+
// Setup mock response with thinking blocks
94+
mockSend.mockResolvedValue({
95+
stream: (async function* () {
96+
yield { messageStart: { role: "assistant" } }
97+
yield {
98+
contentBlockStart: {
99+
contentBlock: {
100+
type: "thinking",
101+
thinking: "Let me think about this...",
102+
},
103+
contentBlockIndex: 0,
104+
},
105+
}
106+
yield {
107+
contentBlockStart: {
108+
start: { text: "Here is my response" },
109+
contentBlockIndex: 1,
110+
},
111+
}
112+
yield { messageStop: { stopReason: "end_turn" } }
113+
})(),
114+
})
115+
116+
// Create message
117+
const messages = [{ role: "user" as const, content: "Test message" }]
118+
const stream = handler.createMessage("", messages)
119+
120+
// Consume stream
121+
const chunks = []
122+
for await (const chunk of stream) {
123+
chunks.push(chunk)
124+
}
125+
126+
// Verify the command was called
127+
expect(ConverseStreamCommand).toHaveBeenCalled()
128+
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
129+
130+
// Extended thinking should be enabled
131+
expect(payload.anthropic_version).toBe("bedrock-20250514")
132+
expect(payload.additionalModelRequestFields).toEqual({
133+
thinking: {
134+
type: "enabled",
135+
budget_tokens: 5000,
136+
},
137+
})
138+
// Temperature and topP should be removed
139+
expect(payload.inferenceConfig.temperature).toBeUndefined()
140+
expect(payload.inferenceConfig.topP).toBeUndefined()
141+
142+
// Verify thinking content was processed
143+
const reasoningChunk = chunks.find((c) => c.type === "reasoning")
144+
expect(reasoningChunk).toBeDefined()
145+
expect(reasoningChunk?.text).toBe("Let me think about this...")
146+
})
147+
148+
it("should NOT enable extended thinking for unsupported models", async () => {
149+
// Use a model that doesn't support reasoning
150+
handler = new AwsBedrockHandler({
151+
...mockOptions,
152+
apiModelId: "anthropic.claude-3-haiku-20240307-v1:0",
153+
enableReasoningEffort: true,
154+
modelMaxThinkingTokens: 5000,
155+
})
156+
157+
// Setup mock response
158+
mockSend.mockResolvedValue({
159+
stream: (async function* () {
160+
yield { messageStart: { role: "assistant" } }
161+
yield {
162+
contentBlockStart: {
163+
start: { text: "Hello" },
164+
contentBlockIndex: 0,
165+
},
166+
}
167+
yield { messageStop: { stopReason: "end_turn" } }
168+
})(),
169+
})
170+
171+
// Create message
172+
const messages = [{ role: "user" as const, content: "Test message" }]
173+
const stream = handler.createMessage("", messages)
174+
175+
// Consume stream
176+
const chunks = []
177+
for await (const chunk of stream) {
178+
chunks.push(chunk)
179+
}
180+
181+
// Verify the command was called
182+
expect(ConverseStreamCommand).toHaveBeenCalled()
183+
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
184+
185+
// Extended thinking should NOT be enabled for unsupported models
186+
expect(payload.anthropic_version).toBeUndefined()
187+
expect(payload.additionalModelRequestFields).toBeUndefined()
188+
expect(payload.inferenceConfig.temperature).toBeDefined()
189+
expect(payload.inferenceConfig.topP).toBeDefined()
190+
})
191+
})
192+
193+
describe("Stream Processing", () => {
194+
it("should handle thinking delta events", async () => {
195+
// Enable reasoning mode
196+
handler = new AwsBedrockHandler({
197+
...mockOptions,
198+
enableReasoningEffort: true,
199+
modelMaxThinkingTokens: 5000,
200+
})
201+
202+
// Setup mock response with thinking deltas
203+
mockSend.mockResolvedValue({
204+
stream: (async function* () {
205+
yield { messageStart: { role: "assistant" } }
206+
yield {
207+
contentBlockDelta: {
208+
delta: {
209+
type: "thinking_delta",
210+
thinking: "First part of thinking...",
211+
},
212+
contentBlockIndex: 0,
213+
},
214+
}
215+
yield {
216+
contentBlockDelta: {
217+
delta: {
218+
type: "thinking_delta",
219+
thinking: " Second part of thinking.",
220+
},
221+
contentBlockIndex: 0,
222+
},
223+
}
224+
yield { messageStop: { stopReason: "end_turn" } }
225+
})(),
226+
})
227+
228+
// Create message
229+
const messages = [{ role: "user" as const, content: "Test message" }]
230+
const stream = handler.createMessage("", messages)
231+
232+
// Consume stream
233+
const chunks = []
234+
for await (const chunk of stream) {
235+
chunks.push(chunk)
236+
}
237+
238+
// Verify thinking deltas were processed
239+
const reasoningChunks = chunks.filter((c) => c.type === "reasoning")
240+
expect(reasoningChunks).toHaveLength(2)
241+
expect(reasoningChunks[0].text).toBe("First part of thinking...")
242+
expect(reasoningChunks[1].text).toBe(" Second part of thinking.")
243+
})
244+
245+
it("should handle signature delta events as reasoning", async () => {
246+
// Enable reasoning mode
247+
handler = new AwsBedrockHandler({
248+
...mockOptions,
249+
enableReasoningEffort: true,
250+
modelMaxThinkingTokens: 5000,
251+
})
252+
253+
// Setup mock response with signature deltas
254+
mockSend.mockResolvedValue({
255+
stream: (async function* () {
256+
yield { messageStart: { role: "assistant" } }
257+
yield {
258+
contentBlockDelta: {
259+
delta: {
260+
type: "signature_delta",
261+
signature: "[Signature content]",
262+
},
263+
contentBlockIndex: 0,
264+
},
265+
}
266+
yield { messageStop: { stopReason: "end_turn" } }
267+
})(),
268+
})
269+
270+
// Create message
271+
const messages = [{ role: "user" as const, content: "Test message" }]
272+
const stream = handler.createMessage("", messages)
273+
274+
// Consume stream
275+
const chunks = []
276+
for await (const chunk of stream) {
277+
chunks.push(chunk)
278+
}
279+
280+
// Verify signature delta was processed as reasoning
281+
const reasoningChunk = chunks.find((c) => c.type === "reasoning")
282+
expect(reasoningChunk).toBeDefined()
283+
expect(reasoningChunk?.text).toBe("[Signature content]")
284+
})
285+
})
286+
})

0 commit comments

Comments
 (0)