diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index 7014ebc42b5..8c9a98352a3 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -392,11 +392,9 @@ def completion( model_response["created"] = time.time() model_response["model"] = model - model_response["usage"] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } + model_response.usage.completion_tokens = completion_tokens + model_response.usage.prompt_tokens = prompt_tokens + model_response.usage.total_tokens = prompt_tokens + completion_tokens return model_response diff --git a/litellm/llms/palm.py b/litellm/llms/palm.py index b4160b63b5e..79a91364943 100644 --- a/litellm/llms/palm.py +++ b/litellm/llms/palm.py @@ -157,11 +157,9 @@ def completion( model_response["created"] = time.time() model_response["model"] = "palm/" + model - model_response["usage"] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } + model_response.usage.completion_tokens = completion_tokens + model_response.usage.prompt_tokens = prompt_tokens + model_response.usage.total_tokens = prompt_tokens + completion_tokens return model_response def embedding(): diff --git a/litellm/llms/petals.py b/litellm/llms/petals.py index a3127eade71..5834129c199 100644 --- a/litellm/llms/petals.py +++ b/litellm/llms/petals.py @@ -176,11 +176,9 @@ def completion( model_response["created"] = time.time() model_response["model"] = model - model_response["usage"] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } + model_response.usage.completion_tokens = completion_tokens + model_response.usage.prompt_tokens = prompt_tokens + model_response.usage.total_tokens = prompt_tokens + completion_tokens return model_response def embedding(): diff --git a/litellm/llms/sagemaker.py b/litellm/llms/sagemaker.py index 962a2fc0e1e..8c999af6396 100644 --- a/litellm/llms/sagemaker.py +++ b/litellm/llms/sagemaker.py @@ -169,11 +169,9 @@ def completion( model_response["created"] = time.time() model_response["model"] = model - model_response["usage"] = { - "prompt_tokens": prompt_tokens, - "completion_tokens": completion_tokens, - "total_tokens": prompt_tokens + completion_tokens, - } + model_response.usage.completion_tokens = completion_tokens + model_response.usage.prompt_tokens = prompt_tokens + model_response.usage.total_tokens = prompt_tokens + completion_tokens return model_response def embedding():