From cd913122528b864049c975d920a130fb81ae38c8 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 22 Sep 2023 20:43:03 -0700 Subject: [PATCH] track vertex ai usage --- litellm/main.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/litellm/main.py b/litellm/main.py index 4b050c34b64..1b2789f75d8 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -784,6 +784,19 @@ def completion( model_response["choices"][0]["message"]["content"] = str(completion_response) model_response["created"] = time.time() model_response["model"] = model + ## CALCULATING USAGE + prompt_tokens = len( + encoding.encode(prompt) + ) + completion_tokens = len( + encoding.encode(model_response["choices"][0]["message"]["content"]) + ) + + model_response["usage"] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } response = model_response elif model in litellm.vertex_text_models or model in litellm.vertex_code_text_models: try: @@ -823,6 +836,19 @@ def completion( model_response["choices"][0]["message"]["content"] = str(completion_response) model_response["created"] = time.time() model_response["model"] = model + ## CALCULATING USAGE + prompt_tokens = len( + encoding.encode(prompt) + ) + completion_tokens = len( + encoding.encode(model_response["choices"][0]["message"]["content"]) + ) + + model_response["usage"] = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": prompt_tokens + completion_tokens, + } response = model_response elif model in litellm.ai21_models: custom_llm_provider = "ai21"