From cfafeb2e2df5823ea8d4960e6b0c565d2e12a00a Mon Sep 17 00:00:00 2001 From: James Sumners Date: Fri, 8 Mar 2024 09:51:57 -0500 Subject: [PATCH 1/4] chore: Changed token_count to only use tokenCountCallback --- lib/llm/bedrock-response.js | 27 ------------------- lib/llm/chat-completion-message.js | 6 ++--- tests/unit/llm/bedrock-response.tap.js | 22 --------------- tests/unit/llm/chat-completion-message.tap.js | 4 --- tests/unit/llm/stream-handler.tap.js | 10 ------- tests/versioned/common.js | 20 ++------------ .../v3/bedrock-chat-completions.tap.js | 2 +- tests/versioned/v3/bedrock-embeddings.tap.js | 6 ++--- 8 files changed, 8 insertions(+), 89 deletions(-) diff --git a/lib/llm/bedrock-response.js b/lib/llm/bedrock-response.js index b79b24f..06e24f5 100644 --- a/lib/llm/bedrock-response.js +++ b/lib/llm/bedrock-response.js @@ -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() { - 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. * @@ -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 diff --git a/lib/llm/chat-completion-message.js b/lib/llm/chat-completion-message.js index 0cc733d..6f0c564 100644 --- a/lib/llm/chat-completion-message.js +++ b/lib/llm/chat-completion-message.js @@ -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) } } diff --git a/tests/unit/llm/bedrock-response.tap.js b/tests/unit/llm/bedrock-response.tap.js index d378690..4dc88ac 100644 --- a/tests/unit/llm/bedrock-response.tap.js +++ b/tests/unit/llm/bedrock-response.tap.js @@ -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) - t.equal(res.outputTokenCount, undefined) t.equal(res.requestId, undefined) t.equal(res.statusCode, 200) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) @@ -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) }) diff --git a/tests/unit/llm/chat-completion-message.tap.js b/tests/unit/llm/chat-completion-message.tap.js index 3f06689..d065f5a 100644 --- a/tests/unit/llm/chat-completion-message.tap.js +++ b/tests/unit/llm/chat-completion-message.tap.js @@ -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) => { @@ -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) => { @@ -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) => { @@ -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( diff --git a/tests/unit/llm/stream-handler.tap.js b/tests/unit/llm/stream-handler.tap.js index beffcf2..3701798 100644 --- a/tests/unit/llm/stream-handler.tap.js +++ b/tests/unit/llm/stream-handler.tap.js @@ -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) - t.equal(br.outputTokenCount, 10) }) tap.test('handles cohere streams', async (t) => { @@ -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) => { @@ -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) => { @@ -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) => { @@ -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) }) diff --git a/tests/versioned/common.js b/tests/versioned/common.js index 7a0d2b1..d81b90e 100644 --- a/tests/versioned/common.js +++ b/tests/versioned/common.js @@ -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', @@ -146,15 +139,6 @@ function assertChatCompletionSummary({ 'error': error } - if (tokenUsage) { - 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') } diff --git a/tests/versioned/v3/bedrock-chat-completions.tap.js b/tests/versioned/v3/bedrock-chat-completions.tap.js index a40bc03..8322145 100644 --- a/tests/versioned/v3/bedrock-chat-completions.tap.js +++ b/tests/versioned/v3/bedrock-chat-completions.tap.js @@ -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 }) tx.end() t.end() diff --git a/tests/versioned/v3/bedrock-embeddings.tap.js b/tests/versioned/v3/bedrock-embeddings.tap.js index 514b2cf..c724590 100644 --- a/tests/versioned/v3/bedrock-embeddings.tap.js +++ b/tests/versioned/v3/bedrock-embeddings.tap.js @@ -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 } From 88bfce8419bd0827a66e02394708ab8319682f00 Mon Sep 17 00:00:00 2001 From: James Sumners Date: Mon, 11 Mar 2024 10:44:07 -0400 Subject: [PATCH 2/4] address feedback --- lib/llm/bedrock-response.js | 27 ++++++++++++++++++++ tests/versioned/common.js | 20 +++++++++++++-- tests/versioned/v3/bedrock-embeddings.tap.js | 6 ++--- 3 files changed, 48 insertions(+), 5 deletions(-) diff --git a/lib/llm/bedrock-response.js b/lib/llm/bedrock-response.js index 06e24f5..b79b24f 100644 --- a/lib/llm/bedrock-response.js +++ b/lib/llm/bedrock-response.js @@ -122,6 +122,25 @@ class BedrockResponse { return this.#id } + /** + * The number of tokens present in the prompt as determined by the remote + * API. + * + * @returns {number|undefined} + */ + get inputTokenCount() { + 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. * @@ -139,6 +158,14 @@ 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 diff --git a/tests/versioned/common.js b/tests/versioned/common.js index d81b90e..7a0d2b1 100644 --- a/tests/versioned/common.js +++ b/tests/versioned/common.js @@ -118,8 +118,15 @@ function assertChatCompletionMessages({ tx, chatMsgs, expectedId, modelId, promp }) } -function assertChatCompletionSummary({ tx, modelId, chatSummary, error = false, numMsgs = 2 }) { - const expectedChatSummary = { +function assertChatCompletionSummary({ + tx, + modelId, + chatSummary, + tokenUsage, + error = false, + numMsgs = 2 +}) { + let 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', @@ -139,6 +146,15 @@ function assertChatCompletionSummary({ tx, modelId, chatSummary, error = false, 'error': error } + if (tokenUsage) { + 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') } diff --git a/tests/versioned/v3/bedrock-embeddings.tap.js b/tests/versioned/v3/bedrock-embeddings.tap.js index c724590..514b2cf 100644 --- a/tests/versioned/v3/bedrock-embeddings.tap.js +++ b/tests/versioned/v3/bedrock-embeddings.tap.js @@ -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': undefined, - 'response.usage.prompt_tokens': undefined, - 'token_count': undefined, + 'response.usage.total_tokens': 13, + 'response.usage.prompt_tokens': 13, + 'token_count': 13, 'input': prompt, 'error': false } From 7119c2fcf74bafa0481ce7f28d47646700aad0c7 Mon Sep 17 00:00:00 2001 From: James Sumners Date: Mon, 11 Mar 2024 12:44:04 -0400 Subject: [PATCH 3/4] address feedback --- tests/unit/llm/bedrock-response.tap.js | 1 + tests/unit/llm/stream-handler.tap.js | 1 + tests/versioned/v3/bedrock-chat-completions.tap.js | 2 +- 3 files changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit/llm/bedrock-response.tap.js b/tests/unit/llm/bedrock-response.tap.js index 4dc88ac..999cd77 100644 --- a/tests/unit/llm/bedrock-response.tap.js +++ b/tests/unit/llm/bedrock-response.tap.js @@ -98,6 +98,7 @@ 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) t.equal(res.requestId, undefined) t.equal(res.statusCode, 200) }) diff --git a/tests/unit/llm/stream-handler.tap.js b/tests/unit/llm/stream-handler.tap.js index 3701798..3044499 100644 --- a/tests/unit/llm/stream-handler.tap.js +++ b/tests/unit/llm/stream-handler.tap.js @@ -117,6 +117,7 @@ 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) }) tap.test('handles cohere streams', async (t) => { diff --git a/tests/versioned/v3/bedrock-chat-completions.tap.js b/tests/versioned/v3/bedrock-chat-completions.tap.js index 8322145..a40bc03 100644 --- a/tests/versioned/v3/bedrock-chat-completions.tap.js +++ b/tests/versioned/v3/bedrock-chat-completions.tap.js @@ -189,7 +189,7 @@ tap.afterEach(async (t) => { chatMsgs }) - t.llmSummary({ tx, modelId, chatSummary, numMsgs: events.length - 1 }) + t.llmSummary({ tx, modelId, chatSummary, tokenUsage: true, numMsgs: events.length - 1 }) tx.end() t.end() From e61bd2fb6d537991e1a24e8f581ee5d498208fd5 Mon Sep 17 00:00:00 2001 From: James Sumners Date: Mon, 11 Mar 2024 13:47:48 -0400 Subject: [PATCH 4/4] address feedback --- tests/unit/llm/bedrock-response.tap.js | 21 +++++++++++++++++++++ tests/unit/llm/stream-handler.tap.js | 9 +++++++++ 2 files changed, 30 insertions(+) diff --git a/tests/unit/llm/bedrock-response.tap.js b/tests/unit/llm/bedrock-response.tap.js index 999cd77..d378690 100644 --- a/tests/unit/llm/bedrock-response.tap.js +++ b/tests/unit/llm/bedrock-response.tap.js @@ -99,6 +99,7 @@ tap.test('non-conforming response is handled gracefully', async (t) => { t.same(res.headers, undefined) t.equal(res.id, undefined) t.equal(res.inputTokenCount, undefined) + t.equal(res.outputTokenCount, undefined) t.equal(res.requestId, undefined) t.equal(res.statusCode, 200) }) @@ -110,6 +111,8 @@ 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) }) @@ -122,6 +125,8 @@ 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) }) @@ -133,6 +138,8 @@ 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) }) @@ -145,6 +152,8 @@ 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) }) @@ -156,6 +165,8 @@ 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) }) @@ -168,6 +179,8 @@ 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) }) @@ -179,6 +192,8 @@ 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) }) @@ -191,6 +206,8 @@ 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) }) @@ -202,6 +219,8 @@ 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) }) @@ -214,6 +233,8 @@ 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) }) diff --git a/tests/unit/llm/stream-handler.tap.js b/tests/unit/llm/stream-handler.tap.js index 3044499..beffcf2 100644 --- a/tests/unit/llm/stream-handler.tap.js +++ b/tests/unit/llm/stream-handler.tap.js @@ -118,6 +118,7 @@ tap.test('handles claude streams', async (t) => { 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 streams', async (t) => { @@ -165,6 +166,8 @@ 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) => { @@ -217,6 +220,8 @@ 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) => { @@ -257,6 +262,8 @@ 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) => { @@ -307,4 +314,6 @@ 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) })