This PR is a fix for Issue 7657 where the usage details for ChatbedRock were not captured properly while it was working fine for ChatbedRockConverse.

I have mentioned the cause for the bug in the issue comment but I am pasting the gist of it here.

Cause of the Bug

The error for this seems to be at langchain/CallbackHandler.py due to mainly two reasons

• There is no key mapping for converting between the model key to the langfuse key
• The model keys for ChatBedrock is the exact same as the OpenAI keys which we are skipping since they are being handled server side (I do not know exactly where it is being handled server side and would love to get your help here)

The problem specifically is with the function _parse_usage_model.

When it comes to ChatBedrockConverse, the input passed to this function looks like the following

{'input_tokens': 11, 'output_tokens': 30, 'total_tokens': 41, 'input_token_details': {'cache_creation': 0, 'cache_read': 0}}

while with ChatBedrock, the input passed to this function looks like the following

{'prompt_tokens': 11, 'completion_tokens': 30, 'cache_read_input_tokens': 0, 'cache_write_input_tokens': 0, 'total_tokens': 41}

Fix for the Bug

• I've added the key mapping for converting between the model key to the langfuse key
• Only the input, output and total keys for ChatBedrock is the same as OpenAI so I've added the remaining keys in OpenAI for the code to be able to differentiate between OpenAI vs other providers

Here is a screenshot of the issue before fixing this

Here is a screenshot of the trace after fixing this issue

I tried running the existing tests but some of them are failing because they require me to provide an OpenAI API key but I don't think they have a bearing on this fix. But do let me know if the team requires me to do some additional testing beyond this.

Disclaimer: Experimental PR review

Greptile Summary

This PR addresses a bug in the LangChain callback handler where token usage metrics were not being properly captured for Amazon's ChatBedrock model. The issue arose because ChatBedrock uses similar token usage keys to OpenAI (prompt_tokens, completion_tokens), which were being skipped as they were assumed to be OpenAI-specific formats handled server-side.

The fix adds proper key mapping for ChatBedrock's format and improves the detection of genuine OpenAI responses by checking for additional OpenAI-specific keys (completion_token_details, prompt_token_details). This ensures ChatBedrock's usage metrics are properly captured while maintaining special handling for OpenAI responses.

Confidence score: 4/5

1. The changes look safe as they only impact token usage reporting for ChatBedrock and preserve existing behavior
2. The fix is well-reasoned and addresses the root cause by distinguishing between OpenAI and ChatBedrock formats
3. Files needing attention:
   • langfuse/langchain/CallbackHandler.py: The changes look good but should be tested with both ChatBedrock and OpenAI to ensure no regressions