diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index a2f0b7cf39c..2a27a9ec7a8 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -189,15 +189,21 @@ class MlflowLogger(CustomLogger): } standard_obj: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") if standard_obj: + token_usage = { + "input_tokens": standard_obj.get("prompt_tokens"), + "output_tokens": standard_obj.get("completion_tokens"), + "total_tokens": standard_obj.get("total_tokens"), + } + cache_read, cache_creation = self._extract_cache_token_usage(standard_obj) + if cache_read is not None: + token_usage["cache_read_input_tokens"] = cache_read + if cache_creation is not None: + token_usage["cache_creation_input_tokens"] = cache_creation attributes.update( { "api_base": standard_obj.get("api_base"), "cache_hit": standard_obj.get("cache_hit"), - "mlflow.chat.tokenUsage": { - "input_tokens": standard_obj.get("prompt_tokens"), - "output_tokens": standard_obj.get("completion_tokens"), - "total_tokens": standard_obj.get("total_tokens"), - }, + "mlflow.chat.tokenUsage": token_usage, "raw_llm_response": standard_obj.get("response"), "response_cost": standard_obj.get("response_cost"), "saved_cache_cost": standard_obj.get("saved_cache_cost"), @@ -217,6 +223,36 @@ class MlflowLogger(CustomLogger): ) return attributes + def _extract_cache_token_usage(self, standard_obj: StandardLoggingPayload) -> "tuple[int | None, int | None]": + """ + Extract cache read and cache creation token counts from the raw response usage. + + The flattened logging payload does not carry cache token fields, but MLflow + needs them to price cached tokens at their discounted rates. Anthropic-style + usage reports top-level cache fields while OpenAI-style usage nests them + under prompt_tokens_details. + """ + response = standard_obj.get("response") + usage = response.get("usage") if isinstance(response, dict) else None + if not isinstance(usage, dict): + return None, None + + details = usage.get("prompt_tokens_details") + if not isinstance(details, dict): + details = None + + cache_read = usage.get("cache_read_input_tokens") + if cache_read is None and details is not None: + cache_read = details.get("cached_tokens") + + cache_creation = usage.get("cache_creation_input_tokens") + if cache_creation is None and details is not None: + cache_creation = details.get("cache_creation_tokens") + if cache_creation is None and details is not None: + cache_creation = details.get("cache_write_tokens") + + return cache_read, cache_creation + def _get_span_type(self, call_type: str | None) -> str: from mlflow.entities import SpanType diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 61010f8531c..2815cfbea9f 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -86,18 +86,14 @@ async def test_mlflow_logging_functionality(): "jobID": "214590dsff09fds", "taskName": "run_page_classification", } - assert ( - tags_param == expected_tags - ), f"Expected tags {expected_tags}, got {tags_param}" + assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}" # Check that prediction parameter was included in inputs inputs_param = call_args.kwargs.get("inputs", {}) - assert ( - "prediction" in inputs_param - ), "Prediction should be included in span inputs" - assert ( - inputs_param["prediction"] == test_prediction - ), f"Expected prediction {test_prediction}, got {inputs_param['prediction']}" + assert "prediction" in inputs_param, "Prediction should be included in span inputs" + assert inputs_param["prediction"] == test_prediction, ( + f"Expected prediction {test_prediction}, got {inputs_param['prediction']}" + ) def test_mlflow_token_usage_attribute_structure(): @@ -138,6 +134,105 @@ def test_mlflow_token_usage_attribute_structure(): } +def test_mlflow_token_usage_includes_anthropic_style_cache_fields(): + """Cache token counts from the raw response usage are lifted into tokenUsage.""" + + mock_mlflow_tracking = MagicMock() + mock_mlflow_tracking.MlflowClient = MagicMock() + + with patch.dict( + "sys.modules", + { + "mlflow": MagicMock(), + "mlflow.tracking": mock_mlflow_tracking, + "mlflow.tracing.utils": MagicMock(), + }, + ): + from litellm.integrations.mlflow import MlflowLogger + + mlflow_logger = MlflowLogger() + + attrs = mlflow_logger._extract_attributes( # type: ignore + { + "litellm_call_id": "123", + "call_type": "completion", + "model": "claude-haiku-4-5", + "standard_logging_object": { + "prompt_tokens": 10500, + "completion_tokens": 200, + "total_tokens": 10700, + "response": { + "usage": { + "prompt_tokens": 10500, + "completion_tokens": 200, + "total_tokens": 10700, + "cache_read_input_tokens": 10000, + "cache_creation_input_tokens": 300, + } + }, + }, + } + ) + + assert attrs["mlflow.chat.tokenUsage"] == { + "input_tokens": 10500, + "output_tokens": 200, + "total_tokens": 10700, + "cache_read_input_tokens": 10000, + "cache_creation_input_tokens": 300, + } + + +def test_mlflow_token_usage_includes_openai_style_cached_tokens(): + """OpenAI-style responses nest cache counts under prompt_tokens_details.""" + + mock_mlflow_tracking = MagicMock() + mock_mlflow_tracking.MlflowClient = MagicMock() + + with patch.dict( + "sys.modules", + { + "mlflow": MagicMock(), + "mlflow.tracking": mock_mlflow_tracking, + "mlflow.tracing.utils": MagicMock(), + }, + ): + from litellm.integrations.mlflow import MlflowLogger + + mlflow_logger = MlflowLogger() + + attrs = mlflow_logger._extract_attributes( # type: ignore + { + "litellm_call_id": "123", + "call_type": "completion", + "model": "gpt-4o", + "standard_logging_object": { + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "response": { + "usage": { + "prompt_tokens": 100, + "completion_tokens": 10, + "total_tokens": 110, + "prompt_tokens_details": { + "cached_tokens": 80, + "audio_tokens": None, + }, + } + }, + }, + } + ) + + assert attrs["mlflow.chat.tokenUsage"] == { + "input_tokens": 100, + "output_tokens": 10, + "total_tokens": 110, + "cache_read_input_tokens": 80, + } + + def _mock_mlflow_modules(): mock_tracking = MagicMock() mock_tracking.MlflowClient = MagicMock() @@ -190,8 +285,5 @@ def test_mlflow_stream_handler_uses_async_complete_response(): ) mlflow_logger._end_span_or_trace.assert_called_once() - assert ( - mlflow_logger._end_span_or_trace.call_args.kwargs["outputs"] - is final_response - ) + assert mlflow_logger._end_span_or_trace.call_args.kwargs["outputs"] is final_response assert "abc123" not in mlflow_logger._stream_id_to_span