phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 03:04:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix completion_tokens on Vertex AI Gemini thinking models

Browse source
This commit is contained in:
Andrew Wesie 2025-04-21 08:57:12 -05:00
parent 62ed5b2f78
commit f87f500a49
1 changed files with 9 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -766,14 +766,20 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
audio_tokens=audio_tokens, audio_tokens=audio_tokens,
text_tokens=text_tokens, text_tokens=text_tokens,
) )
completion_tokens = completion_response["usageMetadata"].get(
"candidatesTokenCount", 0
)
if reasoning_tokens:
# Usage(...) constructor expects that completion_tokens includes the reasoning_tokens.
# However the Vertex AI usage metadata does not include reasoning tokens in candidatesTokenCount.
# Reportedly, this is different from the Gemini API.
completion_tokens += reasoning_tokens
## GET USAGE ## ## GET USAGE ##
usage = Usage( usage = Usage(
prompt_tokens=completion_response["usageMetadata"].get( prompt_tokens=completion_response["usageMetadata"].get(
"promptTokenCount", 0 "promptTokenCount", 0
), ),
completion_tokens=completion_response["usageMetadata"].get( completion_tokens=completion_tokens,
"candidatesTokenCount", 0
),
total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0), total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0),
prompt_tokens_details=prompt_tokens_details, prompt_tokens_details=prompt_tokens_details,
reasoning_tokens=reasoning_tokens, reasoning_tokens=reasoning_tokens,