From 6cae30d712fca9adb23ff06d24c2c5f117254644 Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Sat, 7 Jun 2025 12:23:59 -0600 Subject: [PATCH 1/7] Implement extended reasoning support in AwsBedrockHandler and add corresponding tests --- packages/types/src/providers/bedrock.ts | 3 + .../__tests__/bedrock-reasoning.test.ts | 248 ++++++++++++++++++ src/api/providers/bedrock.ts | 214 +++++++++++++-- .../components/settings/providers/Bedrock.tsx | 6 + 4 files changed, 454 insertions(+), 17 deletions(-) create mode 100644 src/api/providers/__tests__/bedrock-reasoning.test.ts diff --git a/packages/types/src/providers/bedrock.ts b/packages/types/src/providers/bedrock.ts index ce5ea28e95..a15f041252 100644 --- a/packages/types/src/providers/bedrock.ts +++ b/packages/types/src/providers/bedrock.ts @@ -73,6 +73,7 @@ export const bedrockModels = { supportsImages: true, supportsComputerUse: true, supportsPromptCache: true, + supportsReasoningBudget: true, inputPrice: 3.0, outputPrice: 15.0, cacheWritesPrice: 3.75, @@ -87,6 +88,7 @@ export const bedrockModels = { supportsImages: true, supportsComputerUse: true, supportsPromptCache: true, + supportsReasoningBudget: true, inputPrice: 15.0, outputPrice: 75.0, cacheWritesPrice: 18.75, @@ -101,6 +103,7 @@ export const bedrockModels = { supportsImages: true, supportsComputerUse: true, supportsPromptCache: true, + supportsReasoningBudget: true, inputPrice: 3.0, outputPrice: 15.0, cacheWritesPrice: 3.75, diff --git a/src/api/providers/__tests__/bedrock-reasoning.test.ts b/src/api/providers/__tests__/bedrock-reasoning.test.ts new file mode 100644 index 0000000000..b1cee35800 --- /dev/null +++ b/src/api/providers/__tests__/bedrock-reasoning.test.ts @@ -0,0 +1,248 @@ +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 + 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") + }) + }) +}) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 16ce3289aa..68236ed116 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -30,6 +30,8 @@ import { MultiPointStrategy } from "../transform/cache-strategy/multi-point-stra import { ModelInfo as CacheModelInfo } from "../transform/cache-strategy/types" import { convertToBedrockConverseMessages as sharedConverter } from "../transform/bedrock-converse-format" import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index" +import { getModelParams } from "../transform/model-params" +import { shouldUseReasoningBudget } from "../../shared/api" /************************************************************************************ * @@ -40,8 +42,8 @@ import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from ". // Define interface for Bedrock inference config interface BedrockInferenceConfig { maxTokens: number - temperature: number - topP: number + temperature?: number + topP?: number } // Define types for stream events based on AWS SDK @@ -58,10 +60,20 @@ export interface StreamEvent { text?: string } contentBlockIndex?: number + // Extended thinking support + contentBlock?: { + type?: "thinking" | "text" + thinking?: string + text?: string + } } contentBlockDelta?: { delta?: { text?: string + // Extended thinking support + type?: "thinking_delta" | "text_delta" | "signature_delta" + thinking?: string + signature?: string } contentBlockIndex?: number } @@ -256,6 +268,49 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH systemPrompt: string, messages: Anthropic.Messages.MessageParam[], metadata?: ApiHandlerCreateMessageMetadata, + ): ApiStream { + const maxRetries = 3 + let retryCount = 0 + let lastError: unknown + + while (retryCount < maxRetries) { + try { + yield* this.createMessageInternal(systemPrompt, messages, metadata) + return + } catch (error) { + lastError = error + retryCount++ + + // Check if error is retryable + const errorType = this.getErrorType(error) + const retryableErrors = ["THROTTLING", "ABORT", "GENERIC"] + + if (!retryableErrors.includes(errorType) || retryCount >= maxRetries) { + // Not retryable or max retries reached + throw error + } + + // Log retry attempt + logger.info(`Retrying Bedrock request (attempt ${retryCount}/${maxRetries})`, { + ctx: "bedrock", + errorType, + errorMessage: error instanceof Error ? error.message : String(error), + }) + + // Exponential backoff: 1s, 2s, 4s + const delay = Math.pow(2, retryCount - 1) * 1000 + await new Promise((resolve) => setTimeout(resolve, delay)) + } + } + + // If we get here, all retries failed + throw lastError + } + + private async *createMessageInternal( + systemPrompt: string, + messages: Anthropic.Messages.MessageParam[], + metadata?: ApiHandlerCreateMessageMetadata, ): ApiStream { let modelConfig = this.getModel() // Handle cross-region inference @@ -280,20 +335,59 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH conversationId, ) + // Get model parameters including reasoning configuration + const params = getModelParams({ + format: "anthropic", + modelId: modelConfig.id as string, + model: modelConfig.info, + settings: this.options, + }) + // Construct the payload const inferenceConfig: BedrockInferenceConfig = { - maxTokens: modelConfig.info.maxTokens as number, - temperature: this.options.modelTemperature as number, + maxTokens: params.maxTokens || (modelConfig.info.maxTokens as number), + temperature: params.temperature || (this.options.modelTemperature as number), topP: 0.1, } - const payload = { + // Build the base payload + const payload: any = { modelId: modelConfig.id, messages: formatted.messages, system: formatted.system, inferenceConfig, } + // 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 + ) { + // Add the anthropic_version field required for extended thinking + payload.anthropic_version = "bedrock-20250514" + + // Add additionalModelRequestFields with thinking configuration + payload.additionalModelRequestFields = { + thinking: { + type: "enabled", + budget_tokens: params.reasoningBudget, + }, + } + + // Remove temperature, topP, and top_k when thinking is enabled as they are incompatible + delete inferenceConfig.temperature + delete inferenceConfig.topP + + logger.info("Extended thinking enabled for Bedrock request", { + ctx: "bedrock", + modelId: modelConfig.id, + budgetTokens: params.reasoningBudget, + }) + } + // Create AbortController with 10 minute timeout const controller = new AbortController() let timeoutId: NodeJS.Timeout | undefined @@ -397,21 +491,59 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } // Handle content blocks - if (streamEvent.contentBlockStart?.start?.text) { - yield { - type: "text", - text: streamEvent.contentBlockStart.start.text, + if (streamEvent.contentBlockStart) { + // Handle thinking content blocks + if (streamEvent.contentBlockStart.contentBlock?.type === "thinking") { + yield { + type: "reasoning", + text: streamEvent.contentBlockStart.contentBlock.thinking || "", + } + 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 || + "", + } + continue } - continue } // Handle content deltas - if (streamEvent.contentBlockDelta?.delta?.text) { - yield { - type: "text", - text: streamEvent.contentBlockDelta.delta.text, + if (streamEvent.contentBlockDelta?.delta) { + const delta = streamEvent.contentBlockDelta.delta + + // Handle thinking deltas + if (delta.type === "thinking_delta" && delta.thinking) { + yield { + type: "reasoning", + text: delta.thinking, + } + continue + } + + // 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 + } + + // Handle regular text deltas + if (delta.text) { + yield { + type: "text", + text: delta.text, + } + continue } - continue } // Handle message stop if (streamEvent.messageStop) { @@ -509,7 +641,14 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH conversationId?: string, // Optional conversation ID to track cache points across messages ): { system: SystemContentBlock[]; messages: Message[] } { // First convert messages using shared converter for proper image handling - const convertedMessages = sharedConverter(anthropicMessages as Anthropic.Messages.MessageParam[]) + 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) + } // If prompt caching is disabled, return the converted messages directly if (!usePromptCache) { @@ -792,6 +931,35 @@ 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 @@ -905,10 +1073,22 @@ Suggestions: messageTemplate: `Invalid ARN format. ARN should follow the pattern: arn:aws:bedrock:region:account-id:resource-type/resource-name`, logLevel: "error", }, + THINKING_NOT_SUPPORTED: { + patterns: ["thinking", "reasoning", "additionalmodelrequestfields"], + messageTemplate: `Extended thinking/reasoning is not supported for this model or configuration. + +Please verify: +1. You're using a supported model (Claude 3.7 Sonnet, Claude 4 Sonnet, or Claude 4 Opus) +2. Your AWS region supports extended thinking +3. You have the necessary permissions to use this feature + +If the issue persists, try disabling "Enable Reasoning Mode" in the settings.`, + logLevel: "error", + }, // Default/generic error GENERIC: { patterns: [], // Empty patterns array means this is the default - messageTemplate: `Unknown Error`, + messageTemplate: `Bedrock is unable to process your request. Please check your configuration and try again.`, logLevel: "error", }, } diff --git a/webview-ui/src/components/settings/providers/Bedrock.tsx b/webview-ui/src/components/settings/providers/Bedrock.tsx index eb8ca94258..cdaf3fdfd8 100644 --- a/webview-ui/src/components/settings/providers/Bedrock.tsx +++ b/webview-ui/src/components/settings/providers/Bedrock.tsx @@ -8,6 +8,7 @@ import { useAppTranslation } from "@src/i18n/TranslationContext" import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@src/components/ui" import { inputEventTransform, noTransform } from "../transforms" +import { ThinkingBudget } from "../ThinkingBudget" type BedrockProps = { apiConfiguration: ProviderSettings @@ -151,6 +152,11 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedMo )} + ) } From e06ee487244025b91f301eaac78e06d424257550 Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Sat, 7 Jun 2025 21:17:46 -0600 Subject: [PATCH 2/7] 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. --- .../__tests__/bedrock-reasoning.spec.ts | 286 ++++++++++++++++++ .../__tests__/bedrock-reasoning.test.ts | 248 --------------- src/api/providers/bedrock.ts | 186 +++++++----- 3 files changed, 394 insertions(+), 326 deletions(-) create mode 100644 src/api/providers/__tests__/bedrock-reasoning.spec.ts delete mode 100644 src/api/providers/__tests__/bedrock-reasoning.test.ts diff --git a/src/api/providers/__tests__/bedrock-reasoning.spec.ts b/src/api/providers/__tests__/bedrock-reasoning.spec.ts new file mode 100644 index 0000000000..6652398459 --- /dev/null +++ b/src/api/providers/__tests__/bedrock-reasoning.spec.ts @@ -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]") + }) + }) +}) diff --git a/src/api/providers/__tests__/bedrock-reasoning.test.ts b/src/api/providers/__tests__/bedrock-reasoning.test.ts deleted file mode 100644 index b1cee35800..0000000000 --- a/src/api/providers/__tests__/bedrock-reasoning.test.ts +++ /dev/null @@ -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 - 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") - }) - }) -}) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 68236ed116..bf9aa75d9c 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -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 { + 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 { + const text = start?.text || contentBlock?.text + if (text !== undefined) { + yield { + type: "text", + text: text || "", + } + } + } + + /** + * Handles thinking delta events + */ + private *handleThinkingDelta(delta: any): Generator { + 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 { + 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 { + 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 From d2b8f870c5d7871cbfc0224f4454d32d2349d752 Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Mon, 9 Jun 2025 10:49:24 -0600 Subject: [PATCH 3/7] fix: Improve type safety for Bedrock stream generators - Add StreamChunk interface with proper type definitions - Update all generator methods to use Generator - Addresses review feedback for better type safety --- src/api/providers/bedrock.ts | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index bf9aa75d9c..6181a96de8 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -64,6 +64,12 @@ interface BedrockPayload { } } +// Define interface for stream chunks with proper typing +interface StreamChunk { + type: "reasoning" | "text" + text: string +} + // Define types for stream events based on AWS SDK export interface StreamEvent { messageStart?: { @@ -164,7 +170,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles thinking content block events */ - private *handleThinkingContentBlock(contentBlock: any): Generator { + private *handleThinkingContentBlock(contentBlock: any): Generator { if (contentBlock?.type === "thinking" && contentBlock.thinking !== undefined) { yield { type: "reasoning", @@ -176,7 +182,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles text content block events */ - private *handleTextContentBlock(start: any, contentBlock: any): Generator { + private *handleTextContentBlock(start: any, contentBlock: any): Generator { const text = start?.text || contentBlock?.text if (text !== undefined) { yield { @@ -189,7 +195,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles thinking delta events */ - private *handleThinkingDelta(delta: any): Generator { + private *handleThinkingDelta(delta: any): Generator { if (delta.type === "thinking_delta" && delta.thinking) { yield { type: "reasoning", @@ -201,7 +207,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles signature delta events (part of thinking) */ - private *handleSignatureDelta(delta: any): Generator { + private *handleSignatureDelta(delta: any): Generator { if (delta.type === "signature_delta" && delta.signature) { // Signature is part of the thinking process, treat it as reasoning yield { @@ -214,7 +220,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles text delta events */ - private *handleTextDelta(delta: any): Generator { + private *handleTextDelta(delta: any): Generator { if (delta.text) { yield { type: "text", From 03ab6946f3703243f65d7cf7b81338ca72da3b7f Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Mon, 9 Jun 2025 11:06:28 -0600 Subject: [PATCH 4/7] fix: Use ApiStreamChunk type for proper type compatibility - Replace custom StreamChunk interface with ApiStreamChunk import - Update all generator method signatures to use ApiStreamChunk - Fixes type compatibility issues with yield* delegation - Resolves CI test failures related to type mismatches --- src/api/providers/bedrock.ts | 17 +++++++---------- 1 file changed, 7 insertions(+), 10 deletions(-) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 6181a96de8..05183e5b32 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -64,11 +64,8 @@ interface BedrockPayload { } } -// Define interface for stream chunks with proper typing -interface StreamChunk { - type: "reasoning" | "text" - text: string -} +// Import the proper stream chunk types +import { ApiStreamChunk } from "../transform/stream" // Define types for stream events based on AWS SDK export interface StreamEvent { @@ -170,7 +167,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles thinking content block events */ - private *handleThinkingContentBlock(contentBlock: any): Generator { + private *handleThinkingContentBlock(contentBlock: any): Generator { if (contentBlock?.type === "thinking" && contentBlock.thinking !== undefined) { yield { type: "reasoning", @@ -182,7 +179,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles text content block events */ - private *handleTextContentBlock(start: any, contentBlock: any): Generator { + private *handleTextContentBlock(start: any, contentBlock: any): Generator { const text = start?.text || contentBlock?.text if (text !== undefined) { yield { @@ -195,7 +192,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles thinking delta events */ - private *handleThinkingDelta(delta: any): Generator { + private *handleThinkingDelta(delta: any): Generator { if (delta.type === "thinking_delta" && delta.thinking) { yield { type: "reasoning", @@ -207,7 +204,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles signature delta events (part of thinking) */ - private *handleSignatureDelta(delta: any): Generator { + private *handleSignatureDelta(delta: any): Generator { if (delta.type === "signature_delta" && delta.signature) { // Signature is part of the thinking process, treat it as reasoning yield { @@ -220,7 +217,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH /** * Handles text delta events */ - private *handleTextDelta(delta: any): Generator { + private *handleTextDelta(delta: any): Generator { if (delta.text) { yield { type: "text", From 765335b324ec16a61daa138f86b3251ea58e7187 Mon Sep 17 00:00:00 2001 From: Daniel <57051444+daniel-lxs@users.noreply.github.com> Date: Mon, 9 Jun 2025 12:07:37 -0500 Subject: [PATCH 5/7] cleanup --- src/api/providers/bedrock.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 05183e5b32..2cf6b31acd 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -710,7 +710,6 @@ 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[]) - // No need to preserve thinking blocks - the messages are already properly formatted // If prompt caching is disabled, return the converted messages directly if (!usePromptCache) { From d343bff0b7dc4573dcbe0d0017b85a4b769619d8 Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Mon, 9 Jun 2025 12:10:11 -0600 Subject: [PATCH 6/7] refactor(bedrock): consolidate stream processing for better maintainability - Replace multiple handler methods with unified processStreamEvent() - Simplify main stream processing loop from ~30 lines to 1 line - Maintain all existing functionality and test coverage - Address PR reviewer feedback for improved code organization Addresses: Stream processing consolidation suggestion from PR review --- src/api/providers/bedrock.ts | 103 +++++++---------------------------- 1 file changed, 21 insertions(+), 82 deletions(-) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 2cf6b31acd..a111ba6367 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -165,63 +165,30 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } /** - * Handles thinking content block events + * Unified stream event processor for better maintainability */ - private *handleThinkingContentBlock(contentBlock: any): Generator { - if (contentBlock?.type === "thinking" && contentBlock.thinking !== undefined) { - yield { - type: "reasoning", - text: contentBlock.thinking || "", + private *processStreamEvent(streamEvent: StreamEvent): Generator { + // Handle content blocks + if (streamEvent.contentBlockStart) { + const { contentBlock, start } = streamEvent.contentBlockStart + + if (contentBlock?.type === "thinking" && contentBlock.thinking !== undefined) { + yield { type: "reasoning", text: contentBlock.thinking || "" } + } else if (start?.text || contentBlock?.text) { + yield { type: "text", text: start?.text || contentBlock?.text || "" } } } - } - - /** - * Handles text content block events - */ - private *handleTextContentBlock(start: any, contentBlock: any): Generator { - const text = start?.text || contentBlock?.text - if (text !== undefined) { - yield { - type: "text", - text: text || "", - } - } - } - - /** - * Handles thinking delta events - */ - private *handleThinkingDelta(delta: any): Generator { - 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 { - if (delta.type === "signature_delta" && delta.signature) { - // Signature is part of the thinking process, treat it as reasoning - yield { - type: "reasoning", - text: delta.signature, - } - } - } + // Handle content deltas + if (streamEvent.contentBlockDelta?.delta) { + const { delta } = streamEvent.contentBlockDelta - /** - * Handles text delta events - */ - private *handleTextDelta(delta: any): Generator { - if (delta.text) { - yield { - type: "text", - text: delta.text, + if (delta.type === "thinking_delta" && delta.thinking) { + yield { type: "reasoning", text: delta.thinking } + } else if (delta.type === "signature_delta" && delta.signature) { + yield { type: "reasoning", text: delta.signature } + } else if (delta.text) { + yield { type: "text", text: delta.text } } } } @@ -582,36 +549,8 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH continue } - // Handle content blocks - if (streamEvent.contentBlockStart) { - const { contentBlock, start } = streamEvent.contentBlockStart - - // Handle thinking content blocks - if (contentBlock?.type === "thinking") { - yield* this.handleThinkingContentBlock(contentBlock) - continue - } - - // Handle regular text content blocks - if (start?.text || contentBlock?.text) { - yield* this.handleTextContentBlock(start, contentBlock) - continue - } - } - - // Handle content deltas - if (streamEvent.contentBlockDelta?.delta) { - const delta = streamEvent.contentBlockDelta.delta - - // Handle thinking deltas - yield* this.handleThinkingDelta(delta) - - // Handle signature deltas (part of thinking) - yield* this.handleSignatureDelta(delta) - - // Handle regular text deltas - yield* this.handleTextDelta(delta) - } + // Use unified stream event processor + yield* this.processStreamEvent(streamEvent) // Handle message stop if (streamEvent.messageStop) { continue From 13e272c8baa09c5d1f8f94920a4cf7c1eaef3258 Mon Sep 17 00:00:00 2001 From: hannesrudolph Date: Mon, 9 Jun 2025 13:58:45 -0600 Subject: [PATCH 7/7] refactor(bedrock): remove unused ThinkingBudget component from Bedrock --- webview-ui/src/components/settings/providers/Bedrock.tsx | 6 ------ 1 file changed, 6 deletions(-) diff --git a/webview-ui/src/components/settings/providers/Bedrock.tsx b/webview-ui/src/components/settings/providers/Bedrock.tsx index cdaf3fdfd8..eb8ca94258 100644 --- a/webview-ui/src/components/settings/providers/Bedrock.tsx +++ b/webview-ui/src/components/settings/providers/Bedrock.tsx @@ -8,7 +8,6 @@ import { useAppTranslation } from "@src/i18n/TranslationContext" import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@src/components/ui" import { inputEventTransform, noTransform } from "../transforms" -import { ThinkingBudget } from "../ThinkingBudget" type BedrockProps = { apiConfiguration: ProviderSettings @@ -152,11 +151,6 @@ export const Bedrock = ({ apiConfiguration, setApiConfigurationField, selectedMo )} - ) }