(feat) use usage class for vertex ai

This commit is contained in:
ishaan-jaff 2023-10-27 09:38:37 -07:00
parent 692bb3f4e9
commit 3691dc9b08

View file

@ -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