The PR introduces changes to the src/cohere/aws_client.py and tests/test_aws_client.py files.
Summary
The changes add token and billing information to the responses from the AWS client. Specifically, the map_token_counts function is introduced to calculate the input and output token counts from the response headers. This information is then added to the response object under the "meta" field.
Changes
The ApiMeta, ApiMetaTokens, and ApiMetaBilledUnits classes are imported.
The map_token_counts function is added to calculate token counts from response headers.
The map_response_from_bedrock function now adds token count information to the response object.
The test_chat function in tests/test_aws_client.py now asserts that the token and billed unit information is not None.
The PR introduces changes to the
src/cohere/aws_client.pyandtests/test_aws_client.pyfiles.Summary
The changes add token and billing information to the responses from the AWS client. Specifically, the
map_token_countsfunction is introduced to calculate the input and output token counts from the response headers. This information is then added to the response object under the "meta" field.Changes
ApiMeta,ApiMetaTokens, andApiMetaBilledUnitsclasses are imported.map_token_countsfunction is added to calculate token counts from response headers.map_response_from_bedrockfunction now adds token count information to the response object.test_chatfunction intests/test_aws_client.pynow asserts that the token and billed unit information is not None.