diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index d75d35e4531..bfa6c8f8831 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -6603,6 +6603,18 @@ class StandardLoggingPayloadSetup: return clean_metadata + @staticmethod + def _get_rerank_usage_from_meta(meta: object) -> Usage | None: + if not isinstance(meta, dict) or ("billed_units" not in meta and "tokens" not in meta): + return None + billed_units: Final = meta.get("billed_units") + tokens: Final = meta.get("tokens") + meta_total: Final = billed_units.get("total_tokens") if isinstance(billed_units, dict) else None + meta_input: Final = tokens.get("input_tokens") if isinstance(tokens, dict) else None + total_tokens: Final = meta_total if meta_total is not None else (meta_input if meta_input is not None else 0) + prompt_tokens: Final = meta_input if meta_input is not None else (meta_total if meta_total is not None else 0) + return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=total_tokens) + @staticmethod def get_usage_from_response_obj(response_obj: dict | None, combined_usage_object: Usage | None = None) -> Usage: ## BASE CASE ## @@ -6616,6 +6628,10 @@ class StandardLoggingPayloadSetup: ) usage: Final = response_obj.get("usage", None) or {} + if not usage: + rerank_usage: Final = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(response_obj.get("meta")) + if rerank_usage is not None: + return rerank_usage if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)): return Usage( prompt_tokens=0, @@ -6651,6 +6667,9 @@ class StandardLoggingPayloadSetup: return _empty _raw: Final = response_obj.get("usage", None) if _raw is None: + rerank_usage: Final = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(response_obj.get("meta")) + if rerank_usage is not None: + return rerank_usage.model_dump() return _empty if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 4a5729945b4..6a5575058cf 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -203,10 +203,21 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers) def _transform_response(self, response: dict) -> RerankResponse: - # Extract usage information - usage_data: Final = response.get("usage", {}) - _billed_units: Final = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0)) - _tokens: Final = RerankTokens(input_tokens=usage_data.get("total_tokens", 0)) + raw_meta: Final = response.get("meta") + usage_data: Final = response.get("usage") + meta_billed_units: Final = raw_meta.get("billed_units") if isinstance(raw_meta, dict) else None + meta_tokens: Final = raw_meta.get("tokens") if isinstance(raw_meta, dict) else None + usage_total_tokens: Final = usage_data.get("total_tokens", 0) if isinstance(usage_data, dict) else 0 + meta_total: Final = meta_billed_units.get("total_tokens") if isinstance(meta_billed_units, dict) else None + meta_input: Final = meta_tokens.get("input_tokens") if isinstance(meta_tokens, dict) else None + total_tokens: Final = ( + meta_total if meta_total is not None else (meta_input if meta_input is not None else usage_total_tokens) + ) + input_tokens: Final = ( + meta_input if meta_input is not None else (meta_total if meta_total is not None else usage_total_tokens) + ) + _billed_units: Final = RerankBilledUnits(total_tokens=total_tokens) + _tokens: Final = RerankTokens(input_tokens=input_tokens) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) # Extract results diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index bbfae2229ce..df6d95793e2 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -7323,6 +7323,97 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj): assert litellm.error_logs == {} +@pytest.mark.parametrize( + "meta,expected", + [ + (None, None), + ({}, None), + ({"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}}, (42, 42)), + ({"tokens": {"input_tokens": 38}}, (38, 38)), + ({"billed_units": {"total_tokens": 38}}, (38, 38)), + ], +) +def test_get_rerank_usage_from_meta(meta, expected): + """Direct unit test of the helper both get_usage_from_response_obj and + get_usage_as_dict delegate to for the rerank `meta` shape.""" + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + result = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(meta) + + if expected is None: + assert result is None + else: + prompt_tokens, total_tokens = expected + assert result is not None + assert result.prompt_tokens == prompt_tokens + assert result.completion_tokens == 0 + assert result.total_tokens == total_tokens + + +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 + + def test_handle_anthropic_messages_response_logging_preserves_fast_mode_speed(): """/v1/messages non-streaming rebuilds usage by re-transforming the raw Anthropic response. Anthropic's fast-mode multiplier is applied off ``usage.speed``, which the @@ -7662,6 +7753,31 @@ async def test_restore_correlation_context_works_across_asyncio_task_boundary(mo session_id_var.set("") +@pytest.mark.parametrize( + "meta", + [ + {"tokens": {"input_tokens": 38}}, + {"billed_units": {"total_tokens": 38}}, + ], +) +def test_get_usage_as_dict_rerank_meta_partial_fields_cross_fill(meta): + """ + Some rerank servers report only `tokens.input_tokens` or only + `billed_units.total_tokens`, not both. Since rerank has no completion step, + either field alone should populate both prompt_tokens and total_tokens + instead of leaving one of them at zero. + """ + from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup + + response_obj = {"id": "rerank-abc", "results": [], "meta": meta} + + usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj) + + assert usage_dict["prompt_tokens"] == 38 + assert usage_dict["completion_tokens"] == 0 + assert usage_dict["total_tokens"] == 38 + + class TestNonInferenceCallTypesAreNotBilled: """A retrieved response replays the usage of the call that created it, so pricing a read of it double bills the same tokens. Regression tests for LIT-5602.""" diff --git a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 9a62fcf6f0f..bdf0392b25d 100644 --- a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -133,6 +133,47 @@ class TestHostedVLLMRerankTransform: assert result.meta["billed_units"]["total_tokens"] == 42 assert result.meta["tokens"]["input_tokens"] == 42 + def test_transform_response_with_meta(self): + """Some vLLM rerank servers (Cohere-compatible API) return usage under + `meta.billed_units`/`meta.tokens` instead of a top-level `usage` key.""" + response_dict = { + "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59", + "results": [ + {"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}}, + {"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}}, + ], + "meta": { + "billed_units": {"total_tokens": 42}, + "tokens": {"input_tokens": 42}, + }, + } + result = self.config._transform_response(response_dict) + assert result.id == "rerank-d618748e0f5543e8ba09ee7dd131ac59" + assert result.meta["billed_units"]["total_tokens"] == 42 + assert result.meta["tokens"]["input_tokens"] == 42 + + @pytest.mark.parametrize( + "meta", + [ + {"tokens": {"input_tokens": 38}}, + {"billed_units": {"total_tokens": 38}}, + ], + ) + def test_transform_response_with_meta_partial_fields_cross_fill(self, meta): + """Some servers report only one of `tokens`/`billed_units`. Since rerank + has no completion step, either field alone should populate both + `meta.tokens.input_tokens` and `meta.billed_units.total_tokens`.""" + response_dict = { + "id": "rerank-abc", + "results": [ + {"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}}, + ], + "meta": meta, + } + result = self.config._transform_response(response_dict) + assert result.meta["billed_units"]["total_tokens"] == 38 + assert result.meta["tokens"]["input_tokens"] == 38 + def test_transform_response_missing_results(self): response_dict = {"id": "abc123", "usage": {"total_tokens": 10}} with pytest.raises(ValueError, match="No results found in the response="):