Fix: AWS Bedrock 1M context - Move anthropic_beta to additionalModelRequestFields (#7056)

This commit is contained in:
Daniel 2025-08-13 14:26:50 -05:00 committed by GitHub
parent 962df868bd
commit 7ed833cb5d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 30 additions and 19 deletions

View file

@ -635,9 +635,11 @@ describe("AwsBedrockHandler", () => {
expect(mockConverseStreamCommand).toHaveBeenCalled()
const commandArg = mockConverseStreamCommand.mock.calls[0][0] as any
// Should include anthropic_beta parameter but NOT anthropic_version (only for thinking)
expect(commandArg.anthropic_beta).toEqual(["context-1m-2025-08-07"])
expect(commandArg.anthropic_version).toBeUndefined()
// Should include anthropic_beta in additionalModelRequestFields
expect(commandArg.additionalModelRequestFields).toBeDefined()
expect(commandArg.additionalModelRequestFields.anthropic_beta).toEqual(["context-1m-2025-08-07"])
// Should not include anthropic_version since thinking is not enabled
expect(commandArg.additionalModelRequestFields.anthropic_version).toBeUndefined()
})
it("should not include anthropic_beta parameter when 1M context is disabled", async () => {
@ -663,8 +665,8 @@ describe("AwsBedrockHandler", () => {
expect(mockConverseStreamCommand).toHaveBeenCalled()
const commandArg = mockConverseStreamCommand.mock.calls[0][0] as any
// Should not include anthropic_beta parameter
expect(commandArg.anthropic_beta).toBeUndefined()
// Should not include anthropic_beta in additionalModelRequestFields
expect(commandArg.additionalModelRequestFields).toBeUndefined()
})
it("should not include anthropic_beta parameter for non-Claude Sonnet 4 models", async () => {
@ -690,8 +692,8 @@ describe("AwsBedrockHandler", () => {
expect(mockConverseStreamCommand).toHaveBeenCalled()
const commandArg = mockConverseStreamCommand.mock.calls[0][0] as any
// Should not include anthropic_beta parameter for non-Sonnet 4 models
expect(commandArg.anthropic_beta).toBeUndefined()
// Should not include anthropic_beta for non-Sonnet 4 models
expect(commandArg.additionalModelRequestFields).toBeUndefined()
})
it("should enable 1M context window with cross-region inference for Claude Sonnet 4", () => {
@ -738,9 +740,11 @@ describe("AwsBedrockHandler", () => {
mockConverseStreamCommand.mock.calls.length - 1
][0] as any
// Should include anthropic_beta parameter but NOT anthropic_version (only for thinking)
expect(commandArg.anthropic_beta).toEqual(["context-1m-2025-08-07"])
expect(commandArg.anthropic_version).toBeUndefined()
// Should include anthropic_beta in additionalModelRequestFields
expect(commandArg.additionalModelRequestFields).toBeDefined()
expect(commandArg.additionalModelRequestFields.anthropic_beta).toEqual(["context-1m-2025-08-07"])
// Should not include anthropic_version since thinking is not enabled
expect(commandArg.additionalModelRequestFields.anthropic_version).toBeUndefined()
// Model ID should have cross-region prefix
expect(commandArg.modelId).toBe(`us.${BEDROCK_CLAUDE_SONNET_4_MODEL_ID}`)
})

View file

@ -47,12 +47,14 @@ interface BedrockInferenceConfig {
topP?: number
}
// Define interface for Bedrock thinking configuration
interface BedrockThinkingConfig {
thinking: {
// Define interface for Bedrock additional model request fields
// This includes thinking configuration, 1M context beta, and other model-specific parameters
interface BedrockAdditionalModelFields {
thinking?: {
type: "enabled"
budget_tokens: number
}
anthropic_beta?: string[]
[key: string]: any // Add index signature to be compatible with DocumentType
}
@ -63,8 +65,7 @@ interface BedrockPayload {
system?: SystemContentBlock[]
inferenceConfig: BedrockInferenceConfig
anthropic_version?: string
anthropic_beta?: string[]
additionalModelRequestFields?: BedrockThinkingConfig
additionalModelRequestFields?: BedrockAdditionalModelFields
}
// Define specific types for content block events to avoid 'as any' usage
@ -341,7 +342,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
conversationId,
)
let additionalModelRequestFields: BedrockThinkingConfig | undefined
let additionalModelRequestFields: BedrockAdditionalModelFields | undefined
let thinkingEnabled = false
// Determine if thinking should be enabled
@ -382,16 +383,22 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH
const baseModelId = this.parseBaseModelId(modelConfig.id)
const is1MContextEnabled = baseModelId === BEDROCK_CLAUDE_SONNET_4_MODEL_ID && this.options.awsBedrock1MContext
// Add anthropic_beta for 1M context to additionalModelRequestFields
if (is1MContextEnabled) {
if (!additionalModelRequestFields) {
additionalModelRequestFields = {} as BedrockAdditionalModelFields
}
additionalModelRequestFields.anthropic_beta = ["context-1m-2025-08-07"]
}
const payload: BedrockPayload = {
modelId: modelConfig.id,
messages: formatted.messages,
system: formatted.system,
inferenceConfig,
...(additionalModelRequestFields && { additionalModelRequestFields }),
// Add anthropic_version when using thinking features
// Add anthropic_version at top level when using thinking features
...(thinkingEnabled && { anthropic_version: "bedrock-2023-05-31" }),
// Add anthropic_beta when 1M context is enabled
...(is1MContextEnabled && { anthropic_beta: ["context-1m-2025-08-07"] }),
}
// Create AbortController with 10 minute timeout