Community: gather token usage info in BedrockChat during generation

[dmenini]

This PR allows to calculate token usage for prompts and completion directly in the generation method of BedrockChat. The token usage details are then returned together with the generations, so that other downstream tasks can access them easily.

This allows to define a callback for tokens tracking and cost calculation, similarly to what happens with OpenAI (see OpenAICallbackHandler. I plan on adding a BedrockCallbackHandler later.
Right now keeping track of tokens in the callback is already possible, but it requires passing the llm, as done here: https://how.wtf/how-to-count-amazon-bedrock-anthropic-tokens-with-langchain.html. However, I find the approach of this PR cleaner.

Thanks for your reviews. FYI @baskaryan, @hwchase17

[vercel]

The latest updates on your projects. Learn more about Vercel for Git ↗︎

1 Ignored Deployment

Name	Status	Preview	Comments	Updated (UTC)
langchain	⬜️ Ignored (Inspect)	Visit Preview		Mar 28, 2024 6:50pm

[esoler-sage]

Hey! For Bedrock you may want to use the tokens returned in the headers "x-amzn-bedrock-*", like

• x-amzn-bedrock-output-token-count
• x-amzn-bedrock-input-token-count
  instead of using anthropic only tokenizer 🤔?

[dmenini]

Hey! For Bedrock you may want to use the tokens returned in the headers "x-amzn-bedrock-*", like

• x-amzn-bedrock-output-token-count
• x-amzn-bedrock-input-token-count
  instead of using anthropic only tokenizer 🤔?

Thanks for this comment! I didn't know about those headers. It definitely makes sense and I'll rework my implementation accordingly

[dmenini]

The PR is updated! Now token counters are read from the headers in case of simple generation: The resulting LLMOutput will have the fields model_id and usage, containing the token counters.

NB In case of streaming, usage will be an empty dict as extracting the tokens is more complicated. Indeed, the headers are not available, and one would have to extract the token counters from the body, according to the model provider structure. Will do it at a later time :)

[baskaryan]

cc @3coins

[pratik60]

This is great! I was looking to achieve the same thing to enable cost monitoring on the Anthropic Bedrock models. Should we also add total_tokens (sum of prompt_tokens + completion_tokens), to keep it compatible with the OpenAI model?

[dmenini]

Thanks for the comment! I added the total_count

[pratik60]

Thanks! You have linting errors, fyi.

Hopefully one of the project maintainers has a chance to look at this, and get it merged soon!

[Sukitly]

Excellent work!
I'm curious about the differences between the use of "usage" in the response body versus in the headers. According to the Bedrock API documentation, the "usage" is included in the response body as follows:

{
    "id": string,
    "model": string,
    "type": "message",
    "role": "assistant",
    "content": [
        {
            "type": "text",
            "text": string
        }
    ],
    "stop_reason": string,
    "stop_sequence": string,
    "usage": {
        "input_tokens": integer,
        "output_tokens": integer
    }
    
}

Do you guys have any insights on this? @dmenini @esoler-sage @pratik60
Thanks!

[esoler-sage]

I think usage in headers is for sync calls, and async call' last chunk contains usage in body 🤔

Update: Also, the docs you were pointing are using the new messages api from Anthropic, only available on (some) Anthropic models