mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-06 08:18:39 +00:00
Add enhanced error handling and logging for AWS Bedrock custom ARNs
This commit is contained in:
Smartsheet-JB-Brown
parent
224303660b
commit
171037a938
8 changed files with 887 additions and 34 deletions
75
src/api/providers/__tests__/bedrock-custom-arn.test.ts
Normal file
75
src/api/providers/__tests__/bedrock-custom-arn.test.ts
Normal file
|
|
@ -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")
|
||||
})
|
||||
})
|
||||
|
|
@ -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()
|
||||
})
|
||||
})
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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}`)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -28,6 +28,7 @@ export const GLOBAL_STATE_KEYS = [
|
|||
"awsUseCrossRegionInference",
|
||||
"awsProfile",
|
||||
"awsUseProfile",
|
||||
"awsCustomArn",
|
||||
"vertexKeyFile",
|
||||
"vertexJsonCredentials",
|
||||
"vertexProjectId",
|
||||
|
|
|
|||
196
test-custom-arn.js
Normal file
196
test-custom-arn.js
Normal file
|
|
@ -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()
|
||||
|
|
@ -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 = ({
|
|||
</label>
|
||||
<Dropdown
|
||||
id="model-id"
|
||||
value={selectedModelId}
|
||||
value={selectedModelId === "custom-arn" ? "custom-arn" : selectedModelId}
|
||||
onChange={(value) => {
|
||||
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"
|
||||
/>
|
||||
</div>
|
||||
|
||||
{selectedProvider === "bedrock" && selectedModelId === "custom-arn" && (
|
||||
<>
|
||||
<VSCodeTextField
|
||||
value={apiConfiguration?.awsCustomArn || ""}
|
||||
onInput={(e) => {
|
||||
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">
|
||||
<span className="font-medium">Custom ARN</span>
|
||||
</VSCodeTextField>
|
||||
<div className="text-sm text-vscode-descriptionForeground -mt-2">
|
||||
Enter a valid AWS Bedrock ARN for the model you want to use. Format examples:
|
||||
<ul className="list-disc pl-5 mt-1">
|
||||
<li>
|
||||
arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3-sonnet-20240229-v1:0
|
||||
</li>
|
||||
<li>
|
||||
arn:aws:bedrock:us-west-2:123456789012:provisioned-model/my-provisioned-model
|
||||
</li>
|
||||
<li>
|
||||
arn:aws:bedrock:us-east-1:123456789012:default-prompt-router/anthropic.claude:1
|
||||
</li>
|
||||
</ul>
|
||||
Make sure the region in the ARN matches your selected AWS Region above.
|
||||
</div>
|
||||
{apiConfiguration?.awsCustomArn &&
|
||||
(() => {
|
||||
const validation = validateBedrockArn(
|
||||
apiConfiguration.awsCustomArn,
|
||||
apiConfiguration.awsRegion,
|
||||
)
|
||||
|
||||
if (!validation.isValid) {
|
||||
return (
|
||||
<div className="text-sm text-vscode-errorForeground mt-2">
|
||||
{validation.errorMessage ||
|
||||
"Invalid ARN format. Please check the examples above."}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
if (validation.errorMessage) {
|
||||
return (
|
||||
<div className="text-sm text-vscode-errorForeground mt-2">
|
||||
{validation.errorMessage}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
return null
|
||||
})()}
|
||||
=======
|
||||
</>
|
||||
)}
|
||||
<ModelInfoView
|
||||
selectedModelId={selectedModelId}
|
||||
modelInfo={selectedModelInfo}
|
||||
|
|
@ -1333,6 +1401,19 @@ export function normalizeApiConfiguration(apiConfiguration?: ApiConfiguration) {
|
|||
case "anthropic":
|
||||
return getProviderData(anthropicModels, anthropicDefaultModelId)
|
||||
case "bedrock":
|
||||
// Special case for custom ARN
|
||||
if (modelId === "custom-arn") {
|
||||
return {
|
||||
selectedProvider: provider,
|
||||
selectedModelId: "custom-arn",
|
||||
selectedModelInfo: {
|
||||
maxTokens: 5000,
|
||||
contextWindow: 128_000,
|
||||
supportsPromptCache: false,
|
||||
supportsImages: true,
|
||||
},
|
||||
}
|
||||
}
|
||||
return getProviderData(bedrockModels, bedrockDefaultModelId)
|
||||
case "vertex":
|
||||
return getProviderData(vertexModels, vertexDefaultModelId)
|
||||
|
|
|
|||
|
|
@ -80,6 +80,44 @@ export function validateApiConfiguration(apiConfiguration?: ApiConfiguration): s
|
|||
|
||||
return undefined
|
||||
}
|
||||
/**
|
||||
* 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 }
|
||||
*/
|
||||
export 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. Please check the format requirements.",
|
||||
}
|
||||
}
|
||||
|
||||
// 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,
|
||||
}
|
||||
}
|
||||
|
||||
export function validateModelId(
|
||||
apiConfiguration?: ApiConfiguration,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue