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 124

fix(mlflow): include cache token counts in mlflow.chat.tokenUsage attribute

Browse source
This commit is contained in:
Gauthier Piarrette 2026-08-04 04:06:11 -07:00
parent 956d5177d1
commit 7ea2cf73f1
2 changed files with 156 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

67
litellm/integrations/mlflow.py
View file

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

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

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