[Feat] MLFlow Logging - Allow adding tags for ML Flow logging requests (#13108)

* add mlflow tags

* fixes config

* add litellm mlflow

* test_mlflow_request_tags_functionality

* docs ML flow litellm proxy

* docs ml flow

* docs mlflow

docs/my-website/docs/observability/mlflow.md

@@ -17,7 +17,7 @@ MLflow’s integration with LiteLLM supports advanced observability compatible w
 Install MLflow:
 
 ```shell
-pip install mlflow
+pip install "litellm[mlflow]"
 ```
 
 To enable MLflow auto tracing for LiteLLM:
@@ -160,6 +160,102 @@ class CustomAgent:
 
 This approach generates a unified trace, combining your custom Python code with LiteLLM calls.
 
+## LiteLLM Proxy Server 
+
+### Dependencies
+
+For using `mlflow` on LiteLLM Proxy Server, you need to install the `mlflow` package on your docker container.
+
+```shell
+pip install "mlflow>=3.1.4"
+```
+
+### Configuration
+
+Configure MLflow in your LiteLLM proxy configuration file:
+
+```yaml
+model_list:
+  - model_name: openai/*
+    litellm_params:
+      model: openai/*
+
+litellm_settings:
+  success_callback: ["mlflow"]
+  failure_callback: ["mlflow"]
+```
+
+### Environment Variables
+
+For MLflow with Databricks service, set these required environment variables:
+
+```shell
+DATABRICKS_TOKEN="dapixxxxx"
+DATABRICKS_HOST="https://dbc-xxxx.cloud.databricks.com"
+MLFLOW_TRACKING_URI="databricks"
+MLFLOW_REGISTRY_URI="databricks-uc"
+MLFLOW_EXPERIMENT_ID="xxxx"
+```
+
+### Adding Tags for Better Tracing
+
+You can add custom tags to your requests for improved trace organization and filtering in MLflow. Tags help you categorize and search your traces by job ID, task name, or any custom metadata.
+
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+<Tabs>
+<TabItem value="curl" label="curl">
+
+```shell
+curl --location 'http://0.0.0.0:4000/chat/completions' \
+    --header 'Content-Type: application/json' \
+    --header 'Authorization: Bearer sk-1234' \
+    --data '{
+    "model": "gemini-2.5-flash",
+    "messages": [
+        {
+        "role": "user",
+        "content": "what llm are you"
+        }
+    ],
+    "litellm_metadata": {
+        "tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"]
+    }
+}'
+```
+
+</TabItem>
+<TabItem value="openai-python" label="OpenAI Python SDK">
+
+```python
+from openai import OpenAI
+
+# Initialize the OpenAI client pointing to your LiteLLM proxy
+client = OpenAI(
+    api_key="sk-1234",  # Your LiteLLM proxy API key
+    base_url="http://0.0.0.0:4000"  # Your LiteLLM proxy URL
+)
+
+# Make a request with tags in metadata
+response = client.chat.completions.create(
+    model="gemini-2.5-flash",
+    messages=[
+        {
+            "role": "user", 
+            "content": "what llm are you"
+        }
+    ],
+    extra_body={
+        "litellm_metadata": {
+            "tags": ["jobID:214590dsff09fds", "taskName:run_page_classification"]
+        }
+    }
+)
+```
+
+</TabItem>
+</Tabs>
 
 ## Support
 

docs/my-website/docs/proxy/logging.md

@@ -1590,54 +1590,7 @@ curl -X POST 'http://0.0.0.0:4000/chat/completions' \
 
 ## MLflow
 
-
-#### Step1: Install dependencies
-Install the dependencies.
-
-```shell
-pip install litellm mlflow
-```
-
-#### Step 2: Create a `config.yaml` with `mlflow` callback
-
-```yaml
-model_list:
-  - model_name: "*"
-    litellm_params:
-      model: "*"
-litellm_settings:
-  success_callback: ["mlflow"]
-  failure_callback: ["mlflow"]
-```
-
-#### Step 3: Start the LiteLLM proxy
-```shell
-litellm --config config.yaml
-```
-
-#### Step 4: Make a request
-
-```shell
-curl -X POST 'http://0.0.0.0:4000/chat/completions' \
--H 'Content-Type: application/json' \
--d '{
-    "model": "gpt-4o-mini",
-    "messages": [
-      {
-        "role": "user",
-        "content": "What is the capital of France?"
-      }
-    ]
-}'
-```
-
-#### Step 5: Review traces
-
-Run the following command to start MLflow UI and review recorded traces.
-
-```shell
-mlflow ui
-```
+👉 Follow the tutorial [here](../observability/mlflow) to get started with mlflow on LiteLLM Proxy Server
 
 
 

litellm/integrations/mlflow.py

@@ -1,10 +1,15 @@
 import json
 import threading
-from typing import Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from litellm._logging import verbose_logger
 from litellm.integrations.custom_logger import CustomLogger
 
+if TYPE_CHECKING:
+    from litellm.types.utils import StandardLoggingPayload
+else:
+    StandardLoggingPayload = Any
+
 
 class MlflowLogger(CustomLogger):
     def __init__(self):
@@ -178,7 +183,7 @@ class MlflowLogger(CustomLogger):
             "call_type": kwargs.get("call_type"),
             "model": kwargs.get("model"),
         }
-        standard_obj = kwargs.get("standard_logging_object")
+        standard_obj: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object")
         if standard_obj:
             attributes.update(
                 {
@@ -192,6 +197,7 @@ class MlflowLogger(CustomLogger):
                     "raw_llm_response": standard_obj.get("response"),
                     "response_cost": standard_obj.get("response_cost"),
                     "saved_cache_cost": standard_obj.get("saved_cache_cost"),
+                    "request_tags": standard_obj.get("request_tags"),
                 }
             )
         else:
@@ -226,6 +232,7 @@ class MlflowLogger(CustomLogger):
         """
         import mlflow
 
+
         call_type = kwargs.get("call_type", "completion")
         span_name = f"litellm-{call_type}"
         span_type = self._get_span_type(call_type)
@@ -237,7 +244,7 @@ class MlflowLogger(CustomLogger):
         if active_span := mlflow.get_current_active_span():  # type: ignore
             return self._client.start_span(
                 name=span_name,
-                request_id=active_span.request_id,
+                trace_id=active_span.request_id,
                 parent_id=active_span.span_id,
                 span_type=span_type,
                 inputs=inputs,
@@ -250,21 +257,24 @@ class MlflowLogger(CustomLogger):
                 span_type=span_type,
                 inputs=inputs,
                 attributes=attributes,
+                tags=self._transform_tag_list_to_dict(attributes.get("request_tags", [])),
                 start_time_ns=start_time_ns,
             )
+    def _transform_tag_list_to_dict(self, tag_list: list) -> dict:
+        return {tag: "" for tag in tag_list}
 
     def _end_span_or_trace(self, span, outputs, end_time_ns, status):
         """End an MLflow span or a trace."""
         if span.parent_id is None:
             self._client.end_trace(
-                request_id=span.request_id,
+                trace_id=span.request_id,
                 outputs=outputs,
                 status=status,
                 end_time_ns=end_time_ns,
             )
         else:
             self._client.end_span(
-                request_id=span.request_id,
+                trace_id=span.request_id,
                 span_id=span.span_id,
                 outputs=outputs,
                 status=status,

litellm/proxy/proxy_config.yaml

@@ -1,5 +1,9 @@
 model_list:
-  - model_name: vertex_ai/*
+  - model_name: openai/*
     litellm_params:
-      model: gemini/*
+      model: openai/*
 
+
+litellm_settings:
+  success_callback: ["mlflow"]
+  failure_callback: ["mlflow"]

pyproject.toml

@@ -63,6 +63,7 @@ litellm-enterprise = {version = "0.1.16", optional = true}
 diskcache = {version = "^5.6.1", optional = true}
 polars = {version = "^1.31.0", optional = true, python = ">=3.10"}
 semantic-router = {version = "*", optional = true, python = ">=3.9"}
+mlflow = {version = ">3.1.4", optional = true, python = ">=3.10"}
 
 [tool.poetry.extras]
 proxy = [
@@ -108,6 +109,8 @@ caching = ["diskcache"]
 
 semantic-router = ["semantic-router"]
 
+mlflow = ["mlflow"]
+
 [tool.isort]
 profile = "black"
 

tests/test_litellm/integrations/test_mlflow.py

@@ -0,0 +1,60 @@
+import asyncio
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+# Adds the grandparent directory to sys.path to allow importing project modules
+sys.path.insert(0, os.path.abspath("../.."))
+
+import pytest
+
+import litellm
+from litellm.integrations.mlflow import MlflowLogger
+
+
+@pytest.mark.asyncio
+async def test_mlflow_request_tags_functionality():
+    """Test that request_tags are properly extracted and transformed into tags for MLflow traces."""
+    
+    # Mock MLflow client and dependencies
+    mock_client = MagicMock()
+    mock_span = MagicMock()
+    mock_span.parent_id = None  # Simulate root trace
+    mock_span.request_id = "test_trace_id"
+    mock_client.start_trace.return_value = mock_span
+    
+    with patch('mlflow.tracking.MlflowClient', return_value=mock_client), \
+         patch('mlflow.get_current_active_span', return_value=None):
+        
+        # Create MlflowLogger instance
+        mlflow_logger = MlflowLogger()
+        litellm.callbacks = [mlflow_logger]
+        
+        # Test completion with request_tags
+        await litellm.acompletion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "test message"}],
+            mock_response="test response",
+            metadata={
+                "tags": ["tag1", "tag2", "production"]
+            }
+        )
+        
+        # Allow time for async processing
+        await asyncio.sleep(1)
+        
+        # Verify start_trace was called with tags parameter
+        assert mock_client.start_trace.called, "start_trace should have been called"
+        
+        # Get the call arguments
+        call_args = mock_client.start_trace.call_args
+        assert call_args is not None, "start_trace call args should not be None"
+        
+        # Check that tags parameter was included and properly transformed
+        tags_param = call_args.kwargs.get('tags', {})
+        expected_tags = {"tag1": "", "tag2": "", "production": ""}
+        assert tags_param == expected_tags, f"Expected tags {expected_tags}, got {tags_param}"
+        
+        print("✅ Request tags properly transformed and passed to MLflow trace")
+
+