diff --git a/litellm/integrations/newrelic/newrelic.py b/litellm/integrations/newrelic/newrelic.py index 8b097799a27..3a021a1452b 100644 --- a/litellm/integrations/newrelic/newrelic.py +++ b/litellm/integrations/newrelic/newrelic.py @@ -329,9 +329,9 @@ class NewRelicLogger(CustomLogger): return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} return { - "prompt_tokens": usage.get("prompt_tokens", 0), - "completion_tokens": usage.get("completion_tokens", 0), - "total_tokens": usage.get("total_tokens", 0), + "prompt_tokens": usage.get("prompt_tokens") or 0, + "completion_tokens": usage.get("completion_tokens") or 0, + "total_tokens": usage.get("total_tokens") or 0, } def _get_finish_reason(self, response_obj: ModelResponse) -> str: diff --git a/tests/test_litellm/integrations/newrelic/test_newrelic.py b/tests/test_litellm/integrations/newrelic/test_newrelic.py index b53c1d9a28d..535c0e03148 100644 --- a/tests/test_litellm/integrations/newrelic/test_newrelic.py +++ b/tests/test_litellm/integrations/newrelic/test_newrelic.py @@ -426,6 +426,17 @@ class TestExtractUsage: usage = self.logger._extract_usage(response) assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + def test_explicit_none_token_fields_return_zeros(self): + response = { + "usage": { + "prompt_tokens": None, + "completion_tokens": None, + "total_tokens": None, + } + } + usage = self.logger._extract_usage(response) + assert usage == {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + class TestGetFinishReason: def setup_method(self):