Skip to content
This repository has been archived by the owner on Jun 5, 2024. It is now read-only.

chore: Changed token_count to only use tokenCountCallback #272

Merged
merged 4 commits into from
Mar 11, 2024
Merged

chore: Changed token_count to only use tokenCountCallback #272

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cfafeb2
chore: Changed token_count to only use tokenCountCallback
jsumners-nr Mar 8, 2024
88bfce8
address feedback
jsumners-nr Mar 11, 2024
7119c2f
address feedback
jsumners-nr Mar 11, 2024
e61bd2f
address feedback
jsumners-nr Mar 11, 2024
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
27 changes: 0 additions & 27 deletions lib/llm/bedrock-response.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,25 +122,6 @@ class BedrockResponse {
return this.#id
}

/**
* The number of tokens present in the prompt as determined by the remote
* API.
*
* @returns {number|undefined}
*/
get inputTokenCount() {
bizob2828 marked this conversation as resolved.
Show resolved Hide resolved
return this.#tokenCount('x-amzn-bedrock-input-token-count')
}

/**
* The number of tokens in the LLM response as determined by the remote API.
*
* @returns {number|undefined}
*/
get outputTokenCount() {
return this.#tokenCount('x-amzn-bedrock-output-token-count')
}

/**
* UUID assigned to the initial request as returned by the API.
*
Expand All @@ -158,14 +139,6 @@ class BedrockResponse {
get statusCode() {
return this.#innerResponse.statusCode
}

#tokenCount(headerName) {
const headerVal = this.headers?.[headerName]
if (headerVal != null) {
return parseInt(headerVal, 10)
}
return undefined
}
}

module.exports = BedrockResponse
6 changes: 2 additions & 4 deletions lib/llm/chat-completion-message.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,15 +52,13 @@ class LlmChatCompletionMessage extends LlmEvent {
this.#setId(index)
if (this.is_response === true) {
this.role = 'assistant'
this.token_count = this.bedrockResponse.outputTokenCount
if (this.token_count === undefined && typeof tokenCB === 'function') {
if (typeof tokenCB === 'function') {
this.token_count = tokenCB(this.bedrockCommand.modelId, content)
}
} else {
this.role = 'user'
this.content = recordContent === true ? this.bedrockCommand.prompt : undefined
this.token_count = this.bedrockResponse.inputTokenCount
if (this.token_count === undefined && typeof tokenCB === 'function') {
if (typeof tokenCB === 'function') {
this.token_count = tokenCB(this.bedrockCommand.modelId, this.bedrockCommand.prompt)
}
}
Expand Down
22 changes: 0 additions & 22 deletions tests/unit/llm/bedrock-response.tap.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,8 +98,6 @@ tap.test('non-conforming response is handled gracefully', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, undefined)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, undefined)
Copy link
Member

Choose a reason for hiding this comment

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

this should be restored

Copy link
Member

Choose a reason for hiding this comment

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

i meant the entire file, sorry for the confusion

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Why? What are we trying to test by doing so?

Copy link
Member

Choose a reason for hiding this comment

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

i understand this code is going away in #265. so it's probably fine to leave as is but the test coverage will be missing until we rip it all out.

t.equal(res.outputTokenCount, undefined)
t.equal(res.requestId, undefined)
t.equal(res.statusCode, 200)
})
Expand All @@ -111,8 +109,6 @@ tap.test('ai21 malformed responses work', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -125,8 +121,6 @@ tap.test('ai21 complete responses work', async (t) => {
t.equal(res.finishReason, 'done')
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, 'ai21-response-1')
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -138,8 +132,6 @@ tap.test('claude malformed responses work', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -152,8 +144,6 @@ tap.test('claude complete responses work', async (t) => {
t.equal(res.finishReason, 'done')
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -165,8 +155,6 @@ tap.test('cohere malformed responses work', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -179,8 +167,6 @@ tap.test('cohere complete responses work', async (t) => {
t.equal(res.finishReason, 'done')
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, 'cohere-response-1')
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -192,8 +178,6 @@ tap.test('llama2 malformed responses work', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -206,8 +190,6 @@ tap.test('llama2 complete responses work', async (t) => {
t.equal(res.finishReason, 'done')
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -219,8 +201,6 @@ tap.test('titan malformed responses work', async (t) => {
t.equal(res.finishReason, undefined)
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand All @@ -233,8 +213,6 @@ tap.test('titan complete responses work', async (t) => {
t.equal(res.finishReason, 'done')
t.same(res.headers, t.context.response.response.headers)
t.equal(res.id, undefined)
t.equal(res.inputTokenCount, 25)
t.equal(res.outputTokenCount, 25)
t.equal(res.requestId, 'aws-request-1')
t.equal(res.statusCode, 200)
})
Expand Down
4 changes: 0 additions & 4 deletions tests/unit/llm/chat-completion-message.tap.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,6 @@ tap.test('create creates a non-response instance', async (t) => {
t.equal(event.content, 'who are you')
t.equal(event.role, 'user')
t.match(event.id, /[\w-]{36}/)
t.equal(event.token_count, 8)
})

tap.test('create creates a titan response instance', async (t) => {
Expand All @@ -108,7 +107,6 @@ tap.test('create creates a titan response instance', async (t) => {
t.equal(event.content, 'a response')
t.equal(event.role, 'assistant')
t.match(event.id, /[\w-]{36}-0/)
t.equal(event.token_count, 4)
})

tap.test('create creates a cohere response instance', async (t) => {
Expand All @@ -124,7 +122,6 @@ tap.test('create creates a cohere response instance', async (t) => {
t.equal(event.content, 'a response')
t.equal(event.role, 'assistant')
t.match(event.id, /42-0/)
t.equal(event.token_count, 4)
})

tap.test('create creates a ai21 response instance when response.id is undefined', async (t) => {
Expand All @@ -140,7 +137,6 @@ tap.test('create creates a ai21 response instance when response.id is undefined'
t.equal(event.content, 'a response')
t.equal(event.role, 'assistant')
t.match(event.id, /[\w-]{36}-0/)
t.equal(event.token_count, 4)
})

tap.test(
Expand Down
10 changes: 0 additions & 10 deletions tests/unit/llm/stream-handler.tap.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,8 +117,6 @@ tap.test('handles claude streams', async (t) => {
t.equal(br.finishReason, 'done')
t.equal(br.requestId, 'aws-req-1')
t.equal(br.statusCode, 200)
t.equal(br.inputTokenCount, 5)
Copy link
Member

Choose a reason for hiding this comment

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

this should be restored

Copy link
Member

Choose a reason for hiding this comment

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

i meant the entire file, sorry for the confusion

t.equal(br.outputTokenCount, 10)
})

tap.test('handles cohere streams', async (t) => {
Expand Down Expand Up @@ -166,8 +164,6 @@ tap.test('handles cohere streams', async (t) => {
t.equal(br.finishReason, 'done')
t.equal(br.requestId, 'aws-req-1')
t.equal(br.statusCode, 200)
t.equal(br.inputTokenCount, 5)
t.equal(br.outputTokenCount, 10)
})

tap.test('handles cohere embedding streams', async (t) => {
Expand Down Expand Up @@ -220,8 +216,6 @@ tap.test('handles cohere embedding streams', async (t) => {
t.equal(br.finishReason, undefined)
t.equal(br.requestId, 'aws-req-1')
t.equal(br.statusCode, 200)
t.equal(br.inputTokenCount, 5)
t.equal(br.outputTokenCount, 10)
})

tap.test('handles llama2 streams', async (t) => {
Expand Down Expand Up @@ -262,8 +256,6 @@ tap.test('handles llama2 streams', async (t) => {
t.equal(br.finishReason, 'done')
t.equal(br.requestId, 'aws-req-1')
t.equal(br.statusCode, 200)
t.equal(br.inputTokenCount, 5)
t.equal(br.outputTokenCount, 10)
})

tap.test('handles titan streams', async (t) => {
Expand Down Expand Up @@ -314,6 +306,4 @@ tap.test('handles titan streams', async (t) => {
t.equal(br.finishReason, 'done')
t.equal(br.requestId, 'aws-req-1')
t.equal(br.statusCode, 200)
t.equal(br.inputTokenCount, 5)
t.equal(br.outputTokenCount, 10)
})
20 changes: 2 additions & 18 deletions tests/versioned/common.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -118,15 +118,8 @@ function assertChatCompletionMessages({ tx, chatMsgs, expectedId, modelId, promp
})
}

function assertChatCompletionSummary({
tx,
modelId,
chatSummary,
tokenUsage,
error = false,
numMsgs = 2
}) {
let expectedChatSummary = {
function assertChatCompletionSummary({ tx, modelId, chatSummary, error = false, numMsgs = 2 }) {
const expectedChatSummary = {
'id': /[\w]{8}-[\w]{4}-[\w]{4}-[\w]{4}-[\w]{12}/,
'appName': 'New Relic for Node.js tests',
'request_id': 'eda0760a-c3f0-4fc1-9a1e-75559d642866',
Expand All @@ -146,15 +139,6 @@ function assertChatCompletionSummary({
'error': error
}

if (tokenUsage) {
jsumners-nr marked this conversation as resolved.
Show resolved Hide resolved
expectedChatSummary = {
...expectedChatSummary,
'response.usage.total_tokens': 12,
'response.usage.prompt_tokens': 8,
'response.usage.completion_tokens': 4
}
}

this.equal(chatSummary[0].type, 'LlmChatCompletionSummary')
this.match(chatSummary[1], expectedChatSummary, 'should match chat summary message')
}
Expand Down
2 changes: 1 addition & 1 deletion tests/versioned/v3/bedrock-chat-completions.tap.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ tap.afterEach(async (t) => {
chatMsgs
})

t.llmSummary({ tx, modelId, chatSummary, tokenUsage: true, numMsgs: events.length - 1 })
t.llmSummary({ tx, modelId, chatSummary, numMsgs: events.length - 1 })
Copy link
Member

Choose a reason for hiding this comment

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

this should be restored


tx.end()
t.end()
Expand Down
6 changes: 3 additions & 3 deletions tests/versioned/v3/bedrock-embeddings.tap.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,9 +117,9 @@ tap.afterEach(async (t) => {
'ingest_source': 'Node',
'request.model': modelId,
'duration': tx.trace.root.children[0].getDurationInMillis(),
'response.usage.total_tokens': 13,
'response.usage.prompt_tokens': 13,
'token_count': 13,
'response.usage.total_tokens': undefined,
'response.usage.prompt_tokens': undefined,
'token_count': undefined,
'input': prompt,
'error': false
}
Expand Down
Loading