mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-05 08:10:14 +00:00
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.
This commit is contained in:
hannesrudolph
parent
3d794b9f71
commit
388978807a
3 changed files with 394 additions and 326 deletions
286
src/api/providers/__tests__/bedrock-reasoning.spec.ts
Normal file
286
src/api/providers/__tests__/bedrock-reasoning.spec.ts
Normal file
|
|
@ -0,0 +1,286 @@
|
|||
import { vi, describe, it, expect, beforeEach } from "vitest"
|
||||
|
||||
// Mock AWS SDK modules before importing the handler
|
||||
vi.mock("@aws-sdk/credential-providers", () => ({
|
||||
fromIni: vi.fn(),
|
||||
}))
|
||||
|
||||
// Define a shared mock for the send function that will be used by all instances
|
||||
const sharedMockSend = vi.fn()
|
||||
|
||||
vi.mock("@aws-sdk/client-bedrock-runtime", () => ({
|
||||
BedrockRuntimeClient: vi.fn().mockImplementation(() => ({
|
||||
// Ensure all instances of BedrockRuntimeClient use the sharedMockSend
|
||||
send: sharedMockSend,
|
||||
config: { region: "us-east-1" },
|
||||
})),
|
||||
ConverseStreamCommand: vi.fn(), // This will be the mock constructor for ConverseStreamCommand
|
||||
ConverseCommand: vi.fn(),
|
||||
}))
|
||||
|
||||
// Import after mocks are set up
|
||||
import { AwsBedrockHandler } from "../bedrock"
|
||||
// Import ConverseStreamCommand to check its mock constructor (which is vi.fn() from the mock factory)
|
||||
import { ConverseStreamCommand } from "@aws-sdk/client-bedrock-runtime"
|
||||
|
||||
describe("AwsBedrockHandler - Extended Thinking", () => {
|
||||
let handler: AwsBedrockHandler
|
||||
// This will hold the reference to sharedMockSend for use in tests
|
||||
let mockSend: typeof sharedMockSend
|
||||
|
||||
const mockOptions = {
|
||||
awsRegion: "us-east-1",
|
||||
apiModelId: "anthropic.claude-3-7-sonnet-20241029-v1:0",
|
||||
enableReasoningEffort: false, // Default to false
|
||||
modelTemperature: 0.7,
|
||||
}
|
||||
|
||||
beforeEach(() => {
|
||||
// Clear all mocks. This will clear sharedMockSend and the ConverseStreamCommand mock constructor.
|
||||
vi.clearAllMocks()
|
||||
// Assign the shared mock to mockSend so tests can configure it.
|
||||
mockSend = sharedMockSend
|
||||
|
||||
// AwsBedrockHandler will instantiate BedrockRuntimeClient, which will get the sharedMockSend.
|
||||
handler = new AwsBedrockHandler(mockOptions)
|
||||
})
|
||||
|
||||
describe("Extended Thinking Configuration", () => {
|
||||
it("should NOT enable extended thinking by default", async () => {
|
||||
// Setup mock response
|
||||
mockSend.mockResolvedValue({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
start: { text: "Hello" },
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})(),
|
||||
})
|
||||
|
||||
// Create message
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("", messages)
|
||||
|
||||
// Consume stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(ConverseStreamCommand).toHaveBeenCalled()
|
||||
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
|
||||
|
||||
// Extended thinking should NOT be enabled by default
|
||||
expect(payload.anthropic_version).toBeUndefined()
|
||||
expect(payload.additionalModelRequestFields).toBeUndefined()
|
||||
expect(payload.inferenceConfig.temperature).toBeDefined()
|
||||
expect(payload.inferenceConfig.topP).toBeDefined()
|
||||
})
|
||||
|
||||
it("should enable extended thinking when explicitly enabled with reasoning budget", async () => {
|
||||
// Enable reasoning mode with thinking tokens
|
||||
handler = new AwsBedrockHandler({
|
||||
...mockOptions,
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Setup mock response with thinking blocks
|
||||
mockSend.mockResolvedValue({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
contentBlock: {
|
||||
type: "thinking",
|
||||
thinking: "Let me think about this...",
|
||||
},
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
start: { text: "Here is my response" },
|
||||
contentBlockIndex: 1,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})(),
|
||||
})
|
||||
|
||||
// Create message
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("", messages)
|
||||
|
||||
// Consume stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(ConverseStreamCommand).toHaveBeenCalled()
|
||||
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
|
||||
|
||||
// Extended thinking should be enabled
|
||||
expect(payload.anthropic_version).toBe("bedrock-20250514")
|
||||
expect(payload.additionalModelRequestFields).toEqual({
|
||||
thinking: {
|
||||
type: "enabled",
|
||||
budget_tokens: 5000,
|
||||
},
|
||||
})
|
||||
// Temperature and topP should be removed
|
||||
expect(payload.inferenceConfig.temperature).toBeUndefined()
|
||||
expect(payload.inferenceConfig.topP).toBeUndefined()
|
||||
|
||||
// Verify thinking content was processed
|
||||
const reasoningChunk = chunks.find((c) => c.type === "reasoning")
|
||||
expect(reasoningChunk).toBeDefined()
|
||||
expect(reasoningChunk?.text).toBe("Let me think about this...")
|
||||
})
|
||||
|
||||
it("should NOT enable extended thinking for unsupported models", async () => {
|
||||
// Use a model that doesn't support reasoning
|
||||
handler = new AwsBedrockHandler({
|
||||
...mockOptions,
|
||||
apiModelId: "anthropic.claude-3-haiku-20240307-v1:0",
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Setup mock response
|
||||
mockSend.mockResolvedValue({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockStart: {
|
||||
start: { text: "Hello" },
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})(),
|
||||
})
|
||||
|
||||
// Create message
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("", messages)
|
||||
|
||||
// Consume stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(ConverseStreamCommand).toHaveBeenCalled()
|
||||
const payload = (ConverseStreamCommand as any).mock.calls[0][0]
|
||||
|
||||
// Extended thinking should NOT be enabled for unsupported models
|
||||
expect(payload.anthropic_version).toBeUndefined()
|
||||
expect(payload.additionalModelRequestFields).toBeUndefined()
|
||||
expect(payload.inferenceConfig.temperature).toBeDefined()
|
||||
expect(payload.inferenceConfig.topP).toBeDefined()
|
||||
})
|
||||
})
|
||||
|
||||
describe("Stream Processing", () => {
|
||||
it("should handle thinking delta events", async () => {
|
||||
// Enable reasoning mode
|
||||
handler = new AwsBedrockHandler({
|
||||
...mockOptions,
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Setup mock response with thinking deltas
|
||||
mockSend.mockResolvedValue({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
delta: {
|
||||
type: "thinking_delta",
|
||||
thinking: "First part of thinking...",
|
||||
},
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
delta: {
|
||||
type: "thinking_delta",
|
||||
thinking: " Second part of thinking.",
|
||||
},
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})(),
|
||||
})
|
||||
|
||||
// Create message
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("", messages)
|
||||
|
||||
// Consume stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify thinking deltas were processed
|
||||
const reasoningChunks = chunks.filter((c) => c.type === "reasoning")
|
||||
expect(reasoningChunks).toHaveLength(2)
|
||||
expect(reasoningChunks[0].text).toBe("First part of thinking...")
|
||||
expect(reasoningChunks[1].text).toBe(" Second part of thinking.")
|
||||
})
|
||||
|
||||
it("should handle signature delta events as reasoning", async () => {
|
||||
// Enable reasoning mode
|
||||
handler = new AwsBedrockHandler({
|
||||
...mockOptions,
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Setup mock response with signature deltas
|
||||
mockSend.mockResolvedValue({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield {
|
||||
contentBlockDelta: {
|
||||
delta: {
|
||||
type: "signature_delta",
|
||||
signature: "[Signature content]",
|
||||
},
|
||||
contentBlockIndex: 0,
|
||||
},
|
||||
}
|
||||
yield { messageStop: { stopReason: "end_turn" } }
|
||||
})(),
|
||||
})
|
||||
|
||||
// Create message
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("", messages)
|
||||
|
||||
// Consume stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify signature delta was processed as reasoning
|
||||
const reasoningChunk = chunks.find((c) => c.type === "reasoning")
|
||||
expect(reasoningChunk).toBeDefined()
|
||||
expect(reasoningChunk?.text).toBe("[Signature content]")
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -1,248 +0,0 @@
|
|||
import { AwsBedrockHandler } from "../bedrock"
|
||||
import { BedrockRuntimeClient } from "@aws-sdk/client-bedrock-runtime"
|
||||
|
||||
// Mock the AWS SDK
|
||||
jest.mock("@aws-sdk/client-bedrock-runtime")
|
||||
jest.mock("@aws-sdk/credential-providers")
|
||||
|
||||
describe("AwsBedrockHandler - Extended Thinking/Reasoning", () => {
|
||||
let mockClient: jest.Mocked<BedrockRuntimeClient>
|
||||
let mockSend: jest.Mock
|
||||
|
||||
beforeEach(() => {
|
||||
jest.clearAllMocks()
|
||||
mockSend = jest.fn()
|
||||
mockClient = {
|
||||
send: mockSend,
|
||||
config: { region: "us-east-1" },
|
||||
} as any
|
||||
;(BedrockRuntimeClient as jest.Mock).mockImplementation(() => mockClient)
|
||||
})
|
||||
|
||||
describe("Extended Thinking Configuration", () => {
|
||||
it("should NOT include thinking configuration by default", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
awsAccessKey: "test-key",
|
||||
awsSecretKey: "test-secret",
|
||||
awsRegion: "us-east-1",
|
||||
// enableReasoningEffort is NOT set, so reasoning should be disabled
|
||||
})
|
||||
|
||||
// Mock the stream response
|
||||
mockSend.mockResolvedValueOnce({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield { contentBlockStart: { start: { text: "Hello" } } }
|
||||
yield { messageStop: {} }
|
||||
})(),
|
||||
})
|
||||
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("System prompt", messages)
|
||||
|
||||
// Consume the stream
|
||||
for await (const _chunk of stream) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(mockSend).toHaveBeenCalledTimes(1)
|
||||
const command = mockSend.mock.calls[0][0]
|
||||
const payload = command.input
|
||||
|
||||
// Verify thinking is NOT included
|
||||
expect(payload.anthropic_version).toBeUndefined()
|
||||
expect(payload.additionalModelRequestFields).toBeUndefined()
|
||||
expect(payload.inferenceConfig.temperature).toBeDefined()
|
||||
expect(payload.inferenceConfig.topP).toBeDefined()
|
||||
})
|
||||
|
||||
it("should include thinking configuration when explicitly enabled", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
awsAccessKey: "test-key",
|
||||
awsSecretKey: "test-secret",
|
||||
awsRegion: "us-east-1",
|
||||
enableReasoningEffort: true, // Explicitly enable reasoning
|
||||
modelMaxThinkingTokens: 5000, // Set thinking tokens
|
||||
})
|
||||
|
||||
// Mock the stream response
|
||||
mockSend.mockResolvedValueOnce({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield { contentBlockStart: { contentBlock: { type: "thinking", thinking: "Let me think..." } } }
|
||||
yield { contentBlockDelta: { delta: { type: "thinking_delta", thinking: " about this." } } }
|
||||
yield { contentBlockStart: { start: { text: "Here's my answer" } } }
|
||||
yield { messageStop: {} }
|
||||
})(),
|
||||
})
|
||||
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("System prompt", messages)
|
||||
|
||||
// Consume the stream
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(mockSend).toHaveBeenCalledTimes(1)
|
||||
const command = mockSend.mock.calls[0][0]
|
||||
const payload = command.input
|
||||
|
||||
// Verify thinking IS included
|
||||
expect(payload.anthropic_version).toBe("bedrock-20250514")
|
||||
expect(payload.additionalModelRequestFields).toEqual({
|
||||
thinking: {
|
||||
type: "enabled",
|
||||
budget_tokens: 5000,
|
||||
},
|
||||
})
|
||||
// Temperature and topP should be removed when thinking is enabled
|
||||
expect(payload.inferenceConfig.temperature).toBeUndefined()
|
||||
expect(payload.inferenceConfig.topP).toBeUndefined()
|
||||
|
||||
// Verify thinking chunks were properly handled
|
||||
const reasoningChunks = chunks.filter((c) => c.type === "reasoning")
|
||||
expect(reasoningChunks).toHaveLength(2)
|
||||
expect(reasoningChunks[0].text).toBe("Let me think...")
|
||||
expect(reasoningChunks[1].text).toBe(" about this.")
|
||||
})
|
||||
|
||||
it("should NOT enable thinking for non-supported models even if requested", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-haiku-20240307-v1:0", // This model doesn't support reasoning
|
||||
awsAccessKey: "test-key",
|
||||
awsSecretKey: "test-secret",
|
||||
awsRegion: "us-east-1",
|
||||
enableReasoningEffort: true, // Try to enable reasoning
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Mock the stream response
|
||||
mockSend.mockResolvedValueOnce({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
yield { contentBlockStart: { start: { text: "Hello" } } }
|
||||
yield { messageStop: {} }
|
||||
})(),
|
||||
})
|
||||
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("System prompt", messages)
|
||||
|
||||
// Consume the stream
|
||||
for await (const _chunk of stream) {
|
||||
// Just consume
|
||||
}
|
||||
|
||||
// Verify the command was called
|
||||
expect(mockSend).toHaveBeenCalledTimes(1)
|
||||
const command = mockSend.mock.calls[0][0]
|
||||
const payload = command.input
|
||||
|
||||
// Verify thinking is NOT included because model doesn't support it
|
||||
expect(payload.anthropic_version).toBeUndefined()
|
||||
expect(payload.additionalModelRequestFields).toBeUndefined()
|
||||
})
|
||||
|
||||
it("should handle thinking stream events correctly", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-sonnet-4-20250514-v1:0",
|
||||
awsAccessKey: "test-key",
|
||||
awsSecretKey: "test-secret",
|
||||
awsRegion: "us-east-1",
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 8000,
|
||||
})
|
||||
|
||||
// Mock the stream response with various thinking events
|
||||
mockSend.mockResolvedValueOnce({
|
||||
stream: (async function* () {
|
||||
yield { messageStart: { role: "assistant" } }
|
||||
// Thinking block start
|
||||
yield {
|
||||
contentBlockStart: { contentBlock: { type: "thinking", thinking: "Analyzing the request..." } },
|
||||
}
|
||||
// Thinking deltas
|
||||
yield {
|
||||
contentBlockDelta: { delta: { type: "thinking_delta", thinking: "\nThis seems complex." } },
|
||||
}
|
||||
yield {
|
||||
contentBlockDelta: { delta: { type: "thinking_delta", thinking: "\nLet me break it down." } },
|
||||
}
|
||||
// Signature delta (part of thinking)
|
||||
yield {
|
||||
contentBlockDelta: { delta: { type: "signature_delta", signature: "\n[Signature: ABC123]" } },
|
||||
}
|
||||
// Regular text response
|
||||
yield { contentBlockStart: { start: { text: "Based on my analysis" } } }
|
||||
yield { contentBlockDelta: { delta: { text: ", here's the answer." } } }
|
||||
yield { messageStop: {} }
|
||||
})(),
|
||||
})
|
||||
|
||||
const messages = [{ role: "user" as const, content: "Complex question" }]
|
||||
const stream = handler.createMessage("System prompt", messages)
|
||||
|
||||
// Collect all chunks
|
||||
const chunks = []
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
|
||||
// Verify reasoning chunks
|
||||
const reasoningChunks = chunks.filter((c) => c.type === "reasoning")
|
||||
expect(reasoningChunks).toHaveLength(4)
|
||||
expect(reasoningChunks.map((c) => c.text).join("")).toBe(
|
||||
"Analyzing the request...\nThis seems complex.\nLet me break it down.\n[Signature: ABC123]",
|
||||
)
|
||||
|
||||
// Verify text chunks
|
||||
const textChunks = chunks.filter((c) => c.type === "text")
|
||||
expect(textChunks).toHaveLength(2)
|
||||
expect(textChunks.map((c) => c.text).join("")).toBe("Based on my analysis, here's the answer.")
|
||||
})
|
||||
})
|
||||
|
||||
describe("Error Handling for Extended Thinking", () => {
|
||||
it("should provide helpful error message for thinking-related errors", async () => {
|
||||
const handler = new AwsBedrockHandler({
|
||||
apiModelId: "anthropic.claude-3-7-sonnet-20250219-v1:0",
|
||||
awsAccessKey: "test-key",
|
||||
awsSecretKey: "test-secret",
|
||||
awsRegion: "us-east-1",
|
||||
enableReasoningEffort: true,
|
||||
modelMaxThinkingTokens: 5000,
|
||||
})
|
||||
|
||||
// Mock an error response
|
||||
const error = new Error("ValidationException: additionalModelRequestFields.thinking is not supported")
|
||||
mockSend.mockRejectedValueOnce(error)
|
||||
|
||||
const messages = [{ role: "user" as const, content: "Test message" }]
|
||||
const stream = handler.createMessage("System prompt", messages)
|
||||
|
||||
// Collect error chunks
|
||||
const chunks = []
|
||||
try {
|
||||
for await (const chunk of stream) {
|
||||
chunks.push(chunk)
|
||||
}
|
||||
} catch (e) {
|
||||
// Expected to throw
|
||||
}
|
||||
|
||||
// Should have error chunks before throwing
|
||||
expect(chunks).toHaveLength(2)
|
||||
expect(chunks[0].type).toBe("text")
|
||||
if (chunks[0].type === "text") {
|
||||
expect(chunks[0].text).toContain("Extended thinking/reasoning is not supported")
|
||||
}
|
||||
expect(chunks[1].type).toBe("usage")
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
@ -33,6 +33,9 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from ".
|
|||
import { getModelParams } from "../transform/model-params"
|
||||
import { shouldUseReasoningBudget } from "../../shared/api"
|
||||
|
||||
// Constants for Bedrock Extended Thinking
|
||||
const BEDROCK_ANTHROPIC_VERSION = "bedrock-20250514"
|
||||
|
||||
/************************************************************************************
|
||||
*
|
||||
* TYPES
|
||||
|
|
@ -46,6 +49,21 @@ interface BedrockInferenceConfig {
|
|||
topP?: number
|
||||
}
|
||||
|
||||
// Define interface for Bedrock payload
|
||||
interface BedrockPayload {
|
||||
modelId: BedrockModelId | string
|
||||
messages: Message[]
|
||||
system?: SystemContentBlock[]
|
||||
inferenceConfig: BedrockInferenceConfig
|
||||
anthropic_version?: string
|
||||
additionalModelRequestFields?: {
|
||||
thinking?: {
|
||||
type: "enabled"
|
||||
budget_tokens: number
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Define types for stream events based on AWS SDK
|
||||
export interface StreamEvent {
|
||||
messageStart?: {
|
||||
|
|
@ -131,6 +149,80 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
private client: BedrockRuntimeClient
|
||||
private arnInfo: any
|
||||
|
||||
/**
|
||||
* Determines if extended thinking should be enabled based on model support and user settings
|
||||
*/
|
||||
private shouldEnableExtendedThinking(modelInfo: ModelInfo, params: any): boolean {
|
||||
return !!(
|
||||
this.options.enableReasoningEffort &&
|
||||
shouldUseReasoningBudget({ model: modelInfo, settings: this.options }) &&
|
||||
params.reasoning &&
|
||||
params.reasoningBudget
|
||||
)
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles thinking content block events
|
||||
*/
|
||||
private *handleThinkingContentBlock(contentBlock: any): Generator<any, void, unknown> {
|
||||
if (contentBlock?.type === "thinking" && contentBlock.thinking !== undefined) {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: contentBlock.thinking || "",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles text content block events
|
||||
*/
|
||||
private *handleTextContentBlock(start: any, contentBlock: any): Generator<any, void, unknown> {
|
||||
const text = start?.text || contentBlock?.text
|
||||
if (text !== undefined) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: text || "",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles thinking delta events
|
||||
*/
|
||||
private *handleThinkingDelta(delta: any): Generator<any, void, unknown> {
|
||||
if (delta.type === "thinking_delta" && delta.thinking) {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: delta.thinking,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles signature delta events (part of thinking)
|
||||
*/
|
||||
private *handleSignatureDelta(delta: any): Generator<any, void, unknown> {
|
||||
if (delta.type === "signature_delta" && delta.signature) {
|
||||
// Signature is part of the thinking process, treat it as reasoning
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: delta.signature,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handles text delta events
|
||||
*/
|
||||
private *handleTextDelta(delta: any): Generator<any, void, unknown> {
|
||||
if (delta.text) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: delta.text,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
constructor(options: ProviderSettings) {
|
||||
super()
|
||||
this.options = options
|
||||
|
|
@ -351,7 +443,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
}
|
||||
|
||||
// Build the base payload
|
||||
const payload: any = {
|
||||
const payload: BedrockPayload = {
|
||||
modelId: modelConfig.id,
|
||||
messages: formatted.messages,
|
||||
system: formatted.system,
|
||||
|
|
@ -360,14 +452,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
|
||||
// Add extended thinking support ONLY if explicitly enabled by the user
|
||||
// Reasoning is disabled by default as per requirements
|
||||
if (
|
||||
this.options.enableReasoningEffort &&
|
||||
shouldUseReasoningBudget({ model: modelConfig.info, settings: this.options }) &&
|
||||
params.reasoning &&
|
||||
params.reasoningBudget
|
||||
) {
|
||||
if (this.shouldEnableExtendedThinking(modelConfig.info, params) && params.reasoningBudget) {
|
||||
// Add the anthropic_version field required for extended thinking
|
||||
payload.anthropic_version = "bedrock-20250514"
|
||||
payload.anthropic_version = BEDROCK_ANTHROPIC_VERSION
|
||||
|
||||
// Add additionalModelRequestFields with thinking configuration
|
||||
payload.additionalModelRequestFields = {
|
||||
|
|
@ -378,6 +465,8 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
}
|
||||
|
||||
// Remove temperature, topP, and top_k when thinking is enabled as they are incompatible
|
||||
// AWS Bedrock requires these parameters to be undefined when using extended thinking
|
||||
// as the thinking process uses its own internal temperature and sampling parameters
|
||||
delete inferenceConfig.temperature
|
||||
delete inferenceConfig.topP
|
||||
|
||||
|
|
@ -492,23 +581,17 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
|
||||
// Handle content blocks
|
||||
if (streamEvent.contentBlockStart) {
|
||||
const { contentBlock, start } = streamEvent.contentBlockStart
|
||||
|
||||
// Handle thinking content blocks
|
||||
if (streamEvent.contentBlockStart.contentBlock?.type === "thinking") {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: streamEvent.contentBlockStart.contentBlock.thinking || "",
|
||||
}
|
||||
if (contentBlock?.type === "thinking") {
|
||||
yield* this.handleThinkingContentBlock(contentBlock)
|
||||
continue
|
||||
}
|
||||
|
||||
// Handle regular text content blocks
|
||||
if (streamEvent.contentBlockStart.start?.text || streamEvent.contentBlockStart.contentBlock?.text) {
|
||||
yield {
|
||||
type: "text",
|
||||
text:
|
||||
streamEvent.contentBlockStart.start?.text ||
|
||||
streamEvent.contentBlockStart.contentBlock?.text ||
|
||||
"",
|
||||
}
|
||||
if (start?.text || contentBlock?.text) {
|
||||
yield* this.handleTextContentBlock(start, contentBlock)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
|
@ -518,32 +601,13 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
const delta = streamEvent.contentBlockDelta.delta
|
||||
|
||||
// Handle thinking deltas
|
||||
if (delta.type === "thinking_delta" && delta.thinking) {
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: delta.thinking,
|
||||
}
|
||||
continue
|
||||
}
|
||||
yield* this.handleThinkingDelta(delta)
|
||||
|
||||
// Handle signature deltas (part of thinking)
|
||||
if (delta.type === "signature_delta" && delta.signature) {
|
||||
// Signature is part of the thinking process, treat it as reasoning
|
||||
yield {
|
||||
type: "reasoning",
|
||||
text: delta.signature,
|
||||
}
|
||||
continue
|
||||
}
|
||||
yield* this.handleSignatureDelta(delta)
|
||||
|
||||
// Handle regular text deltas
|
||||
if (delta.text) {
|
||||
yield {
|
||||
type: "text",
|
||||
text: delta.text,
|
||||
}
|
||||
continue
|
||||
}
|
||||
yield* this.handleTextDelta(delta)
|
||||
}
|
||||
// Handle message stop
|
||||
if (streamEvent.messageStop) {
|
||||
|
|
@ -643,12 +707,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
// First convert messages using shared converter for proper image handling
|
||||
let convertedMessages = sharedConverter(anthropicMessages as Anthropic.Messages.MessageParam[])
|
||||
|
||||
// Handle extended thinking for tool use
|
||||
// When using tools with extended thinking, we need to preserve the thinking block
|
||||
// from previous assistant messages
|
||||
if (this.options.enableReasoningEffort && modelInfo?.supportsReasoningBudget) {
|
||||
convertedMessages = this.preserveThinkingBlocks(convertedMessages)
|
||||
}
|
||||
// No need to preserve thinking blocks - the messages are already properly formatted
|
||||
|
||||
// If prompt caching is disabled, return the converted messages directly
|
||||
if (!usePromptCache) {
|
||||
|
|
@ -931,35 +990,6 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
|
|||
return content
|
||||
}
|
||||
|
||||
/**
|
||||
* Preserves thinking blocks from previous assistant messages for tool use continuity
|
||||
*/
|
||||
private preserveThinkingBlocks(messages: Message[]): Message[] {
|
||||
// When using extended thinking with tools, we need to preserve the entire
|
||||
// thinking block from previous assistant messages to maintain reasoning continuity
|
||||
return messages.map((message, index) => {
|
||||
if (message.role === "assistant" && index > 0) {
|
||||
// Check if this assistant message follows a tool use pattern
|
||||
const prevMessage = messages[index - 1]
|
||||
if (prevMessage.role === "user" && this.hasToolUseContent(prevMessage)) {
|
||||
// This is likely a response to tool use, preserve any thinking blocks
|
||||
return message
|
||||
}
|
||||
}
|
||||
return message
|
||||
})
|
||||
}
|
||||
|
||||
/**
|
||||
* Checks if a message contains tool use content
|
||||
*/
|
||||
private hasToolUseContent(message: Message): boolean {
|
||||
if (!message.content || !Array.isArray(message.content)) {
|
||||
return false
|
||||
}
|
||||
return message.content.some((block: any) => block.toolUse || block.toolResult)
|
||||
}
|
||||
|
||||
/************************************************************************************
|
||||
*
|
||||
* AMAZON REGIONS
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue