From 015a398713502d5b84c210b87bf106db277cdf3d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 5 Jul 2024 11:01:33 -0700 Subject: [PATCH] fix(vertex_httpx.py): fix assumptions on usagemetadata --- litellm/llms/vertex_httpx.py | 40 ++++++++++++++++++++------------- litellm/types/llms/vertex_ai.py | 4 ++-- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index 31910b7ea39..d56237b73d1 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -603,15 +603,15 @@ class VertexLLM(BaseLLM): ## GET USAGE ## usage = litellm.Usage( - prompt_tokens=completion_response["usageMetadata"][ - "promptTokenCount" - ], + prompt_tokens=completion_response["usageMetadata"].get( + "promptTokenCount", 0 + ), completion_tokens=completion_response["usageMetadata"].get( "candidatesTokenCount", 0 ), - total_tokens=completion_response["usageMetadata"][ - "totalTokenCount" - ], + total_tokens=completion_response["usageMetadata"].get( + "totalTokenCount", 0 + ), ) setattr(model_response, "usage", usage) @@ -647,15 +647,15 @@ class VertexLLM(BaseLLM): ## GET USAGE ## usage = litellm.Usage( - prompt_tokens=completion_response["usageMetadata"][ - "promptTokenCount" - ], + prompt_tokens=completion_response["usageMetadata"].get( + "promptTokenCount", 0 + ), completion_tokens=completion_response["usageMetadata"].get( "candidatesTokenCount", 0 ), - total_tokens=completion_response["usageMetadata"][ - "totalTokenCount" - ], + total_tokens=completion_response["usageMetadata"].get( + "totalTokenCount", 0 + ), ) setattr(model_response, "usage", usage) @@ -705,11 +705,15 @@ class VertexLLM(BaseLLM): ## GET USAGE ## usage = litellm.Usage( - prompt_tokens=completion_response["usageMetadata"]["promptTokenCount"], + prompt_tokens=completion_response["usageMetadata"].get( + "promptTokenCount", 0 + ), completion_tokens=completion_response["usageMetadata"].get( "candidatesTokenCount", 0 ), - total_tokens=completion_response["usageMetadata"]["totalTokenCount"], + total_tokens=completion_response["usageMetadata"].get( + "totalTokenCount", 0 + ), ) setattr(model_response, "usage", usage) @@ -1340,11 +1344,15 @@ class ModelResponseIterator: if "usageMetadata" in processed_chunk: usage = ChatCompletionUsageBlock( - prompt_tokens=processed_chunk["usageMetadata"]["promptTokenCount"], + prompt_tokens=processed_chunk["usageMetadata"].get( + "promptTokenCount", 0 + ), completion_tokens=processed_chunk["usageMetadata"].get( "candidatesTokenCount", 0 ), - total_tokens=processed_chunk["usageMetadata"]["totalTokenCount"], + total_tokens=processed_chunk["usageMetadata"].get( + "totalTokenCount", 0 + ), ) returned_chunk = GenericStreamingChunk( diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 17fc26d60ef..f07767a2dd4 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -239,8 +239,8 @@ class PromptFeedback(TypedDict): class UsageMetadata(TypedDict, total=False): - promptTokenCount: Required[int] - totalTokenCount: Required[int] + promptTokenCount: int + totalTokenCount: int candidatesTokenCount: int