Skip to content

Commit

Permalink
chore: Fixed aws-sdk-v3 bedrock tests (again)
Browse files Browse the repository at this point in the history
  • Loading branch information
jsumners-nr committed May 23, 2024
1 parent 050469e commit 8cef266
Showing 1 changed file with 17 additions and 21 deletions.
38 changes: 17 additions & 21 deletions test/versioned/aws-sdk-v3/bedrock-chat-completions.tap.js
Original file line number Diff line number Diff line change
Expand Up @@ -13,21 +13,6 @@ const createAiResponseServer = require('../../lib/aws-server-stubs/ai-server')
const { FAKE_CREDENTIALS } = require('../../lib/aws-server-stubs')
const { DESTINATIONS } = require('../../../lib/config/attribute-filter')

const pkgVersion = (function () {
try {
const { version } = require('@smithy/smithy-client/package.json')
return version
} catch {
try {
const {
version
} = require('./node_modules/@aws-sdk/client-bedrock-runtime/node_modules/@smithy/smithy-client/package.json')
return version
} catch {}
}
/* eslint-disable-next-line */
}())

const requests = {
ai21: (prompt, modelId) => ({
body: JSON.stringify({ prompt, temperature: 0.5, maxTokens: 100 }),
Expand Down Expand Up @@ -241,9 +226,10 @@ tap.afterEach(async (t) => {
const command = new bedrock.InvokeModelCommand(input)
helper.runInTransaction(agent, async (tx) => {
await client.send(command)
const metrics = agent.metrics.getOrCreateMetric(
`Supportability/Nodejs/ML/Bedrock/${pkgVersion}`
)
const metrics = getPrefixedMetric({
agent,
metricPrefix: 'Supportability/Nodejs/ML/Bedrock'
})
t.equal(metrics.callCount > 0, true)
tx.end()
t.end()
Expand Down Expand Up @@ -588,9 +574,10 @@ tap.test('should not instrument stream when disabled', (t) => {
t.equal(events.length, 0, 'should not create Llm events when streaming is disabled')
const attributes = tx.trace.attributes.get(DESTINATIONS.TRANS_EVENT)
t.equal(attributes.llm, true, 'should assign llm attribute to transaction trace')
const metrics = agent.metrics.getOrCreateMetric(
`Supportability/Nodejs/ML/Bedrock/${pkgVersion}`
)
const metrics = getPrefixedMetric({
agent,
metricPrefix: 'Supportability/Nodejs/ML/Bedrock'
})
t.equal(metrics.callCount > 0, true, 'should set framework metric')
const supportabilityMetrics = agent.metrics.getOrCreateMetric(
`Supportability/Nodejs/ML/Streaming/Disabled`
Expand Down Expand Up @@ -632,3 +619,12 @@ tap.test('should utilize tokenCountCallback when set', (t) => {
t.end()
})
})

function getPrefixedMetric({ agent, metricPrefix }) {
for (const [key, value] of Object.entries(agent.metrics._metrics.unscoped)) {
if (key.startsWith(metricPrefix) === false) {
continue
}
return value
}
}

0 comments on commit 8cef266

Please sign in to comment.