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

style(mlflow): apply ruff format

Browse source
This commit is contained in:
Gauthier Piarrette 2026-08-04 05:21:13 -07:00
parent dbb5929678
commit 88dec6f2ee
2 changed files with 12 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
litellm/integrations/mlflow.py
View file

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

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

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