scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-10 22:41:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 83

fix(vertex_ai): fix toolUsePromptTokenCount in blocked/content-policy response paths

Browse source
This commit is contained in:
unknown 2026-03-26 14:28:38 +05:00
parent 3de5e32f63
commit f07e37ce69
2 changed files with 13 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1603,14 +1603,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
model_response.choices = [choice]
## GET USAGE ##
_usage_metadata = completion_response["usageMetadata"]
usage = Usage(
prompt_tokens=completion_response["usageMetadata"].get(
"promptTokenCount", 0
),
completion_tokens=completion_response["usageMetadata"].get(
"candidatesTokenCount", 0
),
total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0),
prompt_tokens=_usage_metadata.get("promptTokenCount", 0)
+ _usage_metadata.get("toolUsePromptTokenCount", 0),
completion_tokens=_usage_metadata.get("candidatesTokenCount", 0),
total_tokens=_usage_metadata.get("totalTokenCount", 0),
)
setattr(model_response, "usage", usage)
@ -1641,14 +1639,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
model_response.choices = [choice]
## GET USAGE ##
_usage_metadata = completion_response["usageMetadata"]
usage = Usage(
prompt_tokens=completion_response["usageMetadata"].get(
"promptTokenCount", 0
),
completion_tokens=completion_response["usageMetadata"].get(
"candidatesTokenCount", 0
),
total_tokens=completion_response["usageMetadata"].get("totalTokenCount", 0),
prompt_tokens=_usage_metadata.get("promptTokenCount", 0)
+ _usage_metadata.get("toolUsePromptTokenCount", 0),
completion_tokens=_usage_metadata.get("candidatesTokenCount", 0),
total_tokens=_usage_metadata.get("totalTokenCount", 0),
)
setattr(model_response, "usage", usage)

4
litellm/types/llms/vertex_ai.py
View file

@ -287,7 +287,9 @@ class UsageMetadata(TypedDict, total=False):
PromptTokensDetails
] # Alternative key name used in some responses
toolUsePromptTokenCount: int
toolUsePromptTokensDetails: List[PromptTokensDetails]
toolUsePromptTokensDetails: List[
PromptTokensDetails
] # TODO: parse into prompt_tokens_details per modality
class TokenCountDetailsResponse(TypedDict):