Sanitize null token values in provider responses (#16493)

litellm/llms/openai_like/chat/transformation.py

@@ -60,6 +60,23 @@ class OpenAILikeChatConfig(OpenAIGPTConfig):
 
         return message
 
+    @staticmethod
+    def _sanitize_usage_obj(response_json: dict) -> dict:
+        """
+        Checks for a 'usage' object in the response and replaces any None token values with 0.
+        This enforces OpenAI compatibility for providers that might return null.
+
+        This method is future-proof and sanitizes any key ending in '_tokens'.
+        """
+        if "usage" in response_json and isinstance(response_json.get("usage"), dict):
+            usage = response_json["usage"]
+            # Iterate through all keys in the usage dictionary
+            for key, value in usage.items():
+                # Sanitize if the key ends with '_tokens' and its value is None
+                if key.endswith("_tokens") and value is None:
+                    usage[key] = 0
+        return response_json
+
     @staticmethod
     def _transform_response(
         model: str,
@@ -85,6 +102,9 @@ class OpenAILikeChatConfig(OpenAIGPTConfig):
             additional_args={"complete_input_dict": data},
         )
 
+        # Sanitize the usage object at the source
+        response_json = OpenAILikeChatConfig._sanitize_usage_obj(response_json)
+
         if json_mode:
             for choice in response_json["choices"]:
                 message = (

tests/test_litellm/llms/openai_like/chat/test_transformation.py

@@ -0,0 +1,49 @@
+import pytest
+from litellm.llms.openai_like.chat.transformation import OpenAILikeChatConfig
+
+
+def test_sanitize_usage_obj_handles_null_tokens():
+    """
+    Tests that _sanitize_usage_obj correctly converts None values for token counts to 0.
+    """
+    response_json = {
+        "choices": [],
+        "usage": {"prompt_tokens": None, "completion_tokens": 50, "total_tokens": None},
+    }
+
+    sanitized_json = OpenAILikeChatConfig._sanitize_usage_obj(response_json)
+
+    # Assert
+    assert sanitized_json["usage"]["prompt_tokens"] == 0
+    assert sanitized_json["usage"]["completion_tokens"] == 50  # Should remain unchanged
+    assert sanitized_json["usage"]["total_tokens"] == 0
+
+
+def test_sanitize_usage_obj_no_usage():
+    """
+    Tests that the sanitizer handles cases where the 'usage' object is missing.
+    """
+    response_json = {"choices": []}
+
+    sanitized_json = OpenAILikeChatConfig._sanitize_usage_obj(response_json)
+
+    # Assert
+    assert "usage" not in sanitized_json  # Should not add a usage key
+
+
+def test_sanitize_usage_obj_valid_usage():
+    """
+    Tests that the sanitizer does not modify a valid usage object.
+    """
+    response_json = {
+        "choices": [],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
+    }
+
+    # Create a copy to compare against
+    original_json = response_json.copy()
+
+    sanitized_json = OpenAILikeChatConfig._sanitize_usage_obj(response_json)
+
+    # Assert
+    assert sanitized_json == original_json  # The object should be unchanged