From c6f2323ed9327a062074b7741e659e24e09e19c7 Mon Sep 17 00:00:00 2001
From: spions <spionsbbs@gmail.com>
Date: Fri, 7 Aug 2026 17:01:11 +0300
Subject: [PATCH] test: cover rerank usage fallback in the GitHub Actions test
 suite

tests/logging_callback_tests/test_standard_logging_payload.py runs
via CircleCI, which doesn't trigger on this fork's PR, so its coverage
of get_usage_from_response_obj/get_usage_as_dict's new rerank fallback
never reached codecov/patch here. Mirrored the same assertions into
tests/test_litellm/litellm_core_utils/test_litellm_logging.py, which
the "core-utils" GitHub Actions job does run and upload coverage for.
---
 .../test_litellm_logging.py                   | 64 +++++++++++++++++++
 1 file changed, 64 insertions(+)

diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
index 23e0975cd08..a5981e3a98d 100644
--- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
+++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py
@@ -4230,3 +4230,67 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj):
     logging_obj.post_call(original_response='{"ok": true}', input=big_input, api_key="sk-test")
 
     assert litellm.error_logs == {}
+
+
+def test_get_usage_from_response_obj_rerank_meta():
+    """
+    RerankResponse has no top-level `usage` field - token usage lives under
+    `meta.billed_units`/`meta.tokens` (Cohere-style API), which is also what
+    hosted_vllm/other rerank servers return. Without this, /ui/usage always
+    showed 0 tokens for every rerank call.
+    """
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {
+        "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59",
+        "results": [],
+        "meta": {
+            "billed_units": {"total_tokens": 42},
+            "tokens": {"input_tokens": 42},
+        },
+    }
+
+    usage = StandardLoggingPayloadSetup.get_usage_from_response_obj(response_obj)
+
+    assert usage.prompt_tokens == 42
+    assert usage.completion_tokens == 0
+    assert usage.total_tokens == 42
+
+
+def test_get_usage_as_dict_rerank_meta():
+    """
+    get_usage_as_dict() is the function actually called by
+    get_standard_logging_object_payload() (the hot path used to build spend
+    logs / /ui/usage). It must handle the rerank `meta.billed_units`/`meta.tokens`
+    shape the same way get_usage_from_response_obj() does, or /ui/usage keeps
+    showing 0 tokens for rerank calls even after fixing the other function.
+    """
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {
+        "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59",
+        "results": [],
+        "meta": {
+            "billed_units": {"total_tokens": 42},
+            "tokens": {"input_tokens": 42},
+        },
+    }
+
+    usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
+
+    assert usage_dict["prompt_tokens"] == 42
+    assert usage_dict["completion_tokens"] == 0
+    assert usage_dict["total_tokens"] == 42
+
+
+def test_get_usage_as_dict_rerank_meta_no_meta():
+    """A rerank response with no usable `meta` (and no `usage`) falls back to zero."""
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {"id": "rerank-abc", "results": [], "meta": None}
+
+    usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
+
+    assert usage_dict["prompt_tokens"] == 0
+    assert usage_dict["completion_tokens"] == 0
+    assert usage_dict["total_tokens"] == 0