Skip to content

Commit 6acf535

Browse files
bizob2828bizob2828
authored
feat: Added support for region-prefixed Bedrock models (#2947)
1 parent 772f007 commit 6acf535

File tree

5 files changed

+309
-5
lines changed

5 files changed

+309
-5
lines changed

‎lib/llm-events/aws-bedrock/bedrock-command.js

+2-2
Original file line numberDiff line numberDiff line change
@@ -124,11 +124,11 @@ class BedrockCommand {
124124
}
125125

126126
isClaude() {
127-
return this.#modelId.startsWith('anthropic.claude-v')
127+
return this.#modelId.split('.').slice(-2).join('.').startsWith('anthropic.claude-v')
128128
}
129129

130130
isClaude3() {
131-
return this.#modelId.startsWith('anthropic.claude-3')
131+
return this.#modelId.split('.').slice(-2).join('.').startsWith('anthropic.claude-3')
132132
}
133133

134134
isCohere() {

‎test/lib/aws-server-stubs/ai-server/index.js

+27-3
Original file line numberDiff line numberDiff line change
@@ -89,20 +89,44 @@ function handler(req, res) {
8989
case 'anthropic.claude-v1':
9090
case 'anthropic.claude-instant-v1':
9191
case 'anthropic.claude-v2':
92-
case 'anthropic.claude-v2:1': {
92+
case 'anthropic.claude-v2:1':
93+
case 'us.anthropic.claude-v1':
94+
case 'us.anthropic.claude-instant-v1':
95+
case 'us.anthropic.claude-v2':
96+
case 'us.anthropic.claude-v2:1':
97+
case 'eu.anthropic.claude-v1':
98+
case 'eu.anthropic.claude-instant-v1':
99+
case 'eu.anthropic.claude-v2':
100+
case 'eu.anthropic.claude-v2:1':
101+
case 'apac.anthropic.claude-v1':
102+
case 'apac.anthropic.claude-instant-v1':
103+
case 'apac.anthropic.claude-v2':
104+
case 'apac.anthropic.claude-v2:1':{
93105
response = responses.claude.get(payload.prompt)
94106
break
95107
}
96108

97109
case 'anthropic.claude-3-haiku-20240307-v1:0':
98110
case 'anthropic.claude-3-opus-20240229-v1:0':
99-
case 'anthropic.claude-3-sonnet-20240229-v1:0': {
111+
case 'anthropic.claude-3-sonnet-20240229-v1:0':
112+
case 'us.anthropic.claude-3-haiku-20240307-v1:0':
113+
case 'us.anthropic.claude-3-opus-20240229-v1:0':
114+
case 'us.anthropic.claude-3-sonnet-20240229-v1:0':
115+
case 'eu.anthropic.claude-3-haiku-20240307-v1:0':
116+
case 'eu.anthropic.claude-3-opus-20240229-v1:0':
117+
case 'eu.anthropic.claude-3-sonnet-20240229-v1:0':
118+
case 'apac.anthropic.claude-3-haiku-20240307-v1:0':
119+
case 'apac.anthropic.claude-3-opus-20240229-v1:0':
120+
case 'apac.anthropic.claude-3-sonnet-20240229-v1:0': {
100121
response = responses.claude3.get(payload?.messages?.[0]?.content)
101122
break
102123
}
103124

104125
// Chunked claude model
105-
case 'anthropic.claude-3-5-sonnet-20240620-v1:0': {
126+
case 'anthropic.claude-3-5-sonnet-20240620-v1:0':
127+
case 'us.anthropic.claude-3-5-sonnet-20240620-v1:0':
128+
case 'eu.anthropic.claude-3-5-sonnet-20240620-v1:0':
129+
case 'apac.anthropic.claude-3-5-sonnet-20240620-v1:0':{
106130
response = responses.claude3.get(payload?.messages?.[0]?.content?.[0].text)
107131
break
108132
}

‎test/unit/llm-events/aws-bedrock/bedrock-command.test.js

+130
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,13 @@ const claude = {
2424
}
2525
}
2626

27+
const regionClaude = {
28+
modelId: 'us.anthropic.claude-v1',
29+
body: {
30+
prompt: '\n\nHuman: yes\n\nAssistant:'
31+
}
32+
}
33+
2734
const claude35 = {
2835
modelId: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
2936
body: {
@@ -35,13 +42,30 @@ const claude35 = {
3542
}
3643
}
3744

45+
const regionClaude35 = {
46+
modelId: 'us.anthropic.claude-3-5-sonnet-20240620-v1:0',
47+
body: {
48+
messages: [
49+
{ role: 'user', content: [{ type: 'text', text: 'who are' }] },
50+
{ role: 'assistant', content: [{ type: 'text', text: 'researching' }] },
51+
{ role: 'user', content: [{ type: 'text', text: 'you' }] }
52+
]
53+
}
54+
}
3855
const claude3 = {
3956
modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
4057
body: {
4158
messages: [{ role: 'user', content: 'who are you' }]
4259
}
4360
}
4461

62+
const regionClaude3 = {
63+
modelId: 'us.anthropic.claude-3-haiku-20240307-v1:0',
64+
body: {
65+
messages: [{ role: 'user', content: 'who are you' }]
66+
}
67+
}
68+
4569
const cohere = {
4670
modelId: 'cohere.command-text-v14',
4771
body: {
@@ -154,6 +178,17 @@ test('claude minimal command works', async (t) => {
154178
assert.equal(cmd.temperature, undefined)
155179
})
156180

181+
test('region specific claude minimal command works', async (t) => {
182+
t.nr.updatePayload(structuredClone(regionClaude))
183+
const cmd = new BedrockCommand(t.nr.input)
184+
assert.equal(cmd.isClaude(), true)
185+
assert.equal(cmd.maxTokens, undefined)
186+
assert.equal(cmd.modelId, regionClaude.modelId)
187+
assert.equal(cmd.modelType, 'completion')
188+
assert.deepEqual(cmd.prompt, [{ role: 'user', content: claude.body.prompt }])
189+
assert.equal(cmd.temperature, undefined)
190+
})
191+
157192
test('claude complete command works', async (t) => {
158193
const payload = structuredClone(claude)
159194
payload.body.max_tokens_to_sample = 25
@@ -168,6 +203,20 @@ test('claude complete command works', async (t) => {
168203
assert.equal(cmd.temperature, payload.body.temperature)
169204
})
170205

206+
test('region specific claude complete command works', async (t) => {
207+
const payload = structuredClone(regionClaude)
208+
payload.body.max_tokens_to_sample = 25
209+
payload.body.temperature = 0.5
210+
t.nr.updatePayload(payload)
211+
const cmd = new BedrockCommand(t.nr.input)
212+
assert.equal(cmd.isClaude(), true)
213+
assert.equal(cmd.maxTokens, 25)
214+
assert.equal(cmd.modelId, payload.modelId)
215+
assert.equal(cmd.modelType, 'completion')
216+
assert.deepEqual(cmd.prompt, [{ role: 'user', content: payload.body.prompt }])
217+
assert.equal(cmd.temperature, payload.body.temperature)
218+
})
219+
171220
test('claude3 minimal command works', async (t) => {
172221
t.nr.updatePayload(structuredClone(claude3))
173222
const cmd = new BedrockCommand(t.nr.input)
@@ -179,6 +228,17 @@ test('claude3 minimal command works', async (t) => {
179228
assert.equal(cmd.temperature, undefined)
180229
})
181230

231+
test('region specific claude3 minimal command works', async (t) => {
232+
t.nr.updatePayload(structuredClone(regionClaude3))
233+
const cmd = new BedrockCommand(t.nr.input)
234+
assert.equal(cmd.isClaude3(), true)
235+
assert.equal(cmd.maxTokens, undefined)
236+
assert.equal(cmd.modelId, regionClaude3.modelId)
237+
assert.equal(cmd.modelType, 'completion')
238+
assert.deepEqual(cmd.prompt, claude3.body.messages)
239+
assert.equal(cmd.temperature, undefined)
240+
})
241+
182242
test('claude3 complete command works', async (t) => {
183243
const payload = structuredClone(claude3)
184244
payload.body.max_tokens = 25
@@ -193,6 +253,20 @@ test('claude3 complete command works', async (t) => {
193253
assert.equal(cmd.temperature, payload.body.temperature)
194254
})
195255

256+
test('region specific claude3 complete command works', async (t) => {
257+
const payload = structuredClone(regionClaude3)
258+
payload.body.max_tokens = 25
259+
payload.body.temperature = 0.5
260+
t.nr.updatePayload(payload)
261+
const cmd = new BedrockCommand(t.nr.input)
262+
assert.equal(cmd.isClaude3(), true)
263+
assert.equal(cmd.maxTokens, 25)
264+
assert.equal(cmd.modelId, payload.modelId)
265+
assert.equal(cmd.modelType, 'completion')
266+
assert.deepEqual(cmd.prompt, payload.body.messages)
267+
assert.equal(cmd.temperature, payload.body.temperature)
268+
})
269+
196270
test('claude35 minimal command works with claude 3 api', async (t) => {
197271
t.nr.updatePayload(structuredClone(claude3))
198272
const cmd = new BedrockCommand(t.nr.input)
@@ -217,6 +291,19 @@ test('claude35 malformed payload produces reasonable values', async (t) => {
217291
assert.equal(cmd.temperature, undefined)
218292
})
219293

294+
test('region specific claude35 malformed payload produces reasonable values', async (t) => {
295+
const malformedPayload = structuredClone(regionClaude35)
296+
malformedPayload.body = {}
297+
t.nr.updatePayload(malformedPayload)
298+
const cmd = new BedrockCommand(t.nr.input)
299+
assert.equal(cmd.isClaude3(), true)
300+
assert.equal(cmd.maxTokens, undefined)
301+
assert.equal(cmd.modelId, regionClaude35.modelId)
302+
assert.equal(cmd.modelType, 'completion')
303+
assert.deepEqual(cmd.prompt, [])
304+
assert.equal(cmd.temperature, undefined)
305+
})
306+
220307
test('claude35 skips a message that is null in `body.messages`', async (t) => {
221308
const malformedPayload = structuredClone(claude35)
222309
malformedPayload.body.messages = [{ role: 'user', content: 'who are you' }, null]
@@ -226,6 +313,15 @@ test('claude35 skips a message that is null in `body.messages`', async (t) => {
226313
assert.deepEqual(cmd.prompt, [{ role: 'user', content: 'who are you' }])
227314
})
228315

316+
test('region specific claude35 skips a message that is null in `body.messages`', async (t) => {
317+
const malformedPayload = structuredClone(regionClaude35)
318+
malformedPayload.body.messages = [{ role: 'user', content: 'who are you' }, null]
319+
t.nr.updatePayload(malformedPayload)
320+
const cmd = new BedrockCommand(t.nr.input)
321+
assert.equal(cmd.isClaude3(), true)
322+
assert.deepEqual(cmd.prompt, [{ role: 'user', content: 'who are you' }])
323+
})
324+
229325
test('claude35 handles defaulting prompt to empty array when `body.messages` is null', async (t) => {
230326
const malformedPayload = structuredClone(claude35)
231327
malformedPayload.body.messages = null
@@ -235,6 +331,15 @@ test('claude35 handles defaulting prompt to empty array when `body.messages` is
235331
assert.deepEqual(cmd.prompt, [])
236332
})
237333

334+
test('region specific claude35 handles defaulting prompt to empty array when `body.messages` is null', async (t) => {
335+
const malformedPayload = structuredClone(regionClaude35)
336+
malformedPayload.body.messages = null
337+
t.nr.updatePayload(malformedPayload)
338+
const cmd = new BedrockCommand(t.nr.input)
339+
assert.equal(cmd.isClaude3(), true)
340+
assert.deepEqual(cmd.prompt, [])
341+
})
342+
238343
test('claude35 minimal command works', async (t) => {
239344
t.nr.updatePayload(structuredClone(claude35))
240345
const cmd = new BedrockCommand(t.nr.input)
@@ -246,6 +351,17 @@ test('claude35 minimal command works', async (t) => {
246351
assert.equal(cmd.temperature, undefined)
247352
})
248353

354+
test('region specific claude35 minimal command works', async (t) => {
355+
t.nr.updatePayload(structuredClone(regionClaude35))
356+
const cmd = new BedrockCommand(t.nr.input)
357+
assert.equal(cmd.isClaude3(), true)
358+
assert.equal(cmd.maxTokens, undefined)
359+
assert.equal(cmd.modelId, regionClaude35.modelId)
360+
assert.equal(cmd.modelType, 'completion')
361+
assert.deepEqual(cmd.prompt, [{ role: 'user', content: 'who are' }, { role: 'assistant', content: 'researching' }, { role: 'user', content: 'you' }])
362+
assert.equal(cmd.temperature, undefined)
363+
})
364+
249365
test('claude35 complete command works', async (t) => {
250366
const payload = structuredClone(claude35)
251367
payload.body.max_tokens = 25
@@ -260,6 +376,20 @@ test('claude35 complete command works', async (t) => {
260376
assert.equal(cmd.temperature, payload.body.temperature)
261377
})
262378

379+
test('region specific claude35 complete command works', async (t) => {
380+
const payload = structuredClone(regionClaude35)
381+
payload.body.max_tokens = 25
382+
payload.body.temperature = 0.5
383+
t.nr.updatePayload(payload)
384+
const cmd = new BedrockCommand(t.nr.input)
385+
assert.equal(cmd.isClaude3(), true)
386+
assert.equal(cmd.maxTokens, 25)
387+
assert.equal(cmd.modelId, payload.modelId)
388+
assert.equal(cmd.modelType, 'completion')
389+
assert.deepEqual(cmd.prompt, [{ role: 'user', content: 'who are' }, { role: 'assistant', content: 'researching' }, { role: 'user', content: 'you' }])
390+
assert.equal(cmd.temperature, payload.body.temperature)
391+
})
392+
263393
test('cohere minimal command works', async (t) => {
264394
t.nr.updatePayload(structuredClone(cohere))
265395
const cmd = new BedrockCommand(t.nr.input)

‎test/unit/llm-events/aws-bedrock/stream-handler.test.js

+72
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,44 @@ test('handles claude streams', async (t) => {
117117
assert.equal(br.statusCode, 200)
118118
})
119119

120+
test('handles region specific claude streams', async (t) => {
121+
t.nr.passThroughParams.bedrockCommand.isClaude = () => true
122+
t.nr.chunks = [
123+
{ completion: '1', stop_reason: null },
124+
{ completion: '2', stop_reason: 'done', ...t.nr.metrics }
125+
]
126+
const handler = new StreamHandler(t.nr)
127+
128+
assert.equal(handler.generator.name, 'handleClaude')
129+
for await (const event of handler.generator()) {
130+
assert.equal(event.chunk.bytes.constructor, Uint8Array)
131+
}
132+
assert.deepStrictEqual(handler.response, {
133+
response: {
134+
headers: {
135+
'x-amzn-requestid': 'aws-req-1'
136+
},
137+
statusCode: 200
138+
},
139+
output: {
140+
body: new TextEncoder().encode(JSON.stringify({ completion: '12', stop_reason: 'done' }))
141+
}
142+
})
143+
144+
const bc = new BedrockCommand({
145+
modelId: 'us.anthropic.claude-v1',
146+
body: JSON.stringify({
147+
prompt: 'prompt',
148+
maxTokens: 5
149+
})
150+
})
151+
const br = new BedrockResponse({ bedrockCommand: bc, response: handler.response })
152+
assert.equal(br.completions.length, 1)
153+
assert.equal(br.finishReason, 'done')
154+
assert.equal(br.requestId, 'aws-req-1')
155+
assert.equal(br.statusCode, 200)
156+
})
157+
120158
test('handles claude3streams', async (t) => {
121159
t.nr.passThroughParams.bedrockCommand.isClaude3 = () => true
122160
t.nr.chunks = [
@@ -151,6 +189,40 @@ test('handles claude3streams', async (t) => {
151189
assert.equal(br.statusCode, 200)
152190
})
153191

192+
test('handles region specific claude3streams', async (t) => {
193+
t.nr.passThroughParams.bedrockCommand.isClaude3 = () => true
194+
t.nr.chunks = [
195+
{ type: 'content_block_delta', delta: { type: 'text_delta', text: '42' } },
196+
{ type: 'message_delta', delta: { stop_reason: 'done' } },
197+
{ type: 'message_stop', ...t.nr.metrics }
198+
]
199+
const handler = new StreamHandler(t.nr)
200+
201+
assert.equal(handler.generator.name, 'handleClaude3')
202+
for await (const event of handler.generator()) {
203+
assert.equal(event.chunk.bytes.constructor, Uint8Array)
204+
}
205+
const foundBody = JSON.parse(new TextDecoder().decode(handler.response.output.body))
206+
assert.deepStrictEqual(foundBody, {
207+
completions: ['42'],
208+
stop_reason: 'done',
209+
type: 'message_stop'
210+
})
211+
212+
const bc = new BedrockCommand({
213+
modelId: 'us.anthropic.claude-3-haiku-20240307-v1:0',
214+
body: JSON.stringify({
215+
messages: [{ content: 'prompt' }],
216+
maxTokens: 5
217+
})
218+
})
219+
const br = new BedrockResponse({ bedrockCommand: bc, response: handler.response })
220+
assert.equal(br.completions.length, 1)
221+
assert.equal(br.finishReason, 'done')
222+
assert.equal(br.requestId, 'aws-req-1')
223+
assert.equal(br.statusCode, 200)
224+
})
225+
154226
test('handles cohere streams', async (t) => {
155227
t.nr.passThroughParams.bedrockCommand.isCohere = () => true
156228
t.nr.chunks = [

0 commit comments

Comments
 (0)