Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Added support for Claude 3+ Chat API in Bedrock #2870

Merged
merged 6 commits into from
Jan 17, 2025
Merged
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
feat: Unit tests
RyanKadri committed Jan 16, 2025
commit d41cfab560dcc079964a5112e6b4e69cad80e7a1
12 changes: 10 additions & 2 deletions lib/llm-events/aws-bedrock/bedrock-command.js
Original file line number Diff line number Diff line change
@@ -88,13 +88,13 @@ class BedrockCommand {
}
]
} else if (
this.isClaude() === true ||
this.isClaudeTextCompletionApi() === true ||
this.isAi21() === true ||
this.isCohere() === true ||
this.isLlama() === true
) {
return [{ role: 'user', content: this.#body.prompt }]
} else if (this.isClaude3() === true) {
} else if (this.isClaudeMessagesApi() === true) {
return normalizeClaude3Messages(this.#body?.messages ?? [])
Copy link
Member

@bizob2828 bizob2828 Jan 16, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't see a unit nor versioned tests that asserts that this.#body.messages is falsey and defaults to empty array, is this needed? if so, please write a unit test

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was meant to handle the customer passing an invalid body. I realized the default value isn't needed here because I have it handled below anyway. I added a unit test for this case anyway to catch where it's handled lower down

}
return []
@@ -150,6 +150,14 @@ class BedrockCommand {
isTitanEmbed() {
return this.#modelId.startsWith('amazon.titan-embed')
}

isClaudeMessagesApi() {
return (this.isClaude3() === true || this.isClaude() === true) && 'messages' in this.#body
}

isClaudeTextCompletionApi() {
return this.isClaude() === true && 'prompt' in this.#body
}
}

/**
4 changes: 2 additions & 2 deletions lib/llm-events/aws-bedrock/embedding.js
Original file line number Diff line number Diff line change
@@ -10,7 +10,7 @@ const LlmEvent = require('./event')
/**
* @typedef {object} LlmEmbeddingParams
* @augments LlmEventParams
* @property
* @property {string} input - The input message for the embedding call
*/
/**
* @type {LlmEmbeddingParams}
@@ -24,7 +24,7 @@ class LlmEmbedding extends LlmEvent {
const tokenCb = agent?.llm?.tokenCountCallback

this.input = agent.config?.ai_monitoring?.record_content?.enabled
? this.bedrockCommand.prompt
? this.bedrockCommand.prompt?.[0]?.content
: undefined
this.error = params.isError
this.duration = params.segment.getDurationInMillis()
6 changes: 6 additions & 0 deletions test/lib/aws-server-stubs/ai-server/index.js
Original file line number Diff line number Diff line change
@@ -101,6 +101,12 @@ function handler(req, res) {
break
}

// Chunked claude model
case 'anthropic.claude-3-5-sonnet-20240620-v1:0': {
response = responses.claude3.get(payload?.messages?.[0]?.content?.[0].text)
break
}

case 'cohere.command-text-v14':
case 'cohere.command-light-text-v14': {
response = responses.cohere.get(payload.prompt)
34 changes: 34 additions & 0 deletions test/lib/aws-server-stubs/ai-server/responses/claude3.js
Original file line number Diff line number Diff line change
@@ -34,6 +34,40 @@ responses.set('text claude3 ultimate question', {
}
})

responses.set('text claude3 ultimate question chunked', {
headers: {
'content-type': contentType,
'x-amzn-requestid': reqId,
'x-amzn-bedrock-invocation-latency': '926',
'x-amzn-bedrock-output-token-count': '36',
'x-amzn-bedrock-input-token-count': '14'
},
statusCode: 200,
body: {
id: 'msg_bdrk_019V7ABaw8ZZZYuRDSTWK7VE',
type: 'message',
role: 'assistant',
model: 'claude-3-haiku-20240307',
stop_sequence: null,
usage: { input_tokens: 30, output_tokens: 265 },
content: [
{
type: 'text',
text: "Here's a nice picture of a 42"
},
{
type: 'image',
source: {
type: 'base64',
media_type: 'image/jpeg',
data: 'U2hoLiBUaGlzIGlzbid0IHJlYWxseSBhbiBpbWFnZQ=='
}
}
],
stop_reason: 'endoftext'
}
})

responses.set('text claude3 ultimate question streamed', {
headers: {
'content-type': 'application/vnd.amazon.eventstream',
61 changes: 61 additions & 0 deletions test/versioned/aws-sdk-v3/bedrock-chat-completions.test.js
Original file line number Diff line number Diff line change
@@ -50,6 +50,16 @@ const requests = {
}),
modelId
}),
claude3Chunked: (chunks, modelId) => ({
body: JSON.stringify({
anthropic_version: 'bedrock-2023-05-31',
max_tokens: 100,
temperature: 0.5,
system: 'Please respond in the style of Christopher Walken',
messages: chunks
}),
modelId
}),
cohere: (prompt, modelId) => ({
body: JSON.stringify({ prompt, temperature: 0.5, max_tokens: 100 }),
modelId
@@ -466,6 +476,57 @@ test('ai21: should properly create errors on create completion (streamed)', asyn
})
})

test('anthropic-claude-3: should properly create events for chunked messages', async (t) => {
const { bedrock, client, agent, expectedExternalPath } = t.nr
const modelId = 'anthropic.claude-3-5-sonnet-20240620-v1:0'
const prompt = 'text claude3 ultimate question chunked'
const input = requests.claude3Chunked(
[
{
role: 'user',
content: [
{
type: 'text',
text: prompt
}
]
}
],
modelId
)

const command = new bedrock.InvokeModelCommand(input)

const api = helper.getAgentApi()
await helper.runInTransaction(agent, async (tx) => {
api.addCustomAttribute('llm.conversation_id', 'convo-id')
await client.send(command)

assertSegments(
tx.trace.root,
['Llm/completion/Bedrock/InvokeModelCommand', [expectedExternalPath(modelId, 'invoke')]],
{ exact: false }
)

const events = agent.customEventAggregator.events.toArray()
assert.equal(events.length, 3)
const chatSummary = events.filter(([{ type }]) => type === 'LlmChatCompletionSummary')[0]
const chatMsgs = events.filter(([{ type }]) => type === 'LlmChatCompletionMessage')

// Note the <image> placeholder for the image chunk
assertChatCompletionMessages({
modelId,
prompt,
resContent: "Here's a nice picture of a 42\n\n<image>",
tx,
chatMsgs
})

assertChatCompletionSummary({ tx, modelId, chatSummary })
tx.end()
})
})

test('models that do not support streaming should be handled', async (t) => {
const { bedrock, client, agent, expectedExternalPath } = t.nr
const modelId = 'amazon.titan-embed-text-v1'