mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
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.
This commit is contained in:
parent
de365aa6c7
commit
c6f2323ed9
1 changed files with 64 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue