How to get used tokens from bedrock AI

31 Views Asked by At

I want to retrieve used tokens from bedrock AI request. I managed to do this in Python like this

def extract_token_counts_from_response(response):
"""
Extracts token counts from the response's HTTP headers.

Args:
    response (dict): The response object containing 'ResponseMetadata'.

Returns:
    tuple: A tuple containing input token count and output token count.
          Returns None for each if not found.
"""
# Ensure 'ResponseMetadata' and 'HTTPHeaders' exist in the response
if 'ResponseMetadata' in response and 'HTTPHeaders' in response['ResponseMetadata']:
    http_headers = response['ResponseMetadata']['HTTPHeaders']

    # Extracting input and output token counts
    input_token_count = http_headers.get('x-amzn-bedrock-input-token-count')
    output_token_count = http_headers.get('x-amzn-bedrock-output-token-count')

    # Converting token counts to int if they are not None
    input_token_count = int(input_token_count) if input_token_count is not None else None
    output_token_count = int(output_token_count) if output_token_count is not None else None

    return input_token_count, output_token_count
else:
    # Returning None for both counts if the expected structure is not found
    return None, None


response = bedrock.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType)
input_tokens, output_tokens = extract_token_counts_from_response(response)

I can not get this information from nodejs

    const command = new InvokeModelCommand(input);
console.log('bedrock start')
let response;
try {
    response = await client.send(command);

    const [inputTokenCount, outputTokenCount] = extractTokenCountsFromResponse(response);

Response comes without a header. Here is the response

  '$metadata': {
httpStatusCode: 200,
requestId: '869bc908-db4c-4c9d-adc5-c0e643c9d959',
extendedRequestId: undefined,
cfId: undefined,
attempts: 1,
totalRetryDelay: 0

}, contentType: 'application/json', body: Uint8ArrayBlobAdapter(791) [Uint8Array] [ 123, 34, 99, 111, 109, 112, 108, 101, 116, 105, 111, 110, 34, 58, 34, 32, 60, 106, 115, 111, 110, 62, 92, 110, 91, 92, 110, 32, 32, 123, 92, 34, 110, 97, 109, 101, 92, 34, 58, 92, 34, 65, 110, 103, 101, 108, 105, 110, 97, 32, 66, 101, 110, 104, 97, 109, 92, 34, 44, 92, 34, 116, 121, 112, 101, 92, 34, 58, 92, 34, 99, 111, 110, 116, 97, 99, 116, 92, 34, 44, 92, 34, 97, 102, 102, 105, 108, 105, 97, 116, 101, 100, 92, 34, 58, 92, 34, 65, 112, 112, ... 691 more items ] }

Anyone familiar how to get the used tokens?

0

There are 0 best solutions below