From 171037a93878dece89a86665d7540a67dfe283d6 Mon Sep 17 00:00:00 2001 From: Smartsheet-JB-Brown Date: Mon, 10 Mar 2025 11:10:49 -0700 Subject: [PATCH 1/2] Add enhanced error handling and logging for AWS Bedrock custom ARNs --- .../__tests__/bedrock-custom-arn.test.ts | 75 +++ src/api/providers/__tests__/bedrock.test.ts | 29 ++ src/api/providers/bedrock.ts | 491 ++++++++++++++++-- src/shared/api.ts | 2 + src/shared/globalState.ts | 1 + test-custom-arn.js | 196 +++++++ .../src/components/settings/ApiOptions.tsx | 89 +++- webview-ui/src/utils/validate.ts | 38 ++ 8 files changed, 887 insertions(+), 34 deletions(-) create mode 100644 src/api/providers/__tests__/bedrock-custom-arn.test.ts create mode 100644 test-custom-arn.js diff --git a/src/api/providers/__tests__/bedrock-custom-arn.test.ts b/src/api/providers/__tests__/bedrock-custom-arn.test.ts new file mode 100644 index 0000000000..f7dc2870fa --- /dev/null +++ b/src/api/providers/__tests__/bedrock-custom-arn.test.ts @@ -0,0 +1,75 @@ +import { AwsBedrockHandler } from "../bedrock" +import { ApiHandlerOptions } from "../../../shared/api" + +// Mock the AWS SDK +jest.mock("@aws-sdk/client-bedrock-runtime", () => { + const mockSend = jest.fn().mockImplementation(() => { + return Promise.resolve({ + output: new TextEncoder().encode(JSON.stringify({ content: "Test response" })), + }) + }) + + return { + BedrockRuntimeClient: jest.fn().mockImplementation(() => ({ + send: mockSend, + config: { + region: "us-east-1", + }, + })), + ConverseCommand: jest.fn(), + ConverseStreamCommand: jest.fn(), + } +}) + +describe("AwsBedrockHandler with custom ARN", () => { + const mockOptions: ApiHandlerOptions = { + apiModelId: "custom-arn", + awsCustomArn: "arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-sonnet-20240229-v1:0", + awsRegion: "us-east-1", + } + + it("should use the custom ARN as the model ID", async () => { + const handler = new AwsBedrockHandler(mockOptions) + const model = handler.getModel() + + expect(model.id).toBe(mockOptions.awsCustomArn) + expect(model.info).toHaveProperty("maxTokens") + expect(model.info).toHaveProperty("contextWindow") + expect(model.info).toHaveProperty("supportsPromptCache") + }) + + it("should extract region from ARN and use it for client configuration", () => { + // Test with matching region + const handler1 = new AwsBedrockHandler(mockOptions) + expect((handler1 as any).client.config.region).toBe("us-east-1") + + // Test with mismatched region + const mismatchOptions = { + ...mockOptions, + awsRegion: "us-west-2", + } + const handler2 = new AwsBedrockHandler(mismatchOptions) + // Should use the ARN region, not the provided region + expect((handler2 as any).client.config.region).toBe("us-east-1") + }) + + it("should validate ARN format", async () => { + // Invalid ARN format + const invalidOptions = { + ...mockOptions, + awsCustomArn: "invalid-arn-format", + } + + const handler = new AwsBedrockHandler(invalidOptions) + + // completePrompt should throw an error for invalid ARN + await expect(handler.completePrompt("test")).rejects.toThrow("Invalid ARN format") + }) + + it("should complete a prompt successfully with valid ARN", async () => { + const handler = new AwsBedrockHandler(mockOptions) + const response = await handler.completePrompt("test prompt") + + expect(response).toBe("Test response") + }) +}) diff --git a/src/api/providers/__tests__/bedrock.test.ts b/src/api/providers/__tests__/bedrock.test.ts index f1b2c5527f..f778621e9c 100644 --- a/src/api/providers/__tests__/bedrock.test.ts +++ b/src/api/providers/__tests__/bedrock.test.ts @@ -315,5 +315,34 @@ describe("AwsBedrockHandler", () => { expect(modelInfo.info.maxTokens).toBe(5000) expect(modelInfo.info.contextWindow).toBe(128_000) }) + + it("should use custom ARN when provided", () => { + const customArnHandler = new AwsBedrockHandler({ + apiModelId: "anthropic.claude-3-5-sonnet-20241022-v2:0", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "us-east-1", + awsCustomArn: "arn:aws:bedrock:us-east-1::foundation-model/custom-model", + }) + const modelInfo = customArnHandler.getModel() + expect(modelInfo.id).toBe("arn:aws:bedrock:us-east-1::foundation-model/custom-model") + expect(modelInfo.info.maxTokens).toBe(5000) + expect(modelInfo.info.contextWindow).toBe(128_000) + expect(modelInfo.info.supportsPromptCache).toBe(false) + }) + + it("should use default model when custom-arn is selected but no ARN is provided", () => { + const customArnHandler = new AwsBedrockHandler({ + apiModelId: "custom-arn", + awsAccessKey: "test-access-key", + awsSecretKey: "test-secret-key", + awsRegion: "us-east-1", + // No awsCustomArn provided + }) + const modelInfo = customArnHandler.getModel() + // Should fall back to default model + expect(modelInfo.id).not.toBe("custom-arn") + expect(modelInfo.info).toBeDefined() + }) }) }) diff --git a/src/api/providers/bedrock.ts b/src/api/providers/bedrock.ts index 2deb019dc3..76d9364960 100644 --- a/src/api/providers/bedrock.ts +++ b/src/api/providers/bedrock.ts @@ -11,6 +11,47 @@ import { ApiHandlerOptions, BedrockModelId, ModelInfo, bedrockDefaultModelId, be import { ApiStream } from "../transform/stream" import { convertToBedrockConverseMessages } from "../transform/bedrock-converse-format" import { BaseProvider } from "./base-provider" +import { logger } from "../../utils/logging" + +/** + * Validates an AWS Bedrock ARN format and optionally checks if the region in the ARN matches the provided region + * @param arn The ARN string to validate + * @param region Optional region to check against the ARN's region + * @returns An object with validation results: { isValid, arnRegion, errorMessage } + */ +function validateBedrockArn(arn: string, region?: string) { + // Validate ARN format + const arnRegex = /^arn:aws:bedrock:([^:]+):(\d+):(foundation-model|provisioned-model|default-prompt-router)\/(.+)$/ + const match = arn.match(arnRegex) + + if (!match) { + return { + isValid: false, + arnRegion: undefined, + errorMessage: + "Invalid ARN format. ARN should follow the pattern: arn:aws:bedrock:region:account-id:resource-type/resource-name", + } + } + + // Extract region from ARN + const arnRegion = match[1] + + // Check if region in ARN matches provided region (if specified) + if (region && arnRegion !== region) { + return { + isValid: true, + arnRegion, + errorMessage: `Warning: The region in your ARN (${arnRegion}) does not match your selected region (${region}). This may cause access issues. The provider will use the region from the ARN.`, + } + } + + // ARN is valid and region matches (or no region was provided to check against) + return { + isValid: true, + arnRegion, + errorMessage: undefined, + } +} const BEDROCK_DEFAULT_TEMPERATURE = 0.3 @@ -55,8 +96,31 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH super() this.options = options + // Extract region from custom ARN if provided + let region = this.options.awsRegion || "us-east-1" + + // If using custom ARN, extract region from the ARN + if (this.options.awsCustomArn) { + const validation = validateBedrockArn(this.options.awsCustomArn, region) + + if (validation.isValid && validation.arnRegion) { + // If there's a region mismatch warning, log it and use the ARN region + if (validation.errorMessage) { + logger.info( + `Region mismatch: Selected region is ${region}, but ARN region is ${validation.arnRegion}. Using ARN region.`, + { + ctx: "bedrock", + selectedRegion: region, + arnRegion: validation.arnRegion, + }, + ) + region = validation.arnRegion + } + } + } + const clientConfig: BedrockRuntimeClientConfig = { - region: this.options.awsRegion || "us-east-1", + region: region, } if (this.options.awsUseProfile && this.options.awsProfile) { @@ -81,7 +145,41 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH // Handle cross-region inference let modelId: string - if (this.options.awsUseCrossRegionInference) { + + // For custom ARNs, use the ARN directly without modification + if (this.options.awsCustomArn) { + modelId = modelConfig.id + + // Validate ARN format and check region match + const clientRegion = this.client.config.region as string + const validation = validateBedrockArn(modelId, clientRegion) + + if (!validation.isValid) { + logger.error("Invalid ARN format", { + ctx: "bedrock", + modelId, + errorMessage: validation.errorMessage, + }) + yield { + type: "text", + text: `Error: ${validation.errorMessage}`, + } + yield { type: "usage", inputTokens: 0, outputTokens: 0 } + throw new Error("Invalid ARN format") + } + + // Extract region from ARN + const arnRegion = validation.arnRegion! + + // Log warning if there's a region mismatch + if (validation.errorMessage) { + logger.warn(validation.errorMessage, { + ctx: "bedrock", + arnRegion, + clientRegion, + }) + } + } else if (this.options.awsUseCrossRegionInference) { let regionPrefix = (this.options.awsRegion || "").slice(0, 3) switch (regionPrefix) { case "us-": @@ -107,7 +205,7 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH messages: formattedMessages, system: [{ text: systemPrompt }], inferenceConfig: { - maxTokens: modelConfig.info.maxTokens || 5000, + maxTokens: modelConfig.info.maxTokens || 4096, temperature: this.options.modelTemperature ?? BEDROCK_DEFAULT_TEMPERATURE, topP: 0.1, ...(this.options.awsUsePromptCache @@ -121,6 +219,16 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } try { + // Log the payload for debugging custom ARN issues + if (this.options.awsCustomArn) { + logger.debug("Using custom ARN for Bedrock request", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + clientRegion: this.client.config.region, + payload: JSON.stringify(payload, null, 2), + }) + } + const command = new ConverseStreamCommand(payload) const response = await this.client.send(command) @@ -134,7 +242,11 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH try { streamEvent = typeof chunk === "string" ? JSON.parse(chunk) : (chunk as unknown as StreamEvent) } catch (e) { - console.error("Failed to parse stream event:", e) + logger.error("Failed to parse stream event", { + ctx: "bedrock", + error: e instanceof Error ? e : String(e), + chunk: typeof chunk === "string" ? chunk : "binary data", + }) continue } @@ -177,39 +289,257 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH } } } catch (error: unknown) { - console.error("Bedrock Runtime API Error:", error) - // Only access stack if error is an Error object + logger.error("Bedrock Runtime API Error", { + ctx: "bedrock", + error: error instanceof Error ? error : String(error), + }) + + // Enhanced error handling for custom ARN issues + if (this.options.awsCustomArn) { + logger.error("Error occurred with custom ARN", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + }) + + // Check for common ARN-related errors + if (error instanceof Error) { + const errorMessage = error.message.toLowerCase() + + // Access denied errors + if ( + errorMessage.includes("access") && + (errorMessage.includes("model") || errorMessage.includes("denied")) + ) { + logger.error("Permissions issue with custom ARN", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + errorType: "access_denied", + clientRegion: this.client.config.region, + }) + yield { + type: "text", + text: `Error: You don't have access to the model with the specified ARN. Please verify: + +1. The ARN is correct and points to a valid model +2. Your AWS credentials have permission to access this model (check IAM policies) +3. The region in the ARN (${this.client.config.region}) matches the region where the model is deployed +4. If using a provisioned model, ensure it's active and not in a failed state +5. If using a custom model, ensure your account has been granted access to it`, + } + } + // Model not found errors + else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) { + logger.error("Invalid ARN or non-existent model", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + errorType: "not_found", + }) + yield { + type: "text", + text: `Error: The specified ARN does not exist or is invalid. Please check: + +1. The ARN format is correct (arn:aws:bedrock:region:account-id:resource-type/resource-name) +2. The model exists in the specified region +3. The account ID in the ARN is correct +4. The resource type is one of: foundation-model, provisioned-model, or default-prompt-router`, + } + } + // Throttling errors + else if ( + errorMessage.includes("throttl") || + errorMessage.includes("rate") || + errorMessage.includes("limit") + ) { + logger.error("Throttling or rate limit issue with Bedrock", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + errorType: "throttling", + }) + yield { + type: "text", + text: `Error: Request was throttled or rate limited. Please try: + +1. Reducing the frequency of requests +2. If using a provisioned model, check its throughput settings +3. Contact AWS support to request a quota increase if needed`, + } + } + // Other errors + else { + logger.error("Unspecified error with custom ARN", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + errorStack: error.stack, + errorMessage: error.message, + }) + yield { + type: "text", + text: `Error with custom ARN: ${error.message} + +Please check: +1. Your AWS credentials are valid and have the necessary permissions +2. The ARN format is correct +3. The region in the ARN matches the region where you're making the request`, + } + } + } else { + yield { + type: "text", + text: `Unknown error occurred with custom ARN. Please check your AWS credentials and ARN format.`, + } + } + } else { + // Standard error handling for non-ARN cases + if (error instanceof Error) { + logger.error("Standard Bedrock error", { + ctx: "bedrock", + errorStack: error.stack, + errorMessage: error.message, + }) + yield { + type: "text", + text: `Error: ${error.message}`, + } + } else { + logger.error("Unknown Bedrock error", { + ctx: "bedrock", + error: String(error), + }) + yield { + type: "text", + text: "An unknown error occurred", + } + } + } + + // Always yield usage info + yield { + type: "usage", + inputTokens: 0, + outputTokens: 0, + } + + // Re-throw the error if (error instanceof Error) { - console.error("Error stack:", error.stack) - yield { - type: "text", - text: `Error: ${error.message}`, - } - yield { - type: "usage", - inputTokens: 0, - outputTokens: 0, - } throw error } else { - const unknownError = new Error("An unknown error occurred") - yield { - type: "text", - text: unknownError.message, - } - yield { - type: "usage", - inputTokens: 0, - outputTokens: 0, - } - throw unknownError + throw new Error("An unknown error occurred") } } } override getModel(): { id: BedrockModelId | string; info: ModelInfo } { + // If custom ARN is provided, use it + if (this.options.awsCustomArn) { + // Custom ARNs should not be modified with region prefixes + // as they already contain the full resource path + + // Check if the ARN contains information about the model type + // This helps set appropriate token limits for models behind prompt routers + const arnLower = this.options.awsCustomArn.toLowerCase() + + // Determine model info based on ARN content + let modelInfo: ModelInfo + + if (arnLower.includes("claude-3-7-sonnet") || arnLower.includes("claude-3.7-sonnet")) { + // Claude 3.7 Sonnet has 8192 tokens in Bedrock + modelInfo = { + maxTokens: 8192, + contextWindow: 200_000, + supportsPromptCache: false, + supportsImages: true, + supportsComputerUse: true, + } + } else if (arnLower.includes("claude-3-5-sonnet") || arnLower.includes("claude-3.5-sonnet")) { + // Claude 3.5 Sonnet has 8192 tokens in Bedrock + modelInfo = { + maxTokens: 8192, + contextWindow: 200_000, + supportsPromptCache: false, + supportsImages: true, + supportsComputerUse: true, + } + } else if (arnLower.includes("claude-3-opus") || arnLower.includes("claude-3.0-opus")) { + // Claude 3 Opus has 4096 tokens in Bedrock + modelInfo = { + maxTokens: 4096, + contextWindow: 200_000, + supportsPromptCache: false, + supportsImages: true, + } + } else if (arnLower.includes("claude-3-haiku") || arnLower.includes("claude-3.0-haiku")) { + // Claude 3 Haiku has 4096 tokens in Bedrock + modelInfo = { + maxTokens: 4096, + contextWindow: 200_000, + supportsPromptCache: false, + supportsImages: true, + } + } else if (arnLower.includes("claude-3-5-haiku") || arnLower.includes("claude-3.5-haiku")) { + // Claude 3.5 Haiku has 8192 tokens in Bedrock + modelInfo = { + maxTokens: 8192, + contextWindow: 200_000, + supportsPromptCache: false, + supportsImages: false, + } + } else if (arnLower.includes("claude")) { + // Generic Claude model with conservative token limit + modelInfo = { + maxTokens: 4096, + contextWindow: 128_000, + supportsPromptCache: false, + supportsImages: true, + } + } else if (arnLower.includes("llama3") || arnLower.includes("llama-3")) { + // Llama 3 models typically have 8192 tokens in Bedrock + modelInfo = { + maxTokens: 8192, + contextWindow: 128_000, + supportsPromptCache: false, + supportsImages: arnLower.includes("90b") || arnLower.includes("11b"), + } + } else if (arnLower.includes("nova-pro")) { + // Amazon Nova Pro + modelInfo = { + maxTokens: 5000, + contextWindow: 300_000, + supportsPromptCache: false, + supportsImages: true, + } + } else { + // Default for unknown models or prompt routers + modelInfo = { + maxTokens: 4096, + contextWindow: 128_000, + supportsPromptCache: false, + supportsImages: true, + } + } + + // If modelMaxTokens is explicitly set in options, override the default + if (this.options.modelMaxTokens && this.options.modelMaxTokens > 0) { + modelInfo.maxTokens = this.options.modelMaxTokens + } + + return { + id: this.options.awsCustomArn, + info: modelInfo, + } + } + const modelId = this.options.apiModelId if (modelId) { + // Special case for custom ARN option + if (modelId === "custom-arn") { + // This should not happen as we should have awsCustomArn set + // but just in case, return a default model + return { + id: bedrockDefaultModelId, + info: bedrockModels[bedrockDefaultModelId], + } + } + // For tests, allow any model ID if (process.env.NODE_ENV === "test") { return { @@ -239,7 +569,43 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH // Handle cross-region inference let modelId: string - if (this.options.awsUseCrossRegionInference) { + + // For custom ARNs, use the ARN directly without modification + if (this.options.awsCustomArn) { + modelId = modelConfig.id + logger.debug("Using custom ARN in completePrompt", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + }) + + // Validate ARN format and check region match + const clientRegion = this.client.config.region as string + const validation = validateBedrockArn(modelId, clientRegion) + + if (!validation.isValid) { + logger.error("Invalid ARN format in completePrompt", { + ctx: "bedrock", + modelId, + errorMessage: validation.errorMessage, + }) + throw new Error( + validation.errorMessage || + "Invalid ARN format. ARN should follow the pattern: arn:aws:bedrock:region:account-id:resource-type/resource-name", + ) + } + + // Extract region from ARN + const arnRegion = validation.arnRegion! + + // Log warning if there's a region mismatch + if (validation.errorMessage) { + logger.warn(validation.errorMessage, { + ctx: "bedrock", + arnRegion, + clientRegion, + }) + } + } else if (this.options.awsUseCrossRegionInference) { let regionPrefix = (this.options.awsRegion || "").slice(0, 3) switch (regionPrefix) { case "us-": @@ -265,12 +631,21 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH }, ]), inferenceConfig: { - maxTokens: modelConfig.info.maxTokens || 5000, + maxTokens: modelConfig.info.maxTokens || 4096, temperature: this.options.modelTemperature ?? BEDROCK_DEFAULT_TEMPERATURE, topP: 0.1, }, } + // Log the payload for debugging custom ARN issues + if (this.options.awsCustomArn) { + logger.debug("Bedrock completePrompt request details", { + ctx: "bedrock", + clientRegion: this.client.config.region, + payload: JSON.stringify(payload, null, 2), + }) + } + const command = new ConverseCommand(payload) const response = await this.client.send(command) @@ -282,11 +657,67 @@ export class AwsBedrockHandler extends BaseProvider implements SingleCompletionH return output.content } } catch (parseError) { - console.error("Failed to parse Bedrock response:", parseError) + logger.error("Failed to parse Bedrock response", { + ctx: "bedrock", + error: parseError instanceof Error ? parseError : String(parseError), + }) } } return "" } catch (error) { + // Enhanced error handling for custom ARN issues + if (this.options.awsCustomArn) { + logger.error("Error occurred with custom ARN in completePrompt", { + ctx: "bedrock", + customArn: this.options.awsCustomArn, + error: error instanceof Error ? error : String(error), + }) + + if (error instanceof Error) { + const errorMessage = error.message.toLowerCase() + + // Access denied errors + if ( + errorMessage.includes("access") && + (errorMessage.includes("model") || errorMessage.includes("denied")) + ) { + throw new Error( + `Bedrock custom ARN error: You don't have access to the model with the specified ARN. Please verify: +1. The ARN is correct and points to a valid model +2. Your AWS credentials have permission to access this model (check IAM policies) +3. The region in the ARN matches the region where the model is deployed +4. If using a provisioned model, ensure it's active and not in a failed state`, + ) + } + // Model not found errors + else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) { + throw new Error( + `Bedrock custom ARN error: The specified ARN does not exist or is invalid. Please check: +1. The ARN format is correct (arn:aws:bedrock:region:account-id:resource-type/resource-name) +2. The model exists in the specified region +3. The account ID in the ARN is correct +4. The resource type is one of: foundation-model, provisioned-model, or default-prompt-router`, + ) + } + // Throttling errors + else if ( + errorMessage.includes("throttl") || + errorMessage.includes("rate") || + errorMessage.includes("limit") + ) { + throw new Error( + `Bedrock custom ARN error: Request was throttled or rate limited. Please try: +1. Reducing the frequency of requests +2. If using a provisioned model, check its throughput settings +3. Contact AWS support to request a quota increase if needed`, + ) + } else { + throw new Error(`Bedrock custom ARN error: ${error.message}`) + } + } + } + + // Standard error handling if (error instanceof Error) { throw new Error(`Bedrock completion error: ${error.message}`) } diff --git a/src/shared/api.ts b/src/shared/api.ts index 98d595cd03..986a412354 100644 --- a/src/shared/api.ts +++ b/src/shared/api.ts @@ -39,6 +39,7 @@ export interface ApiHandlerOptions { awspromptCacheId?: string awsProfile?: string awsUseProfile?: boolean + awsCustomArn?: string vertexKeyFile?: string vertexJsonCredentials?: string vertexProjectId?: string @@ -99,6 +100,7 @@ export const API_CONFIG_KEYS: GlobalStateKey[] = [ // "awspromptCacheId", // NOT exist on GlobalStateKey "awsProfile", "awsUseProfile", + "awsCustomArn", "vertexKeyFile", "vertexJsonCredentials", "vertexProjectId", diff --git a/src/shared/globalState.ts b/src/shared/globalState.ts index 540b7e72be..e3522e1c0b 100644 --- a/src/shared/globalState.ts +++ b/src/shared/globalState.ts @@ -28,6 +28,7 @@ export const GLOBAL_STATE_KEYS = [ "awsUseCrossRegionInference", "awsProfile", "awsUseProfile", + "awsCustomArn", "vertexKeyFile", "vertexJsonCredentials", "vertexProjectId", diff --git a/test-custom-arn.js b/test-custom-arn.js new file mode 100644 index 0000000000..dd22ed69b0 --- /dev/null +++ b/test-custom-arn.js @@ -0,0 +1,196 @@ +// Test script to verify AWS Bedrock functionality with custom ARNs +// This file should be deleted after testing + +// IMPORTANT: Before running this script, make sure you have: +// 1. Configured an AWS profile in your AWS credentials file (~/.aws/credentials) +// 2. For prompt routing, created a prompt router in AWS Bedrock (https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-routing.html) +// 3. For prompt routing, have the prompt router ARN in the format: arn:aws:bedrock:region:account-id:default-prompt-router/router-name + +const { BedrockRuntimeClient, ConverseCommand } = require("@aws-sdk/client-bedrock-runtime") +const { fromIni } = require("@aws-sdk/credential-providers") + +// The model ID or ARN provided by the user (not stored in source code) +const modelIdOrArn = process.env.CUSTOM_ARN +// The AWS profile to use for authentication +const awsProfile = process.env.AWS_PROFILE || "default" + +if (!modelIdOrArn) { + console.error("Please provide a model ID or ARN via the CUSTOM_ARN environment variable") + process.exit(1) +} + +console.log(`Using AWS profile: ${awsProfile}`) + +// Check if the input is an ARN or a model ID +const arnRegex = + /^arn:aws:bedrock:([^:]+):(\d+):(foundation-model|provisioned-model|default-prompt-router|prompt-router)\/(.+)$/ +const match = modelIdOrArn.match(arnRegex) +const isArn = !!match + +// If it's not an ARN, assume it's a model ID +if (!isArn) { + console.log(`Using model ID: ${modelIdOrArn}`) +} + +// Use us-west-2 region by default +const defaultRegion = "us-west-2" +// Always use the default region, ignoring the region in the ARN +const region = defaultRegion + +if (isArn) { + console.log(`Using region: ${region} with AWS profile "${awsProfile}" (overriding ARN region: ${match[1]})`) +} else { + console.log(`Using region: ${region} with AWS profile "${awsProfile}"`) +} + +// Create a client with the specified AWS profile +let client +try { + client = new BedrockRuntimeClient({ + region: region, + credentials: fromIni({ + profile: awsProfile, + }), + }) + console.log("Successfully created Bedrock client") +} catch (error) { + console.error("Error creating Bedrock client:", error) + process.exit(1) +} + +// Use the input as the model ID +if (isArn) { + console.log(`Using custom ARN as model ID: ${modelIdOrArn}`) +} else { + console.log(`Using standard model ID: ${modelIdOrArn}`) +} + +const payload = { + modelId: modelIdOrArn, + messages: [ + { + role: "user", + content: [ + { + text: isArn + ? "Hello, can you verify that this prompt router ARN is working correctly? This is a test of AWS Bedrock Intelligent Prompt Routing." + : `Hello, can you verify that this model ID is working correctly with the specified AWS profile?`, + }, + ], + }, + ], + inferenceConfig: { + // For Claude models, use appropriate token limits based on model type + // Claude 3.7 Sonnet: 8192, Claude 3.5 Sonnet: 8192, Claude 3 Opus: 4096, Claude 3 Haiku: 4096 + maxTokens: 4096, // Conservative default that works for all Claude models + temperature: 0.3, + topP: 0.1, + }, +} + +console.log( + isArn + ? "Sending request to Bedrock API using prompt router ARN..." + : "Sending request to Bedrock API using standard model ID...", +) + +async function testCustomArn() { + try { + const command = new ConverseCommand(payload) + const response = await client.send(command) + + // Handle the response format where output is an object + if (response.output && typeof response.output === "object") { + if (response.output.message && response.output.message.content) { + console.log("Success! Received response:") + console.log(JSON.stringify(response)) + console.log(response.output.message.content) + return + } + } + // Handle the response format where output is a Uint8Array + else if (response.output && response.output instanceof Uint8Array) { + try { + const outputStr = new TextDecoder().decode(response.output) + const output = JSON.parse(outputStr) + if (output.content) { + console.log("Success! Received response:") + console.log(output.content) + return + } + } catch (parseError) { + console.error("Failed to parse Bedrock response:", parseError) + } + } + console.error("No valid response content received") + } catch (error) { + console.error(isArn ? "Error occurred with custom ARN:" : "Error occurred with model ID:", error) + + if (error.message) { + const errorMessage = error.message.toLowerCase() + + // Access denied errors + if ( + errorMessage.includes("access") && + (errorMessage.includes("model") || errorMessage.includes("denied")) + ) { + if (isArn) { + console.error("\nThis appears to be a permissions issue with the prompt router ARN. Please verify:") + console.error("1. The ARN is correct and points to a valid prompt router") + console.error( + `2. Your AWS credentials (${awsProfile} profile) have permission to access this prompt router`, + ) + console.error("3. The region in the ARN matches the region where the prompt router is deployed") + console.error("4. The prompt router is properly configured and active") + } else { + console.error("\nThis appears to be a permissions issue with the model. Please verify:") + console.error( + `1. Your AWS credentials (${awsProfile} profile) have permission to access this model`, + ) + console.error("2. The model exists in the specified region") + console.error("3. The model is available for use with your account") + } + } + // Model not found errors + else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) { + if (isArn) { + console.error("\nThis appears to be an invalid prompt router ARN. Please check:") + console.error( + "1. The ARN format is correct (arn:aws:bedrock:region:account-id:default-prompt-router/router-name)", + ) + console.error("2. The prompt router exists in the specified region") + console.error("3. The account ID in the ARN is correct") + } else { + console.error("\nThis appears to be an invalid model ID. Please check:") + console.error("1. The model ID is correct") + console.error("2. The model exists in the specified region") + } + } + // Validation errors + else if (errorMessage.includes("validation")) { + if (isArn) { + console.error("\nThis appears to be a validation error with the prompt router ARN. Please check:") + console.error("1. The ARN format is correct") + console.error("2. The prompt router is properly configured") + console.error("3. The request payload is valid for prompt routing") + } else { + console.error("\nThis appears to be a validation error with the model ID. Please check:") + console.error("1. The model ID format is correct") + console.error("2. The request payload is valid for this model") + } + } + // Throttling errors + else if ( + errorMessage.includes("throttl") || + errorMessage.includes("rate") || + errorMessage.includes("limit") + ) { + console.error("\nThis appears to be a throttling or rate limit issue. Please try:") + console.error("1. Reducing the frequency of requests") + console.error("2. Contact AWS support to request a quota increase if needed") + } + } + } +} + +testCustomArn() diff --git a/webview-ui/src/components/settings/ApiOptions.tsx b/webview-ui/src/components/settings/ApiOptions.tsx index f7982080c8..656a6831cb 100644 --- a/webview-ui/src/components/settings/ApiOptions.tsx +++ b/webview-ui/src/components/settings/ApiOptions.tsx @@ -41,7 +41,7 @@ import { VSCodeButtonLink } from "../common/VSCodeButtonLink" import { ModelInfoView } from "./ModelInfoView" import { ModelPicker } from "./ModelPicker" import { TemperatureControl } from "./TemperatureControl" -import { validateApiConfiguration, validateModelId } from "@/utils/validate" +import { validateApiConfiguration, validateModelId, validateBedrockArn } from "@/utils/validate" import { ApiErrorMessage } from "./ApiErrorMessage" import { ThinkingBudget } from "./ThinkingBudget" @@ -1267,14 +1267,82 @@ const ApiOptions = ({ { - setApiConfigurationField("apiModelId", typeof value == "string" ? value : value?.value) + const modelValue = typeof value == "string" ? value : value?.value + setApiConfigurationField("apiModelId", modelValue) + + // Clear custom ARN if not using custom ARN option + if (modelValue !== "custom-arn" && selectedProvider === "bedrock") { + setApiConfigurationField("awsCustomArn", "") + } }} - options={selectedProviderModelOptions} + options={[ + ...selectedProviderModelOptions, + ...(selectedProvider === "bedrock" + ? [{ value: "custom-arn", label: "Use custom ARN..." }] + : []), + ]} className="w-full" /> + + {selectedProvider === "bedrock" && selectedModelId === "custom-arn" && ( + <> + { + const value = (e.target as HTMLInputElement).value + setApiConfigurationField("awsCustomArn", value) + }} + placeholder="Enter ARN (e.g. arn:aws:bedrock:us-east-1:123456789012:foundation-model/my-model)" + className="w-full"> + Custom ARN + +
+ Enter a valid AWS Bedrock ARN for the model you want to use. Format examples: +
    +
  • + arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-sonnet-20240229-v1:0 +
    • +
  • + arn:aws:bedrock:us-west-2:123456789012:provisioned-model/my-provisioned-model +
    • +
  • + arn:aws:bedrock:us-east-1:123456789012:default-prompt-router/anthropic.claude:1 +
    • +
+ Make sure the region in the ARN matches your selected AWS Region above. +
+ {apiConfiguration?.awsCustomArn && + (() => { + const validation = validateBedrockArn( + apiConfiguration.awsCustomArn, + apiConfiguration.awsRegion, + ) + + if (!validation.isValid) { + return ( +
+ {validation.errorMessage || + "Invalid ARN format. Please check the examples above."} +
+ ) + } + + if (validation.errorMessage) { + return ( +
+ {validation.errorMessage} +
+ ) + } + + return null + })()} + ======= + + )} Date: Mon, 10 Mar 2025 22:15:52 -0400 Subject: [PATCH 2/2] Cleanup --- src/api/providers/__tests__/bedrock.test.ts | 2 +- test-custom-arn.js | 196 -------------------- 2 files changed, 1 insertion(+), 197 deletions(-) delete mode 100644 test-custom-arn.js diff --git a/src/api/providers/__tests__/bedrock.test.ts b/src/api/providers/__tests__/bedrock.test.ts index f778621e9c..45d5270237 100644 --- a/src/api/providers/__tests__/bedrock.test.ts +++ b/src/api/providers/__tests__/bedrock.test.ts @@ -326,7 +326,7 @@ describe("AwsBedrockHandler", () => { }) const modelInfo = customArnHandler.getModel() expect(modelInfo.id).toBe("arn:aws:bedrock:us-east-1::foundation-model/custom-model") - expect(modelInfo.info.maxTokens).toBe(5000) + expect(modelInfo.info.maxTokens).toBe(4096) expect(modelInfo.info.contextWindow).toBe(128_000) expect(modelInfo.info.supportsPromptCache).toBe(false) }) diff --git a/test-custom-arn.js b/test-custom-arn.js deleted file mode 100644 index dd22ed69b0..0000000000 --- a/test-custom-arn.js +++ /dev/null @@ -1,196 +0,0 @@ -// Test script to verify AWS Bedrock functionality with custom ARNs -// This file should be deleted after testing - -// IMPORTANT: Before running this script, make sure you have: -// 1. Configured an AWS profile in your AWS credentials file (~/.aws/credentials) -// 2. For prompt routing, created a prompt router in AWS Bedrock (https://docs.aws.amazon.com/bedrock/latest/userguide/prompt-routing.html) -// 3. For prompt routing, have the prompt router ARN in the format: arn:aws:bedrock:region:account-id:default-prompt-router/router-name - -const { BedrockRuntimeClient, ConverseCommand } = require("@aws-sdk/client-bedrock-runtime") -const { fromIni } = require("@aws-sdk/credential-providers") - -// The model ID or ARN provided by the user (not stored in source code) -const modelIdOrArn = process.env.CUSTOM_ARN -// The AWS profile to use for authentication -const awsProfile = process.env.AWS_PROFILE || "default" - -if (!modelIdOrArn) { - console.error("Please provide a model ID or ARN via the CUSTOM_ARN environment variable") - process.exit(1) -} - -console.log(`Using AWS profile: ${awsProfile}`) - -// Check if the input is an ARN or a model ID -const arnRegex = - /^arn:aws:bedrock:([^:]+):(\d+):(foundation-model|provisioned-model|default-prompt-router|prompt-router)\/(.+)$/ -const match = modelIdOrArn.match(arnRegex) -const isArn = !!match - -// If it's not an ARN, assume it's a model ID -if (!isArn) { - console.log(`Using model ID: ${modelIdOrArn}`) -} - -// Use us-west-2 region by default -const defaultRegion = "us-west-2" -// Always use the default region, ignoring the region in the ARN -const region = defaultRegion - -if (isArn) { - console.log(`Using region: ${region} with AWS profile "${awsProfile}" (overriding ARN region: ${match[1]})`) -} else { - console.log(`Using region: ${region} with AWS profile "${awsProfile}"`) -} - -// Create a client with the specified AWS profile -let client -try { - client = new BedrockRuntimeClient({ - region: region, - credentials: fromIni({ - profile: awsProfile, - }), - }) - console.log("Successfully created Bedrock client") -} catch (error) { - console.error("Error creating Bedrock client:", error) - process.exit(1) -} - -// Use the input as the model ID -if (isArn) { - console.log(`Using custom ARN as model ID: ${modelIdOrArn}`) -} else { - console.log(`Using standard model ID: ${modelIdOrArn}`) -} - -const payload = { - modelId: modelIdOrArn, - messages: [ - { - role: "user", - content: [ - { - text: isArn - ? "Hello, can you verify that this prompt router ARN is working correctly? This is a test of AWS Bedrock Intelligent Prompt Routing." - : `Hello, can you verify that this model ID is working correctly with the specified AWS profile?`, - }, - ], - }, - ], - inferenceConfig: { - // For Claude models, use appropriate token limits based on model type - // Claude 3.7 Sonnet: 8192, Claude 3.5 Sonnet: 8192, Claude 3 Opus: 4096, Claude 3 Haiku: 4096 - maxTokens: 4096, // Conservative default that works for all Claude models - temperature: 0.3, - topP: 0.1, - }, -} - -console.log( - isArn - ? "Sending request to Bedrock API using prompt router ARN..." - : "Sending request to Bedrock API using standard model ID...", -) - -async function testCustomArn() { - try { - const command = new ConverseCommand(payload) - const response = await client.send(command) - - // Handle the response format where output is an object - if (response.output && typeof response.output === "object") { - if (response.output.message && response.output.message.content) { - console.log("Success! Received response:") - console.log(JSON.stringify(response)) - console.log(response.output.message.content) - return - } - } - // Handle the response format where output is a Uint8Array - else if (response.output && response.output instanceof Uint8Array) { - try { - const outputStr = new TextDecoder().decode(response.output) - const output = JSON.parse(outputStr) - if (output.content) { - console.log("Success! Received response:") - console.log(output.content) - return - } - } catch (parseError) { - console.error("Failed to parse Bedrock response:", parseError) - } - } - console.error("No valid response content received") - } catch (error) { - console.error(isArn ? "Error occurred with custom ARN:" : "Error occurred with model ID:", error) - - if (error.message) { - const errorMessage = error.message.toLowerCase() - - // Access denied errors - if ( - errorMessage.includes("access") && - (errorMessage.includes("model") || errorMessage.includes("denied")) - ) { - if (isArn) { - console.error("\nThis appears to be a permissions issue with the prompt router ARN. Please verify:") - console.error("1. The ARN is correct and points to a valid prompt router") - console.error( - `2. Your AWS credentials (${awsProfile} profile) have permission to access this prompt router`, - ) - console.error("3. The region in the ARN matches the region where the prompt router is deployed") - console.error("4. The prompt router is properly configured and active") - } else { - console.error("\nThis appears to be a permissions issue with the model. Please verify:") - console.error( - `1. Your AWS credentials (${awsProfile} profile) have permission to access this model`, - ) - console.error("2. The model exists in the specified region") - console.error("3. The model is available for use with your account") - } - } - // Model not found errors - else if (errorMessage.includes("not found") || errorMessage.includes("does not exist")) { - if (isArn) { - console.error("\nThis appears to be an invalid prompt router ARN. Please check:") - console.error( - "1. The ARN format is correct (arn:aws:bedrock:region:account-id:default-prompt-router/router-name)", - ) - console.error("2. The prompt router exists in the specified region") - console.error("3. The account ID in the ARN is correct") - } else { - console.error("\nThis appears to be an invalid model ID. Please check:") - console.error("1. The model ID is correct") - console.error("2. The model exists in the specified region") - } - } - // Validation errors - else if (errorMessage.includes("validation")) { - if (isArn) { - console.error("\nThis appears to be a validation error with the prompt router ARN. Please check:") - console.error("1. The ARN format is correct") - console.error("2. The prompt router is properly configured") - console.error("3. The request payload is valid for prompt routing") - } else { - console.error("\nThis appears to be a validation error with the model ID. Please check:") - console.error("1. The model ID format is correct") - console.error("2. The request payload is valid for this model") - } - } - // Throttling errors - else if ( - errorMessage.includes("throttl") || - errorMessage.includes("rate") || - errorMessage.includes("limit") - ) { - console.error("\nThis appears to be a throttling or rate limit issue. Please try:") - console.error("1. Reducing the frequency of requests") - console.error("2. Contact AWS support to request a quota increase if needed") - } - } - } -} - -testCustomArn()