From 7ea2cf73f11e92c6e5702e0b406a1666e763161f Mon Sep 17 00:00:00 2001 From: Gauthier Piarrette Date: Tue, 4 Aug 2026 04:06:11 -0700 Subject: [PATCH 1/4] fix(mlflow): include cache token counts in mlflow.chat.tokenUsage attribute --- litellm/integrations/mlflow.py | 67 +++++++++++-- .../test_litellm/integrations/test_mlflow.py | 99 +++++++++++++++++++ 2 files changed, 156 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index f41de320843..aadf7d49d17 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -62,7 +62,10 @@ class MlflowLogger(CustomLogger): inputs = self._construct_input(kwargs) input_messages = inputs.get("messages", []) - output_messages = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])] + output_messages = [ + c.message.model_dump(exclude_none=True) + for c in getattr(response_obj, "choices", []) + ] if messages := [*input_messages, *output_messages]: set_span_chat_messages(span, messages) if tools := inputs.get("tools"): @@ -129,7 +132,9 @@ class MlflowLogger(CustomLogger): # If this is the final chunk, end the span. The final chunk # has the assembled streaming response (key differs between sync/async paths). - final_response = kwargs.get("complete_streaming_response") or kwargs.get("async_complete_streaming_response") + final_response = kwargs.get("complete_streaming_response") or kwargs.get( + "async_complete_streaming_response" + ) if final_response: end_time_ns = int(end_time.timestamp() * 1e9) @@ -153,7 +158,9 @@ class MlflowLogger(CustomLogger): span.add_event( SpanEvent( name="streaming_chunk", - attributes={"delta": json.dumps(choice.delta.model_dump, default=str)}, + attributes={ + "delta": json.dumps(choice.delta.model_dump, default=str) + }, ) ) except Exception: @@ -187,17 +194,21 @@ class MlflowLogger(CustomLogger): "call_type": kwargs.get("call_type"), "model": kwargs.get("model"), } - standard_obj: StandardLoggingPayload | None = kwargs.get("standard_logging_object") + standard_obj: 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"), + } + token_usage.update(self._extract_cache_token_usage(standard_obj)) 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 +228,40 @@ class MlflowLogger(CustomLogger): ) return attributes + def _extract_cache_token_usage(self, standard_obj) -> dict: + """ + Extract cache 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 {} + + details = usage.get("prompt_tokens_details") + details = details if isinstance(details, dict) else {} + + cache_read = usage.get("cache_read_input_tokens") + if cache_read is None: + cache_read = details.get("cached_tokens") + + cache_creation = usage.get("cache_creation_input_tokens") + if cache_creation is None: + cache_creation = details.get("cache_creation_tokens") + if cache_creation is None: + cache_creation = details.get("cache_write_tokens") + + cache_token_usage = {} + if cache_read is not None: + cache_token_usage["cache_read_input_tokens"] = cache_read + if cache_creation is not None: + cache_token_usage["cache_creation_input_tokens"] = cache_creation + return cache_token_usage + def _get_span_type(self, call_type: str | None) -> str: from mlflow.entities import SpanType @@ -260,7 +305,9 @@ class MlflowLogger(CustomLogger): span_type=span_type, inputs=inputs, attributes=attributes, - tags=self._transform_tag_list_to_dict(attributes.get("request_tags", [])), + tags=self._transform_tag_list_to_dict( + attributes.get("request_tags", []) + ), start_time_ns=start_time_ns, ) diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 32358641984..8d88f456731 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -141,6 +141,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() From dbb592967893bb765ef3513686875eb2b21144f5 Mon Sep 17 00:00:00 2001 From: Gauthier Piarrette Date: Tue, 4 Aug 2026 05:12:37 -0700 Subject: [PATCH 2/4] fix(mlflow): type the cache token extraction helper --- litellm/integrations/mlflow.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index aadf7d49d17..80766c3fee6 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -228,7 +228,9 @@ class MlflowLogger(CustomLogger): ) return attributes - def _extract_cache_token_usage(self, standard_obj) -> dict: + def _extract_cache_token_usage( + self, standard_obj: StandardLoggingPayload + ) -> dict[str, int]: """ Extract cache token counts from the raw response usage. From 88dec6f2eec93e224f8b1f3438a0db7607ce4903 Mon Sep 17 00:00:00 2001 From: Gauthier Piarrette Date: Tue, 4 Aug 2026 05:21:13 -0700 Subject: [PATCH 3/4] style(mlflow): apply ruff format --- litellm/integrations/mlflow.py | 25 +++++-------------- .../test_litellm/integrations/test_mlflow.py | 19 +++++--------- 2 files changed, 12 insertions(+), 32 deletions(-) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 80766c3fee6..0965d650f89 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -62,10 +62,7 @@ class MlflowLogger(CustomLogger): inputs = self._construct_input(kwargs) input_messages = inputs.get("messages", []) - output_messages = [ - c.message.model_dump(exclude_none=True) - for c in getattr(response_obj, "choices", []) - ] + output_messages = [c.message.model_dump(exclude_none=True) for c in getattr(response_obj, "choices", [])] if messages := [*input_messages, *output_messages]: set_span_chat_messages(span, messages) if tools := inputs.get("tools"): @@ -132,9 +129,7 @@ class MlflowLogger(CustomLogger): # If this is the final chunk, end the span. The final chunk # has the assembled streaming response (key differs between sync/async paths). - final_response = kwargs.get("complete_streaming_response") or kwargs.get( - "async_complete_streaming_response" - ) + final_response = kwargs.get("complete_streaming_response") or kwargs.get("async_complete_streaming_response") if final_response: end_time_ns = int(end_time.timestamp() * 1e9) @@ -158,9 +153,7 @@ class MlflowLogger(CustomLogger): span.add_event( SpanEvent( name="streaming_chunk", - attributes={ - "delta": json.dumps(choice.delta.model_dump, default=str) - }, + attributes={"delta": json.dumps(choice.delta.model_dump, default=str)}, ) ) except Exception: @@ -194,9 +187,7 @@ class MlflowLogger(CustomLogger): "call_type": kwargs.get("call_type"), "model": kwargs.get("model"), } - standard_obj: StandardLoggingPayload | None = kwargs.get( - "standard_logging_object" - ) + standard_obj: StandardLoggingPayload | None = kwargs.get("standard_logging_object") if standard_obj: token_usage = { "input_tokens": standard_obj.get("prompt_tokens"), @@ -228,9 +219,7 @@ class MlflowLogger(CustomLogger): ) return attributes - def _extract_cache_token_usage( - self, standard_obj: StandardLoggingPayload - ) -> dict[str, int]: + def _extract_cache_token_usage(self, standard_obj: StandardLoggingPayload) -> dict[str, int]: """ Extract cache token counts from the raw response usage. @@ -307,9 +296,7 @@ class MlflowLogger(CustomLogger): span_type=span_type, inputs=inputs, attributes=attributes, - tags=self._transform_tag_list_to_dict( - attributes.get("request_tags", []) - ), + tags=self._transform_tag_list_to_dict(attributes.get("request_tags", [])), start_time_ns=start_time_ns, ) diff --git a/tests/test_litellm/integrations/test_mlflow.py b/tests/test_litellm/integrations/test_mlflow.py index 8d88f456731..a781d4327d6 100644 --- a/tests/test_litellm/integrations/test_mlflow.py +++ b/tests/test_litellm/integrations/test_mlflow.py @@ -89,18 +89,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(): @@ -292,8 +288,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 From 120cc79b2e86ea1cb467d796531af03ef4d24359 Mon Sep 17 00:00:00 2001 From: Gauthier Piarrette Date: Tue, 4 Aug 2026 05:31:40 -0700 Subject: [PATCH 4/4] fix(mlflow): return cache token counts as a tuple to satisfy the type discipline gate --- litellm/integrations/mlflow.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/litellm/integrations/mlflow.py b/litellm/integrations/mlflow.py index 0965d650f89..392b6785c7e 100644 --- a/litellm/integrations/mlflow.py +++ b/litellm/integrations/mlflow.py @@ -194,7 +194,11 @@ class MlflowLogger(CustomLogger): "output_tokens": standard_obj.get("completion_tokens"), "total_tokens": standard_obj.get("total_tokens"), } - token_usage.update(self._extract_cache_token_usage(standard_obj)) + 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"), @@ -219,9 +223,9 @@ class MlflowLogger(CustomLogger): ) return attributes - def _extract_cache_token_usage(self, standard_obj: StandardLoggingPayload) -> dict[str, int]: + def _extract_cache_token_usage(self, standard_obj: StandardLoggingPayload) -> "tuple[int | None, int | None]": """ - Extract cache token counts from the raw response usage. + 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 @@ -231,27 +235,23 @@ class MlflowLogger(CustomLogger): response = standard_obj.get("response") usage = response.get("usage") if isinstance(response, dict) else None if not isinstance(usage, dict): - return {} + return None, None details = usage.get("prompt_tokens_details") - details = details if isinstance(details, dict) else {} + if not isinstance(details, dict): + details = None cache_read = usage.get("cache_read_input_tokens") - if cache_read is None: + 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: + if cache_creation is None and details is not None: cache_creation = details.get("cache_creation_tokens") - if cache_creation is None: + if cache_creation is None and details is not None: cache_creation = details.get("cache_write_tokens") - cache_token_usage = {} - if cache_read is not None: - cache_token_usage["cache_read_input_tokens"] = cache_read - if cache_creation is not None: - cache_token_usage["cache_creation_input_tokens"] = cache_creation - return cache_token_usage + return cache_read, cache_creation def _get_span_type(self, call_type: str | None) -> str: from mlflow.entities import SpanType