From 3691dc9b08ae061270e8a1c2d707100a35ef7103 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 27 Oct 2023 09:38:37 -0700 Subject: [PATCH] (feat) use usage class for vertex ai --- litellm/llms/vertex_ai.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/litellm/llms/vertex_ai.py b/litellm/llms/vertex_ai.py index f124a088ed4..9db661c3012 100644 --- a/litellm/llms/vertex_ai.py +++ b/litellm/llms/vertex_ai.py @@ -144,11 +144,9 @@ def completion( encoding.encode(model_response["choices"][0]["message"].get("content", "")) ) - 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