konnosDashboard
bedrock.ts119 lines · main
1import { createAmazonBedrock } from '@ai-sdk/amazon-bedrock'
2import { createCredentialChain, fromNodeProviderChain } from '@aws-sdk/credential-providers'
3import { CredentialsProviderError } from '@smithy/property-provider'
4import { awsCredentialsProvider } from '@vercel/functions/oidc'
5import { LanguageModel } from 'ai'
6
7import { BedrockModel } from './model.utils'
8import { selectWeightedKey } from './util'
9
10const credentialProvider = createCredentialChain(
11 // Vercel OIDC provider will be used for staging/production
12 vercelOidcProvider,
13
14 // AWS profile will be used for local development
15 fromNodeProviderChain({
16 profile: process.env.AWS_BEDROCK_PROFILE,
17 })
18)
19
20/**
21 * Creates a Vercel OIDC provider for AWS credentials.
22 *
23 * Wraps `awsCredentialsProvider` to properly handle errors
24 * so that it can be used in a credential chain.
25 */
26async function vercelOidcProvider() {
27 try {
28 return await awsCredentialsProvider({
29 roleArn: process.env.AWS_BEDROCK_ROLE_ARN!,
30 })()
31 } catch (error) {
32 const message = error instanceof Error ? error.message : 'Failed to create Vercel OIDC provider'
33
34 // Re-throw using the correct error type and `tryNextLink` option
35 throw new CredentialsProviderError(message, {
36 tryNextLink: true,
37 })
38 }
39}
40
41export async function checkAwsCredentials() {
42 try {
43 const credentials = await credentialProvider()
44 return !!credentials
45 } catch (error) {
46 return false
47 }
48}
49
50export const bedrockRegionMap = {
51 use1: 'us-east-1',
52 use2: 'us-east-2',
53 usw2: 'us-west-2',
54 euc1: 'eu-central-1',
55} as const
56
57export type BedrockRegion = keyof typeof bedrockRegionMap
58
59export const regionPrefixMap: Record<BedrockRegion, string> = {
60 use1: 'us',
61 use2: 'us',
62 usw2: 'us',
63 euc1: 'eu',
64}
65
66export type RegionWeights = Record<BedrockRegion, number>
67
68/**
69 * Weights for distributing requests across Bedrock regions.
70 * Weights are proportional to our rate limits per model per region.
71 */
72const modelRegionWeights: Record<BedrockModel, RegionWeights> = {
73 ['anthropic.claude-3-7-sonnet-20250219-v1:0']: {
74 use1: 40,
75 use2: 10,
76 usw2: 10,
77 euc1: 10,
78 },
79 ['openai.gpt-oss-120b-1:0']: {
80 use1: 0,
81 use2: 0,
82 usw2: 30,
83 euc1: 0,
84 },
85}
86
87/**
88 * Creates a Bedrock client that routes requests to different regions
89 * based on a routing key, with optional OpenAI support.
90 *
91 * Used to load balance requests across multiple regions depending on
92 * their capacities.
93 */
94export function createRoutedBedrock(routingKey?: string) {
95 return async (modelId: BedrockModel): Promise<LanguageModel> => {
96 const regionWeights = modelRegionWeights[modelId]
97
98 // Select the Bedrock region based on the routing key and the model
99 const bedrockRegion = routingKey
100 ? await selectWeightedKey(routingKey, regionWeights)
101 : // There's a few places where getModel is called without a routing key
102 // Will cause disproportionate load on use1 region
103 regionWeights['use1'] > 0
104 ? 'use1'
105 : 'usw2'
106
107 const bedrock = createAmazonBedrock({
108 credentialProvider,
109 region: bedrockRegionMap[bedrockRegion],
110 })
111
112 // Cross-region models require the region prefix
113 const activeRegions = Object.values(regionWeights).filter((weight) => weight > 0).length
114 const modelName = activeRegions > 1 ? `${regionPrefixMap[bedrockRegion]}.${modelId}` : modelId
115
116 const model = bedrock(modelName)
117 return model
118 }
119}