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.py
andtests/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
ApiMeta
,ApiMetaTokens
, andApiMetaBilledUnits
classes are imported.map_token_counts
function is added to calculate token counts from response headers.map_response_from_bedrock
function now adds token count information to the response object.test_chat
function intests/test_aws_client.py
now asserts that the token and billed unit information is not None.