BedrockAnthropicTokenUsageCallbackHandler does not function with langchain_aws

[austinmw]

langchain_community.callbacks.bedrock_anthropic_callback.BedrockAnthropicTokenUsageCallbackHandler appears to be broken with langchain_aws models.

Example:

# Latest versions
#%pip install -U langchain_community==0.2.0 langchain_core==0.2.0 langchain_aws==0.1.4

from langchain_community.callbacks.bedrock_anthropic_callback import BedrockAnthropicTokenUsageCallbackHandler
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_aws import ChatBedrock

region = "us-east-1"

model = ChatBedrock(
    model_id="anthropic.claude-3-sonnet-20240229-v1:0",
    region_name="us-east-1",
)

# Create an instance of the callback handler
token_usage_callback = BedrockAnthropicTokenUsageCallbackHandler()

# Pass the callback handler to the underlying LLM model
model.callbacks = [token_usage_callback]

prompt = PromptTemplate(
    template="List 5 colors",
    input_variables=[],
)

# Create an instance of the callback handler
token_usage_callback = BedrockAnthropicTokenUsageCallbackHandler()

# Pass the callback handler to the underlying LLM model
model.callbacks = [token_usage_callback]

# Create the processing chain
chain = prompt | model | StrOutputParser()

response = chain.invoke({})
print(response)

print(token_usage_callback)

Output:

Here are 5 colors:

1. Red
2. Blue
3. Yellow
4. Green
5. Purple
Tokens Used: 0
    Prompt Tokens: 0
    Completion Tokens: 0
Successful Requests: 1
Total Cost (USD): $0.0

Side note: It might be a good idea for langchain_aws to take over the BedrockAnthropicTokenUsageCallbackHandler callback instead in order to keep it more up to date?

[3coins]

@austinmw PR 20 added support for token count with callbacks. It might not be compatible with the BedrockAnthropicTokenUsageCallbackHandler but you can use the sample in the PR to get the token count. Having a dedicated callback handler does add some utility, plz feel free to add one in this repo if you have cycles to do so.