Implementation:BerriAI Litellm MLflow Logger

Attribute	Value
Sources	litellm/integrations/mlflow.py
Domains	Logging, Tracing, ML Experiment Tracking, Integrations
Last Updated	2026-02-15 16:00 GMT

Overview

The MlflowLogger logs LLM call events as MLflow spans and traces, providing full tracing integration with the MLflow experiment tracking platform.

Description

MlflowLogger extends CustomLogger to create MLflow spans for each LLM call, capturing inputs (messages, tools), outputs (response objects), attributes (model, token usage, cost, cache status), and status codes. It supports both streaming and non-streaming responses, with streaming calls consolidated into a single span across all chunks. The logger automatically detects whether to create a new trace or nest a child span under an existing active span. It also sets MLflow-specific chat attributes (messages and tools) using MLflow's tracing utilities and handles both success and failure events, recording exceptions as SpanEvent objects.

Usage

Import and register MlflowLogger when you want to trace LLM calls in MLflow. Requires the mlflow package to be installed.

Code Reference

Source Location

litellm/integrations/mlflow.py

Signature

class MlflowLogger(CustomLogger):
    def __init__(self)

Import

from litellm.integrations.mlflow import MlflowLogger

I/O Contract

Inputs

Parameter	Type	Required	Description
(none)	--	--	Constructor requires no arguments. Uses MlflowClient() internally.

Key Methods

Method	Returns	Description
log_success_event(kwargs, response_obj, start_time, end_time)	None	Sync success handler. Delegates to _handle_success.
async_log_success_event(kwargs, response_obj, start_time, end_time)	None	Async success handler. Delegates to _handle_success.
log_failure_event(kwargs, response_obj, start_time, end_time)	None	Sync failure handler. Records exception as SpanEvent.
async_log_failure_event(kwargs, response_obj, start_time, end_time)	None	Async failure handler. Records exception as SpanEvent.
_handle_stream_event(kwargs, response_obj, start_time, end_time)	None	Handles streaming responses by creating one span per stream and adding chunk events.

Outputs

Output	Type	Description
Side effect	MLflow Span/Trace	Creates MLflow spans and traces via MlflowClient.start_trace() / start_span().

Usage Examples

import litellm

# Register the MLflow callback
litellm.success_callback = ["mlflow"]
litellm.failure_callback = ["mlflow"]

response = litellm.completion(
    model="gpt-4",
    messages=[{"role": "user", "content": "Hello!"}],
)
# A new MLflow trace is created with an LLM-type span

# With streaming
response = litellm.completion(
    model="gpt-4",
    messages=[{"role": "user", "content": "Tell me a story"}],
    stream=True,
)
for chunk in response:
    print(chunk)
# A single MLflow span is created for the entire stream

Related Pages

• BerriAI_Litellm_Braintrust_Logger - another evaluation/tracing integration
• BerriAI_Litellm_Literal_AI_Logger - another observability platform integration
• BerriAI_Litellm_PostHog_Logger - another analytics integration