Skip to content

Commit

Permalink
test: Converted llm-events tests to use node:test (newrelic#2535)
Browse files Browse the repository at this point in the history
  • Loading branch information
@amychisholm03
amychisholm03 authored Sep 4, 2024
1 parent 762511b commit ebfa2e9
Show file tree
Hide file tree
Showing 19 changed files with 1,168 additions and 1,146 deletions.
318 changes: 160 additions & 158 deletions test/unit/llm-events/aws-bedrock/bedrock-command.test.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,8 @@

'use strict'

const tap = require('tap')
const test = require('node:test')
const assert = require('node:assert')
const structuredClone = require('./clone')
const BedrockCommand = require('../../../../lib/llm-events/aws-bedrock/bedrock-command')

Expand Down Expand Up @@ -73,19 +74,20 @@ const titanEmbed = {
}
}

tap.beforeEach((t) => {
t.context.input = {
test.beforeEach((ctx) => {
ctx.nr = {}
ctx.nr.input = {
body: JSON.stringify('{"foo":"foo"}')
}

t.context.updatePayload = (payload) => {
t.context.input.modelId = payload.modelId
t.context.input.body = JSON.stringify(payload.body)
ctx.nr.updatePayload = (payload) => {
ctx.nr.input.modelId = payload.modelId
ctx.nr.input.body = JSON.stringify(payload.body)
}
})

tap.test('non-conforming command is handled gracefully', async (t) => {
const cmd = new BedrockCommand(t.context.input)
test('non-conforming command is handled gracefully', async (t) => {
const cmd = new BedrockCommand(t.nr.input)
for (const model of [
'Ai21',
'Claude',
Expand All @@ -96,210 +98,210 @@ tap.test('non-conforming command is handled gracefully', async (t) => {
'Titan',
'TitanEmbed'
]) {
t.equal(cmd[`is${model}`](), false)
assert.equal(cmd[`is${model}`](), false)
}
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, '')
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, undefined)
t.equal(cmd.temperature, undefined)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, '')
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, undefined)
assert.equal(cmd.temperature, undefined)
})

tap.test('ai21 minimal command works', async (t) => {
t.context.updatePayload(structuredClone(ai21))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isAi21(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, ai21.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, ai21.body.prompt)
t.equal(cmd.temperature, undefined)
test('ai21 minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(ai21))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isAi21(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, ai21.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, ai21.body.prompt)
assert.equal(cmd.temperature, undefined)
})

tap.test('ai21 complete command works', async (t) => {
test('ai21 complete command works', async (t) => {
const payload = structuredClone(ai21)
payload.body.maxTokens = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isAi21(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.prompt)
t.equal(cmd.temperature, payload.body.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isAi21(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.prompt)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('claude minimal command works', async (t) => {
t.context.updatePayload(structuredClone(claude))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isClaude(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, claude.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, claude.body.prompt)
t.equal(cmd.temperature, undefined)
test('claude minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(claude))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isClaude(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, claude.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, claude.body.prompt)
assert.equal(cmd.temperature, undefined)
})

tap.test('claude complete command works', async (t) => {
test('claude complete command works', async (t) => {
const payload = structuredClone(claude)
payload.body.max_tokens_to_sample = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isClaude(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.prompt)
t.equal(cmd.temperature, payload.body.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isClaude(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.prompt)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('claude3 minimal command works', async (t) => {
t.context.updatePayload(structuredClone(claude3))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isClaude3(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, claude3.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, claude3.body.messages[0].content)
t.equal(cmd.temperature, undefined)
test('claude3 minimal command works', 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)
})

tap.test('claude3 complete command works', async (t) => {
test('claude3 complete command works', async (t) => {
const payload = structuredClone(claude3)
payload.body.max_tokens = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isClaude3(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.messages[0].content)
t.equal(cmd.temperature, payload.body.temperature)
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)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('cohere minimal command works', async (t) => {
t.context.updatePayload(structuredClone(cohere))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isCohere(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, cohere.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, cohere.body.prompt)
t.equal(cmd.temperature, undefined)
test('cohere minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(cohere))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isCohere(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, cohere.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, cohere.body.prompt)
assert.equal(cmd.temperature, undefined)
})

tap.test('cohere complete command works', async (t) => {
test('cohere complete command works', async (t) => {
const payload = structuredClone(cohere)
payload.body.max_tokens = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isCohere(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.prompt)
t.equal(cmd.temperature, payload.body.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isCohere(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.prompt)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('cohere embed minimal command works', async (t) => {
t.context.updatePayload(structuredClone(cohereEmbed))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isCohereEmbed(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, cohereEmbed.modelId)
t.equal(cmd.modelType, 'embedding')
t.same(cmd.prompt, cohereEmbed.body.texts.join(' '))
t.equal(cmd.temperature, undefined)
test('cohere embed minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(cohereEmbed))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isCohereEmbed(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, cohereEmbed.modelId)
assert.equal(cmd.modelType, 'embedding')
assert.deepStrictEqual(cmd.prompt, cohereEmbed.body.texts.join(' '))
assert.equal(cmd.temperature, undefined)
})

tap.test('llama2 minimal command works', async (t) => {
t.context.updatePayload(structuredClone(llama2))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isLlama(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, llama2.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, llama2.body.prompt)
t.equal(cmd.temperature, undefined)
test('llama2 minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(llama2))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isLlama(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, llama2.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, llama2.body.prompt)
assert.equal(cmd.temperature, undefined)
})

tap.test('llama2 complete command works', async (t) => {
test('llama2 complete command works', async (t) => {
const payload = structuredClone(llama2)
payload.body.max_gen_length = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isLlama(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.prompt)
t.equal(cmd.temperature, payload.body.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isLlama(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.prompt)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('llama3 minimal command works', async (t) => {
t.context.updatePayload(structuredClone(llama3))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isLlama(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, llama3.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, llama3.body.prompt)
t.equal(cmd.temperature, undefined)
test('llama3 minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(llama3))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isLlama(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, llama3.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, llama3.body.prompt)
assert.equal(cmd.temperature, undefined)
})

tap.test('llama3 complete command works', async (t) => {
test('llama3 complete command works', async (t) => {
const payload = structuredClone(llama3)
payload.body.max_gen_length = 25
payload.body.temperature = 0.5
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isLlama(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.prompt)
t.equal(cmd.temperature, payload.body.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isLlama(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.prompt)
assert.equal(cmd.temperature, payload.body.temperature)
})

tap.test('titan minimal command works', async (t) => {
t.context.updatePayload(structuredClone(titan))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isTitan(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, titan.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, titan.body.inputText)
t.equal(cmd.temperature, undefined)
test('titan minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(titan))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isTitan(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, titan.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, titan.body.inputText)
assert.equal(cmd.temperature, undefined)
})

tap.test('titan complete command works', async (t) => {
test('titan complete command works', async (t) => {
const payload = structuredClone(titan)
payload.body.textGenerationConfig = {
maxTokenCount: 25,
temperature: 0.5
}
t.context.updatePayload(payload)
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isTitan(), true)
t.equal(cmd.maxTokens, 25)
t.equal(cmd.modelId, payload.modelId)
t.equal(cmd.modelType, 'completion')
t.equal(cmd.prompt, payload.body.inputText)
t.equal(cmd.temperature, payload.body.textGenerationConfig.temperature)
t.nr.updatePayload(payload)
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isTitan(), true)
assert.equal(cmd.maxTokens, 25)
assert.equal(cmd.modelId, payload.modelId)
assert.equal(cmd.modelType, 'completion')
assert.equal(cmd.prompt, payload.body.inputText)
assert.equal(cmd.temperature, payload.body.textGenerationConfig.temperature)
})

tap.test('titan embed minimal command works', async (t) => {
t.context.updatePayload(structuredClone(titanEmbed))
const cmd = new BedrockCommand(t.context.input)
t.equal(cmd.isTitanEmbed(), true)
t.equal(cmd.maxTokens, undefined)
t.equal(cmd.modelId, titanEmbed.modelId)
t.equal(cmd.modelType, 'embedding')
t.equal(cmd.prompt, titanEmbed.body.inputText)
t.equal(cmd.temperature, undefined)
test('titan embed minimal command works', async (t) => {
t.nr.updatePayload(structuredClone(titanEmbed))
const cmd = new BedrockCommand(t.nr.input)
assert.equal(cmd.isTitanEmbed(), true)
assert.equal(cmd.maxTokens, undefined)
assert.equal(cmd.modelId, titanEmbed.modelId)
assert.equal(cmd.modelType, 'embedding')
assert.equal(cmd.prompt, titanEmbed.body.inputText)
assert.equal(cmd.temperature, undefined)
})
Loading

0 comments on commit ebfa2e9

Please sign in to comment.