def get_tokens_from_response(
    response: Union["CompletionResponse", "ChatResponse"],
) -> Tuple[int, int]:
    """Get the token counts from a raw response."""
    raw_response = response.raw
    if not isinstance(raw_response, dict):
        raw_response = dict(raw_response or {})

    usage = raw_response.get("usage", raw_response.get("usage_metadata", {}))
    if usage is None:
        usage = response.additional_kwargs

    if not usage:
        return 0, 0

    if not isinstance(usage, dict):
        usage = usage.model_dump()

    possible_input_keys = ("prompt_tokens", "input_tokens", "prompt_token_count")
    possible_output_keys = (
        "completion_tokens",
        "output_tokens",
        "candidates_token_count",
    )

    prompt_tokens = 0
    for input_key in possible_input_keys:
        if input_key in usage:
            prompt_tokens = usage[input_key]
            break

    completion_tokens = 0
    for output_key in possible_output_keys:
        if output_key in usage:
            completion_tokens = usage[output_key]
            break

    return prompt_tokens, completion_tokens
Get the token counts from a raw response.
llama-index-core/llama_index/core/callbacks/token_counting.py