feat: Updated token calculation for aws-bedrock LLM Events

[svetlanabrennan]

Description

This PR updates how we calculate token_count and collect token usage attributes for openai.

See spec details here.

Token Calculation

Customers can choose to provide an llm_token_count_callback which agents MUST invoke to obtain token counts. If no callback is registered, then agents MUST pull token count data from the usage metadata objects returned in the LLM response.

• token_count should be set on LlmChatCompletionMessage event.
• token_count: set to equal zero and attach to the message event as an attribute IF (in this order):
  • A. the user set a tokenCallback and it's valid (not null and greater than 0)
  • B. there are token usage metadata in the response

Otherwise omit this attribute completely. This signifies the pipeline tokenizer whether or not it should run (this is done for backwards compatibility).

• See "response.usage.* attributes: " section below for details on the token usage data we look for.
• Record token_count attribute even if ai_monitoring.record_content.enabled is false.

response.usage.* attributes:

The relevant token usage attributes MUST be attached to the LlmChatCompletionSummary and LlmEmbedding events.

These will be set by obtaining the token values from the tokenCallback or from the token usage metadata in the response object or response headers (in this order).

• The LlmChatCompletionSummary event should add response.usage.prompt_tokens, response.usage.completion_tokens and response.usage.total_tokens

  • If the user set a tokenCallback, use that to obtain the prompt token count and set that value as response.usage.prompt_tokens on the summary event.
  • Then do that again for obtaining the completion token and set that value as response.usage.completion_tokens on the summary event.
  • Then add up both of those values to get the total token count and set response.usage.total_tokens on the summary event.
  • If prompt tokens and completion tokens are not obtained from the tokenCallback or the response object or the response headers, don’t set them at all.

• The LlmEmbedding event should only add response.usage.total_tokens.

  • use the tokenCallback to get the input tokens and set response.usage.total_tokens
  • If there was no tokenCallback registered, obtain total tokens from the response object.
  • For the embedding event, if total tokens is not obtained from the tokenCallback or the response object or the response headers, don’t set it at all.

Additional changes

I updated the stream handlers for Claude3 and Llama in order to get the token usage (input and output) from the stream and add to the currentBody before returning the response. The other handlers like Titan already has the token usage in the parsed event so no additional logic had to be added there.

How to Test

Run npm run test.

Related Issues

Partially fixes issue #3042

[codecov]

Codecov Report

❌ Patch coverage is 92.06349% with 15 lines in your changes missing coverage. Please review.
✅ Project coverage is 89.54%. Comparing base (b8678cc) to head (9b84934).
⚠️ Report is 13 commits behind head on main.

Files with missing lines	Patch %	Lines
lib/llm-events/aws-bedrock/embedding.js	64.86%	13 Missing ⚠️
lib/llm-events/aws-bedrock/stream-handler.js	90.90%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (b8678cc) and HEAD (9b84934). Click for more details.

HEAD has 4 uploads less than BASE

Flag	BASE (b8678cc)	HEAD (9b84934)
unit-tests-20.x	1	0
unit-tests-22.x	1	0
unit-tests-24.x	1	0
integration-tests-cjs-24.x	1	0

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3445      +/-   ##
==========================================
- Coverage   97.75%   89.54%   -8.22%     
==========================================
  Files         404      413       +9     
  Lines       54575    55474     +899     
  Branches        1        1              
==========================================
- Hits        53348    49672    -3676     
- Misses       1227     5802    +4575     

Flag	Coverage Δ
integration-tests-cjs-20.x	73.32% <14.81%> (-0.63%)	⬇️
integration-tests-cjs-22.x	73.36% <14.81%> (-0.61%)	⬇️
integration-tests-cjs-24.x	?
integration-tests-esm-20.x	51.76% <14.81%> (+0.01%)	⬆️
integration-tests-esm-22.x	51.81% <14.81%> (+0.01%)	⬆️
integration-tests-esm-24.x	53.02% <18.51%> (+<0.01%)	⬆️
unit-tests-20.x	?
unit-tests-22.x	?
unit-tests-24.x	?
versioned-tests-20.x	81.09% <92.06%> (-0.13%)	⬇️
versioned-tests-22.x	81.10% <92.06%> (-0.13%)	⬇️
versioned-tests-24.x	81.03% <92.06%> (-0.14%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.