From 46ce4995b86b162b93515cc5d4e2c8db9afad065 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 12 Sep 2024 14:49:43 -0700 Subject: [PATCH] fix type errors --- litellm/types/utils.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 6ed3954f83..b5f1ba0805 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -473,6 +473,7 @@ class Usage(CompletionUsage): prompt_tokens: Optional[int] = None, completion_tokens: Optional[int] = None, total_tokens: Optional[int] = None, + reasoning_tokens: Optional[int] = None, **params, ): ## DEEPSEEK PROMPT TOKEN HANDLING ## - follow the anthropic format, of having prompt tokens be just the non-cached token input. Enables accurate cost-tracking - Relevant issue: https://github.com/BerriAI/litellm/issues/5285 @@ -482,12 +483,16 @@ class Usage(CompletionUsage): and prompt_tokens is not None ): prompt_tokens = params["prompt_cache_miss_tokens"] - data = { - "prompt_tokens": prompt_tokens or 0, - "completion_tokens": completion_tokens or 0, - "total_tokens": total_tokens or 0, - } - super().__init__(**data) + + completion_tokens_details = CompletionTokensDetails( + reasoning_tokens=reasoning_tokens + ) + super().__init__( + prompt_tokens=prompt_tokens or 0, + completion_tokens=completion_tokens or 0, + total_tokens=total_tokens or 0, + completion_tokens_details=completion_tokens_details or None, + ) ## ANTHROPIC MAPPING ## if "cache_creation_input_tokens" in params and isinstance(