From a901a2499e73e0f1142054ead53e7b99df94b201 Mon Sep 17 00:00:00 2001
From: cmcadams-newrelic <157145722+cmcadams-newrelic@users.noreply.github.com>
Date: Tue, 12 Nov 2024 09:08:50 -0500
Subject: [PATCH] feat: Added support for Bedrock invoke model user request
 (#2721)

Co-authored-by: James Sumners <jsumners@newrelic.com>
---
 lib/llm-events/aws-bedrock/bedrock-command.js |  6 ++-
 .../aws-bedrock/bedrock-command.test.js       | 43 +++++++++++++++++++
 2 files changed, 48 insertions(+), 1 deletion(-)

diff --git a/lib/llm-events/aws-bedrock/bedrock-command.js b/lib/llm-events/aws-bedrock/bedrock-command.js
index 937b48c2c0..58ebc14a10 100644
--- a/lib/llm-events/aws-bedrock/bedrock-command.js
+++ b/lib/llm-events/aws-bedrock/bedrock-command.js
@@ -85,7 +85,11 @@ class BedrockCommand {
       result = this.#body.prompt
     } else if (this.isClaude3() === true) {
       result = this.#body?.messages?.reduce((acc, curr) => {
-        acc += curr?.content ?? ''
+        if (typeof curr?.content === 'string') {
+          acc += curr?.content
+        } else if (Object.keys(curr?.content).length) {
+          acc += curr?.content.text ?? ''
+        }
         return acc
       }, '')
     }
diff --git a/test/unit/llm-events/aws-bedrock/bedrock-command.test.js b/test/unit/llm-events/aws-bedrock/bedrock-command.test.js
index 639d11ef54..8097a4fad5 100644
--- a/test/unit/llm-events/aws-bedrock/bedrock-command.test.js
+++ b/test/unit/llm-events/aws-bedrock/bedrock-command.test.js
@@ -24,6 +24,13 @@ const claude = {
   }
 }
 
+const claude35 = {
+  modelId: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
+  body: {
+    messages: [{ role: 'user', content: { type: 'text', text: 'who are you' } }]
+  }
+}
+
 const claude3 = {
   modelId: 'anthropic.claude-3-haiku-20240307-v1:0',
   body: {
@@ -182,6 +189,42 @@ test('claude3 complete command works', async (t) => {
   assert.equal(cmd.temperature, payload.body.temperature)
 })
 
+test('claude35 minimal command works with claude 3 api', async (t) => {
+  t.nr.updatePayload(structuredClone(claude3))
+  const cmd = new BedrockCommand(t.nr.input)
+  assert.equal(cmd.isClaude3(), true)
+  assert.equal(cmd.maxTokens, undefined)
+  assert.equal(cmd.modelId, claude3.modelId)
+  assert.equal(cmd.modelType, 'completion')
+  assert.equal(cmd.prompt, claude3.body.messages[0].content)
+  assert.equal(cmd.temperature, undefined)
+})
+
+test('claude35 minimal command works', async (t) => {
+  t.nr.updatePayload(structuredClone(claude35))
+  const cmd = new BedrockCommand(t.nr.input)
+  assert.equal(cmd.isClaude3(), true)
+  assert.equal(cmd.maxTokens, undefined)
+  assert.equal(cmd.modelId, claude35.modelId)
+  assert.equal(cmd.modelType, 'completion')
+  assert.equal(cmd.prompt, claude35.body.messages[0].content.text)
+  assert.equal(cmd.temperature, undefined)
+})
+
+test('claude35 complete command works', async (t) => {
+  const payload = structuredClone(claude35)
+  payload.body.max_tokens = 25
+  payload.body.temperature = 0.5
+  t.nr.updatePayload(payload)
+  const cmd = new BedrockCommand(t.nr.input)
+  assert.equal(cmd.isClaude3(), true)
+  assert.equal(cmd.maxTokens, 25)
+  assert.equal(cmd.modelId, payload.modelId)
+  assert.equal(cmd.modelType, 'completion')
+  assert.equal(cmd.prompt, payload.body.messages[0].content.text)
+  assert.equal(cmd.temperature, payload.body.temperature)
+})
+
 test('cohere minimal command works', async (t) => {
   t.nr.updatePayload(structuredClone(cohere))
   const cmd = new BedrockCommand(t.nr.input)