From b033082523e587e50abcb2d2978811f5f683088e Mon Sep 17 00:00:00 2001 From: Yukky <67864326+Yoshino-Yukitaro@users.noreply.github.com> Date: Fri, 21 Mar 2025 05:03:51 +0900 Subject: [PATCH] Reflect Cross-region inference option in `ap-xx` region (#1842) * Fix: Enable cross-region inference for 'ap-xx' region in AwsBedrockHandler.completePrompt * Fix: Enable cross-region inference for 'ap-xx' region in AwsBedrockHandler.createMessage * Create itchy-waves-move.md --------- Co-authored-by: Matt Rubens --- .changeset/itchy-waves-move.md | 5 + src/api/providers/__tests__/bedrock.test.ts | 317 +++++++++++++++++++- src/api/providers/bedrock.ts | 6 + 3 files changed, 327 insertions(+), 1 deletion(-) create mode 100644 .changeset/itchy-waves-move.md diff --git a/.changeset/itchy-waves-move.md b/.changeset/itchy-waves-move.md new file mode 100644 index 0000000000..4daea1527e --- /dev/null +++ b/.changeset/itchy-waves-move.md @@ -0,0 +1,5 @@ +--- +"roo-cline": patch +--- + +Reflect Cross-region inference option in `ap-xx` region diff --git a/src/api/providers/__tests__/bedrock.test.ts b/src/api/providers/__tests__/bedrock.test.ts index 0094c3f12b..e9ba74ac6b 100644 --- a/src/api/providers/__tests__/bedrock.test.ts +++ b/src/api/providers/__tests__/bedrock.test.ts @@ -165,6 +165,222 @@ describe("AwsBedrockHandler", () => { ) }) + it("should handle cross-region inference for us-xx region", async () => { + const handlerWithProfile = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "us-east-1", + awsUseCrossRegionInference: true, + }) + + // Mock AWS SDK invoke + const mockStream = { + [Symbol.asyncIterator]: async function* () { + yield { + metadata: { + usage: { + inputTokens: 10, + outputTokens: 5, + }, + }, + } + }, + } + + const mockInvoke = jest.fn().mockResolvedValue({ + stream: mockStream, + }) + + handlerWithProfile["client"] = { + send: mockInvoke, + } as unknown as BedrockRuntimeClient + + const stream = handlerWithProfile.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + expect(chunks[0]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + }) + + expect(mockInvoke).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "us.anthropic.claude-3-5-sonnet-20241022-v2:0", + }), + }), + ) + }) + + it("should handle cross-region inference for eu-xx region", async () => { + const handlerWithProfile = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20240620-v1:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "eu-west-1", + awsUseCrossRegionInference: true, + }) + + // Mock AWS SDK invoke + const mockStream = { + [Symbol.asyncIterator]: async function* () { + yield { + metadata: { + usage: { + inputTokens: 10, + outputTokens: 5, + }, + }, + } + }, + } + + const mockInvoke = jest.fn().mockResolvedValue({ + stream: mockStream, + }) + + handlerWithProfile["client"] = { + send: mockInvoke, + } as unknown as BedrockRuntimeClient + + const stream = handlerWithProfile.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + expect(chunks[0]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + }) + + expect(mockInvoke).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "eu.anthropic.claude-3-5-sonnet-20240620-v1:0", + }), + }), + ) + }) + + it("should handle cross-region inference for ap-xx region", async () => { + const handlerWithProfile = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "ap-northeast-1", + awsUseCrossRegionInference: true, + }) + + // Mock AWS SDK invoke + const mockStream = { + [Symbol.asyncIterator]: async function* () { + yield { + metadata: { + usage: { + inputTokens: 10, + outputTokens: 5, + }, + }, + } + }, + } + + const mockInvoke = jest.fn().mockResolvedValue({ + stream: mockStream, + }) + + handlerWithProfile["client"] = { + send: mockInvoke, + } as unknown as BedrockRuntimeClient + + const stream = handlerWithProfile.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + expect(chunks[0]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + }) + + expect(mockInvoke).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "apac.anthropic.claude-3-5-sonnet-20241022-v2:0", + }), + }), + ) + }) + + it("should handle cross-region inference for other region", async () => { + const handlerWithProfile = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "ca-central-1", + awsUseCrossRegionInference: true, + }) + + // Mock AWS SDK invoke + const mockStream = { + [Symbol.asyncIterator]: async function* () { + yield { + metadata: { + usage: { + inputTokens: 10, + outputTokens: 5, + }, + }, + } + }, + } + + const mockInvoke = jest.fn().mockResolvedValue({ + stream: mockStream, + }) + + handlerWithProfile["client"] = { + send: mockInvoke, + } as unknown as BedrockRuntimeClient + + const stream = handlerWithProfile.createMessage(systemPrompt, mockMessages) + const chunks = [] + + for await (const chunk of stream) { + chunks.push(chunk) + } + + expect(chunks.length).toBeGreaterThan(0) + expect(chunks[0]).toEqual({ + type: "usage", + inputTokens: 10, + outputTokens: 5, + }) + + expect(mockInvoke).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "anthropic.claude-3-sonnet-20240229-v1:0", + }), + }), + ) + }) + it("should handle API errors", async () => { // Mock AWS SDK invoke with error const mockInvoke = jest.fn().mockRejectedValue(new Error("AWS Bedrock error")) @@ -260,7 +476,7 @@ describe("AwsBedrockHandler", () => { expect(result).toBe("") }) - it("should handle cross-region inference", async () => { + it("should handle cross-region inference for us-xx region", async () => { handler = new AwsBedrockHandler({ apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", awsAccessKey: "test-access-key", @@ -292,6 +508,105 @@ describe("AwsBedrockHandler", () => { }), ) }) + + it("should handle cross-region inference for eu-xx region", async () => { + handler = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20240620-v1:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "eu-west-1", + awsUseCrossRegionInference: true, + }) + + const mockResponse = { + output: new TextEncoder().encode( + JSON.stringify({ + content: "Test response", + }), + ), + } + + const mockSend = jest.fn().mockResolvedValue(mockResponse) + handler["client"] = { + send: mockSend, + } as unknown as BedrockRuntimeClient + + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Test response") + expect(mockSend).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "eu.anthropic.claude-3-5-sonnet-20240620-v1:0", + }), + }), + ) + }) + + it("should handle cross-region inference for ap-xx region", async () => { + handler = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "ap-northeast-1", + awsUseCrossRegionInference: true, + }) + + const mockResponse = { + output: new TextEncoder().encode( + JSON.stringify({ + content: "Test response", + }), + ), + } + + const mockSend = jest.fn().mockResolvedValue(mockResponse) + handler["client"] = { + send: mockSend, + } as unknown as BedrockRuntimeClient + + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Test response") + expect(mockSend).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "apac.anthropic.claude-3-5-sonnet-20241022-v2:0", + }), + }), + ) + }) + + it("should handle cross-region inference for other regions", async () => { + handler = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-sonnet-20240229-v1:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "ca-central-1", + awsUseCrossRegionInference: true, + }) + + const mockResponse = { + output: new TextEncoder().encode( + JSON.stringify({ + content: "Test response", + }), + ), + } + + const mockSend = jest.fn().mockResolvedValue(mockResponse) + handler["client"] = { + send: mockSend, + } as unknown as BedrockRuntimeClient + + const result = await handler.completePrompt("Test prompt") + expect(result).toBe("Test response") + expect(mockSend).toHaveBeenCalledWith( + expect.objectContaining({ + input: expect.objectContaining({ + modelId: "anthropic.claude-3-sonnet-20240229-v1:0", + }), + }), + ) + }) }) describe("getModel", () => { diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 1637fe29f3..4696c1dc91 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -211,6 +211,9 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH case "eu-": modelId = `eu.${modelConfig.id}` break + case "ap-": + modelId = `apac.${modelConfig.id}` + break default: modelId = modelConfig.id break @@ -610,6 +613,9 @@ Please check: case "eu-": modelId = `eu.${modelConfig.id}` break + case "ap-": + modelId = `apac.${modelConfig.id}` + break default: modelId = modelConfig.id break