scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 154

Merge 120cc79b2e into eb0e3f8c18

Browse source
This commit is contained in:
Gauthier Piarrette 2026-08-27 19:18:23 -05:00 • committed by GitHub
commit 508809f138
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 146 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
litellm/integrations/mlflow.py
View file

@ -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

118
tests/test_litellm/integrations/test_mlflow.py
View file

@ -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